ml/src/evaluation/plots.py

"""
Visualization module.
Generates training curves, confusion matrices, and per-class metric charts.
"""

from pathlib import Path
from typing import Optional

import numpy as np
import matplotlib
matplotlib.use('Agg')  # Non-interactive backend
import matplotlib.pyplot as plt
import matplotlib.ticker as mticker
import seaborn as sns


# Style defaults
plt.style.use('seaborn-v0_8-whitegrid')
COLORS = {
    'train': '#2196F3',
    'val': '#FF5722',
    'primary': '#4CAF50',
    'secondary': '#FFC107',
}


def plot_training_curves(
    history: dict,
    save_path: str | Path = None,
    title: str = 'Training Curves',
    figsize: tuple = (14, 5),
) -> None:
    """Plot loss and accuracy curves for training and validation."""
    fig, axes = plt.subplots(1, 2, figsize=figsize)

    epochs = range(1, len(history['train_loss']) + 1)

    # Loss
    axes[0].plot(epochs, history['train_loss'], color=COLORS['train'], label='Train', linewidth=2)
    axes[0].plot(epochs, history['val_loss'], color=COLORS['val'], label='Validation', linewidth=2)
    axes[0].set_title(f'{title} - Loss', fontsize=13, fontweight='bold')
    axes[0].set_xlabel('Epoch')
    axes[0].set_ylabel('Loss')
    axes[0].legend()
    axes[0].grid(True, alpha=0.3)

    # Accuracy
    axes[1].plot(epochs, history['train_accuracy'], color=COLORS['train'], label='Train', linewidth=2)
    axes[1].plot(epochs, history['val_accuracy'], color=COLORS['val'], label='Validation', linewidth=2)
    axes[1].set_title(f'{title} - Accuracy', fontsize=13, fontweight='bold')
    axes[1].set_xlabel('Epoch')
    axes[1].set_ylabel('Accuracy')
    axes[1].legend()
    axes[1].grid(True, alpha=0.3)
    axes[1].yaxis.set_major_formatter(mticker.PercentFormatter(1.0))

    plt.tight_layout()

    if save_path:
        Path(save_path).parent.mkdir(parents=True, exist_ok=True)
        plt.savefig(save_path, dpi=150, bbox_inches='tight')
        print(f"  Saved training curves to {save_path}")

    plt.close()


def plot_confusion_matrix(
    y_true: list[int],
    y_pred: list[int],
    class_names: list[str],
    save_path: str | Path = None,
    title: str = 'Confusion Matrix',
    figsize: tuple = None,
    normalize: bool = True,
) -> None:
    """Plot confusion matrix heatmap."""
    from sklearn.metrics import confusion_matrix as sk_cm

    cm = sk_cm(y_true, y_pred)
    n_classes = len(class_names)

    if figsize is None:
        size = max(10, n_classes * 0.5)
        figsize = (size, size)

    if normalize:
        cm_normalized = cm.astype('float') / cm.sum(axis=1)[:, np.newaxis]
        cm_display = cm_normalized
        fmt = '.2f'
    else:
        cm_display = cm
        fmt = 'd'

    fig, ax = plt.subplots(figsize=figsize)

    sns.heatmap(
        cm_display,
        annot=True,
        fmt=fmt,
        cmap='Blues',
        xticklabels=class_names,
        yticklabels=class_names,
        ax=ax,
        linewidths=0.5,
        cbar_kws={'shrink': 0.8},
    )

    ax.set_title(title, fontsize=14, fontweight='bold', pad=15)
    ax.set_xlabel('Predicted', fontsize=12)
    ax.set_ylabel('True', fontsize=12)

    # Rotate labels
    plt.xticks(rotation=45, ha='right', fontsize=8)
    plt.yticks(rotation=0, fontsize=8)

    plt.tight_layout()

    if save_path:
        Path(save_path).parent.mkdir(parents=True, exist_ok=True)
        plt.savefig(save_path, dpi=150, bbox_inches='tight')
        print(f"  Saved confusion matrix to {save_path}")

    plt.close()


def plot_per_class_f1(
    per_class_metrics: dict,
    class_names: list[str],
    save_path: str | Path = None,
    title: str = 'Per-Class F1 Score',
    figsize: tuple = None,
) -> None:
    """Plot horizontal bar chart of per-class F1 scores."""
    f1_scores = []
    names = []

    for name in class_names:
        if name in per_class_metrics:
            f1_scores.append(per_class_metrics[name]['f1'])
            names.append(name)

    # Sort by F1 score
    sorted_pairs = sorted(zip(names, f1_scores), key=lambda x: x[1])
    names, f1_scores = zip(*sorted_pairs) if sorted_pairs else ([], [])

    n = len(names)
    if figsize is None:
        figsize = (10, max(6, n * 0.35))

    fig, ax = plt.subplots(figsize=figsize)

    colors = [plt.cm.RdYlGn(score) for score in f1_scores]

    bars = ax.barh(range(n), f1_scores, color=colors, edgecolor='white', linewidth=0.5)

    ax.set_yticks(range(n))
    ax.set_yticklabels(names, fontsize=9)
    ax.set_xlabel('F1 Score', fontsize=11)
    ax.set_title(title, fontsize=13, fontweight='bold')
    ax.set_xlim(0, 1.05)

    # Add value labels
    for bar, score in zip(bars, f1_scores):
        ax.text(bar.get_width() + 0.01, bar.get_y() + bar.get_height() / 2,
                f'{score:.2f}', va='center', fontsize=8)

    ax.axvline(x=np.mean(list(f1_scores)), color='red', linestyle='--',
               alpha=0.7, label=f'Mean: {np.mean(list(f1_scores)):.2f}')
    ax.legend(fontsize=10)

    plt.tight_layout()

    if save_path:
        Path(save_path).parent.mkdir(parents=True, exist_ok=True)
        plt.savefig(save_path, dpi=150, bbox_inches='tight')
        print(f"  Saved per-class F1 chart to {save_path}")

    plt.close()


def plot_sample_predictions(
    images: list,
    true_labels: list[str],
    pred_labels: list[str],
    confidences: list[float],
    save_path: str | Path = None,
    title: str = 'Sample Predictions',
    n_samples: int = 16,
) -> None:
    """Plot grid of sample predictions with true vs predicted labels."""
    n = min(n_samples, len(images))
    cols = 4
    rows = (n + cols - 1) // cols

    fig, axes = plt.subplots(rows, cols, figsize=(cols * 3.5, rows * 3.5))
    if rows == 1:
        axes = [axes]
    axes = [ax for row in axes for ax in (row if hasattr(row, '__iter__') else [row])]

    for i in range(n):
        ax = axes[i]
        img = images[i]
        if hasattr(img, 'numpy'):
            img = img.numpy()
        if img.shape[0] == 3:
            img = np.transpose(img, (1, 2, 0))
        img = np.clip(img, 0, 1)

        ax.imshow(img)
        correct = true_labels[i] == pred_labels[i]
        color = 'green' if correct else 'red'
        ax.set_title(
            f"True: {true_labels[i]}\nPred: {pred_labels[i]} ({confidences[i]:.0%})",
            fontsize=8, color=color, fontweight='bold',
        )
        ax.axis('off')

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

    fig.suptitle(title, fontsize=14, fontweight='bold', y=1.02)
    plt.tight_layout()

    if save_path:
        Path(save_path).parent.mkdir(parents=True, exist_ok=True)
        plt.savefig(save_path, dpi=150, bbox_inches='tight')

    plt.close()