"""
Plotting utilities shared across all three courseworks.

All functions save plots to files and optionally display them.
Designed to support report generation without manual plotting code.
"""

import os
import numpy as np
import matplotlib
matplotlib.use('Agg')  # Non-interactive backend for saving plots
import matplotlib.pyplot as plt
import seaborn as sns


def plot_loss_curves(train_losses, val_losses, title="Training & Validation Loss",
                     save_path=None):
    """
    Plot training and validation loss curves over epochs.

    Args:
        train_losses: List of training loss values per epoch
        val_losses: List of validation loss values per epoch
        title: Plot title
        save_path: Path to save figure (if None, shows plot)
    """
    fig, ax = plt.subplots(1, 1, figsize=(8, 5))
    epochs = range(1, len(train_losses) + 1)
    ax.plot(epochs, train_losses, 'b-', label='Train Loss', linewidth=2)
    ax.plot(epochs, val_losses, 'r-', label='Val Loss', linewidth=2)
    ax.set_xlabel('Epoch', fontsize=12)
    ax.set_ylabel('Loss', fontsize=12)
    ax.set_title(title, fontsize=14)
    ax.legend(fontsize=11)
    ax.grid(True, alpha=0.3)
    plt.tight_layout()
    _save_or_show(fig, save_path)


def plot_accuracy_curves(train_accs, val_accs, title="Training & Validation Accuracy",
                         save_path=None):
    """
    Plot training and validation accuracy curves over epochs.

    Args:
        train_accs: List of training accuracy values per epoch
        val_accs: List of validation accuracy values per epoch
        title: Plot title
        save_path: Path to save figure (if None, shows plot)
    """
    fig, ax = plt.subplots(1, 1, figsize=(8, 5))
    epochs = range(1, len(train_accs) + 1)
    ax.plot(epochs, train_accs, 'b-', label='Train Accuracy', linewidth=2)
    ax.plot(epochs, val_accs, 'r-', label='Val Accuracy', linewidth=2)
    ax.set_xlabel('Epoch', fontsize=12)
    ax.set_ylabel('Accuracy', fontsize=12)
    ax.set_title(title, fontsize=14)
    ax.legend(fontsize=11)
    ax.grid(True, alpha=0.3)
    ax.set_ylim([0, 1.0])
    plt.tight_layout()
    _save_or_show(fig, save_path)


def plot_confusion_matrix(cm, class_names, title="Confusion Matrix", save_path=None):
    """
    Plot a confusion matrix as a heatmap.

    Args:
        cm: (num_classes, num_classes) confusion matrix array
        class_names: List of class name strings
        title: Plot title
        save_path: Path to save figure
    """
    fig, ax = plt.subplots(1, 1, figsize=(10, 8))
    sns.heatmap(cm, annot=True, fmt='d', cmap='Blues',
                xticklabels=class_names, yticklabels=class_names, ax=ax)
    ax.set_xlabel('Predicted', fontsize=12)
    ax.set_ylabel('True', fontsize=12)
    ax.set_title(title, fontsize=14)
    plt.xticks(rotation=45, ha='right')
    plt.yticks(rotation=0)
    plt.tight_layout()
    _save_or_show(fig, save_path)


def plot_sample_predictions(images, true_labels, pred_labels, class_names,
                            save_path=None, n=16):
    """
    Plot a grid of sample images with their predicted and true labels.
    Correct predictions are shown in green, incorrect in red.

    Args:
        images: (N, 784) array of flattened images (or (N, 28, 28))
        true_labels: (N,) true class indices
        pred_labels: (N,) predicted class indices
        class_names: List of class name strings
        save_path: Path to save figure
        n: Number of samples to show (default: 16, arranged as 4x4)
    """
    n = min(n, len(images))
    cols = 4
    rows = (n + cols - 1) // cols
    fig, axes = plt.subplots(rows, cols, figsize=(12, 3 * rows))
    axes = axes.flatten()

    for i in range(n):
        img = images[i].reshape(28, 28) if images[i].ndim == 1 else images[i]
        axes[i].imshow(img, cmap='gray')
        axes[i].axis('off')

        true_name = class_names[true_labels[i]]
        pred_name = class_names[pred_labels[i]]
        correct = true_labels[i] == pred_labels[i]
        color = 'green' if correct else 'red'
        axes[i].set_title(f"True: {true_name}\nPred: {pred_name}",
                          fontsize=9, color=color)

    # Hide unused axes
    for i in range(n, len(axes)):
        axes[i].axis('off')

    plt.suptitle("Sample Predictions", fontsize=14)
    plt.tight_layout()
    _save_or_show(fig, save_path)


def plot_loss_comparison(results_dict, title="Loss Function Comparison",
                         save_path=None):
    """
    Overlay training curves for multiple loss functions (used in CW2 report).

    Left panel: Training accuracy over epochs — directly shows convergence speed.
    Loss functions with smaller gradient magnitude (e.g., squared loss due to the
    softmax Jacobian factor) converge to high accuracy more slowly, making the
    plateau effect visible as a lower starting accuracy and a slower rise.

    Right panel: Validation accuracy with auto-zoomed Y-axis — shows final
    performance differences across the relevant accuracy range.

    Args:
        results_dict: Dict mapping loss_name -> dict with keys:
            'train_losses': list of per-epoch losses
            'train_accs': list of per-epoch training accuracies
            'val_accs': list of per-epoch validation accuracies
        title: Plot title
        save_path: Path to save figure
    """
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(14, 5))

    colors = {'cross_entropy': '#1f77b4', 'hinge': '#ff7f0e',
              'exponential': '#2ca02c', 'squared': '#d62728'}
    default_colors = plt.cm.Set1(np.linspace(0, 1, len(results_dict)))

    all_train_accs = []
    all_val_accs = []
    for idx, (name, results) in enumerate(results_dict.items()):
        epochs = range(1, len(results['train_accs']) + 1)
        color = colors.get(name, default_colors[idx])

        # -- Left: Training accuracy (convergence speed & plateau effect) --
        train_accs = results['train_accs']
        ax1.plot(epochs, train_accs, color=color, label=name, linewidth=2)
        all_train_accs.extend(train_accs)

        # -- Right: Validation accuracy --
        val_accs = results['val_accs']
        ax2.plot(epochs, val_accs, color=color, label=name, linewidth=2)
        all_val_accs.extend(val_accs)

    # Left: auto-zoom training accuracy Y-axis
    ax1.set_xlabel('Epoch', fontsize=12)
    ax1.set_ylabel('Training Accuracy', fontsize=12)
    ax1.set_title('Training Accuracy Curves', fontsize=13)
    ax1.legend(fontsize=10)
    ax1.grid(True, alpha=0.3)
    if all_train_accs:
        ta_min = min(all_train_accs)
        ta_max = max(all_train_accs)
        margin = max(0.03, (ta_max - ta_min) * 0.15)
        ax1.set_ylim([max(0, ta_min - margin), min(1.0, ta_max + margin)])

    # Right: auto-zoom validation accuracy Y-axis
    ax2.set_xlabel('Epoch', fontsize=12)
    ax2.set_ylabel('Validation Accuracy', fontsize=12)
    ax2.set_title('Validation Accuracy Curves', fontsize=13)
    ax2.legend(fontsize=10)
    ax2.grid(True, alpha=0.3)
    if all_val_accs:
        va_min = min(all_val_accs)
        va_max = max(all_val_accs)
        margin = max(0.03, (va_max - va_min) * 0.15)
        ax2.set_ylim([max(0, va_min - margin), min(1.0, va_max + margin)])

    plt.suptitle(title, fontsize=14)
    plt.tight_layout()
    _save_or_show(fig, save_path)


def plot_weight_visualization(weights, class_names, save_path=None):
    """
    Visualize learned weight templates for logistic regression.
    Reshapes each class's weight vector from (784,) to (28, 28) image.

    Args:
        weights: (num_classes, 784) or (784, num_classes) weight matrix
        class_names: List of class name strings
        save_path: Path to save figure
    """
    # Ensure shape is (num_classes, 784)
    if weights.shape[0] == 784:
        weights = weights.T

    num_classes = weights.shape[0]
    fig, axes = plt.subplots(2, 5, figsize=(15, 6))
    axes = axes.flatten()

    for i in range(num_classes):
        w = weights[i].reshape(28, 28)
        im = axes[i].imshow(w, cmap='RdBu_r', vmin=-np.abs(w).max(), vmax=np.abs(w).max())
        axes[i].set_title(class_names[i], fontsize=11)
        axes[i].axis('off')

    plt.suptitle("Learned Weight Templates (Logistic Regression)", fontsize=14)
    plt.tight_layout()
    _save_or_show(fig, save_path)


def plot_hyperparameter_search(param_name, param_values, train_accs, val_accs,
                               save_path=None):
    """
    Plot train/val accuracy vs a hyperparameter value (used in CW3 report).

    Args:
        param_name: Name of the hyperparameter (e.g., "Hidden Dimension")
        param_values: List of hyperparameter values tested
        train_accs: List of final training accuracies for each value
        val_accs: List of final validation accuracies for each value
        save_path: Path to save figure
    """
    fig, ax = plt.subplots(1, 1, figsize=(8, 5))

    x = range(len(param_values))
    ax.plot(x, train_accs, 'bo-', label='Train Accuracy', linewidth=2, markersize=8)
    ax.plot(x, val_accs, 'ro-', label='Val Accuracy', linewidth=2, markersize=8)
    ax.set_xticks(x)
    ax.set_xticklabels([str(v) for v in param_values], rotation=45, ha='right')
    ax.set_xlabel(param_name, fontsize=12)
    ax.set_ylabel('Accuracy', fontsize=12)
    ax.set_title(f'Accuracy vs {param_name}', fontsize=14)
    ax.legend(fontsize=11)
    ax.grid(True, alpha=0.3)
    ax.set_ylim([0, 1.0])
    plt.tight_layout()
    _save_or_show(fig, save_path)


def plot_training_summary(train_losses, val_losses, train_accs, val_accs,
                          title="Training Summary", save_path=None):
    """
    Combined plot with loss curves (left) and accuracy curves (right) in one figure.
    This is the recommended plot to include in your report.

    Args:
        train_losses: List of training loss values per epoch
        val_losses: List of validation loss values per epoch
        train_accs: List of training accuracy values per epoch
        val_accs: List of validation accuracy values per epoch
        title: Plot title
        save_path: Path to save figure
    """
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(14, 5))
    epochs = range(1, len(train_losses) + 1)

    # Loss curves
    ax1.plot(epochs, train_losses, 'b-', label='Train Loss', linewidth=2)
    ax1.plot(epochs, val_losses, 'r-', label='Val Loss', linewidth=2)
    ax1.set_xlabel('Epoch', fontsize=12)
    ax1.set_ylabel('Loss', fontsize=12)
    ax1.set_title('Loss', fontsize=13)
    ax1.legend(fontsize=10)
    ax1.grid(True, alpha=0.3)

    # Accuracy curves
    ax2.plot(epochs, train_accs, 'b-', label='Train Accuracy', linewidth=2)
    ax2.plot(epochs, val_accs, 'r-', label='Val Accuracy', linewidth=2)
    ax2.set_xlabel('Epoch', fontsize=12)
    ax2.set_ylabel('Accuracy', fontsize=12)
    ax2.set_title('Accuracy', fontsize=13)
    ax2.legend(fontsize=10)
    ax2.grid(True, alpha=0.3)
    ax2.set_ylim([0, 1.0])

    plt.suptitle(title, fontsize=14)
    plt.tight_layout()
    _save_or_show(fig, save_path)


def _save_or_show(fig, save_path):
    """Save figure to file or display it."""
    if save_path is not None:
        os.makedirs(os.path.dirname(save_path), exist_ok=True) if os.path.dirname(save_path) else None
        fig.savefig(save_path, dpi=150, bbox_inches='tight')
        print(f"  Plot saved to: {save_path}")
        plt.close(fig)
    else:
        plt.show()
