src/wav2vec2/evaluate.py

"""
Responsibilities:
1. Compute all metrics on the test set
(Accuracy, Weighted F1, UAR, Per-class F1)
2. Plot confusion matrix (normalized)
3. Plot training curves (loss, accuracy, UAR per epoch)
4. Save a full results summary to a text file

Metrics explained:
- Accuracy : % of correct predictions (biased toward neutral)
- UAR      : Unweighted Average Recall = macro recall
THE standard metric in SER — treats all classes
equally regardless of size
- Weighted F1: F1 weighted by class size (for reference)
- Per-class F1: shows which emotions are hardest

Imports used by: main.py
"""

import os
import json
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.ticker as ticker
from sklearn.metrics import (
accuracy_score,
f1_score,
recall_score,
classification_report,
confusion_matrix,
)

from .config import (
RESULTS_DIR,
IDX_TO_EMOTION,
NUM_CLASSES,
get_results_path,
)

# Emotion names in index order (0–3) for plot labels

EMOTION_LABELS = [IDX_TO_EMOTION[i] for i in range(NUM_CLASSES)]

# ─────────────────────────────────────────────────────────────────────────────

# PART A — Metrics Computation

# ─────────────────────────────────────────────────────────────────────────────

def compute_metrics(all_labels, all_preds):
    """
    Compute all evaluation metrics from raw label and prediction lists.

    ```
    Args:
        all_labels (list[int]): ground truth class indices
        all_preds  (list[int]): predicted class indices

    Returns:
        dict with keys: accuracy, uar, weighted_f1, per_class_f1,
                        classification_report (string)
    """
    accuracy     = accuracy_score(all_labels, all_preds)
    uar          = recall_score(all_labels, all_preds, average="macro",    zero_division=0)
    weighted_f1  = f1_score(all_labels,     all_preds, average="weighted", zero_division=0)
    per_class_f1 = f1_score(all_labels,     all_preds, average=None,       zero_division=0)
    report       = classification_report(
        all_labels, all_preds,
        target_names=EMOTION_LABELS,
        zero_division=0
    )

    metrics = {
        "accuracy":              round(accuracy,    4),
        "uar":                   round(uar,         4),
        "weighted_f1":           round(weighted_f1, 4),
        "per_class_f1": {
            EMOTION_LABELS[i]: round(float(per_class_f1[i]), 4)
            for i in range(len(EMOTION_LABELS))
        },
        "classification_report": report,
    }

    return metrics


def print_metrics(metrics, split_name="Test"):
    """Print a clean metrics summary to the console."""
    print("\n" + "=" * 55)
    print(f"Results on {split_name} Set")
    print("=" * 55)
    print(f"  Accuracy    : {metrics['accuracy']:.4f}  ({metrics['accuracy']*100:.2f}%)")
    print(f"  UAR         : {metrics['uar']:.4f}  ({metrics['uar']*100:.2f}%)")
    print(f"              (UAR = Unweighted Average Recall — main SER metric)")
    print(f"  Weighted F1 : {metrics['weighted_f1']:.4f}")
    print(f"\n  Per-class F1:")
    for emotion, f1 in metrics["per_class_f1"].items():
        print(f"    {emotion:8s}: {f1:.4f}")
    print(f"\n  Full Report:\n{metrics['classification_report']}")
    print("=" * 55)

# ─────────────────────────────────────────────────────────────────────────────

# PART B — Confusion Matrix Plot

# ─────────────────────────────────────────────────────────────────────────────

def plot_confusion_matrix(all_labels, all_preds, dataset_name="iemocap", save=True):
    """
    Plot a normalized confusion matrix.

    ```
    Normalized means each row sums to 1 (i.e., recall per class).
    This makes it easy to see which emotions the model confuses.
    E.g., if happy and excited look similar, happy→neutral will have a
    high off-diagonal value.

    Args:
        all_labels (list[int]): ground truth
        all_preds  (list[int]): predictions
        save (bool): save to RESULTS_DIR/confusion_matrix.png
    """
    cm = confusion_matrix(all_labels, all_preds)

    cm_norm = np.divide(
        cm.astype(float),
        cm.sum(axis=1, keepdims=True),
        where=cm.sum(axis=1, keepdims=True) != 0
    )

    fig, axes = plt.subplots(1, 2, figsize=(14, 5))

    for ax, data, title, fmt in zip(
        axes,
        [cm, cm_norm],
        ["Confusion Matrix (Counts)", "Confusion Matrix (Normalized — Recall per Class)"],
        ["d", ".2f"]
    ):
        im = ax.imshow(data, interpolation="nearest", cmap="Blues")
        ax.set_title(title, fontsize=13, fontweight="bold", pad=12)
        ax.set_xlabel("Predicted", fontsize=11)
        ax.set_ylabel("True", fontsize=11)
        ax.set_xticks(range(NUM_CLASSES))
        ax.set_yticks(range(NUM_CLASSES))
        ax.set_xticklabels(EMOTION_LABELS, rotation=45, ha="right")
        ax.set_yticklabels(EMOTION_LABELS)

        thresh = data.max() / 2.0
        for i in range(NUM_CLASSES):
            for j in range(NUM_CLASSES):
                val = data[i, j]
                text = f"{val:{fmt}}" if fmt == ".2f" else f"{int(val)}"
                color = "white" if val > thresh else "black"
                ax.text(j, i, text, ha="center", va="center",
                        color=color, fontsize=10, fontweight="bold")

        plt.colorbar(im, ax=ax)

    plt.tight_layout()

    if save:
        path = get_results_path(dataset_name, "confusion_matrix.png")
        plt.savefig(path, dpi=150, bbox_inches="tight")
        print(f"\n[Saved] Confusion matrix → {path}")

    plt.show()


# ─────────────────────────────────────────────────────────────────────────────

# PART C — Training Curves Plot

# ─────────────────────────────────────────────────────────────────────────────

def plot_training_curves(history, dataset_name="iemocap", save=True):
    """
    Plot training and validation curves across epochs:
    - Loss (left)
    - Accuracy (middle)
    - UAR / Unweighted Average Recall (right)

    ```
    The UAR plot is the most important — it shows how well the model
    handles class imbalance across training epochs.

    Args:
        history (dict): from person3_train.py train() — keys are
                        train_loss, val_loss, train_acc, val_acc,
                        train_uar, val_uar
        save (bool): save to RESULTS_DIR/training_curves.png
    """
    epochs = range(1, len(history["train_loss"]) + 1)

    fig, axes = plt.subplots(1, 3, figsize=(16, 5))
    fig.suptitle("Wav2Vec2 Fine-Tuning — Training Curves", fontsize=14, fontweight="bold")

    metrics = [
        ("loss",  "Loss",                "lower"),
        ("acc",   "Accuracy",            "upper"),
        ("uar",   "UAR (Macro Recall)",  "upper"),
    ]

    for ax, (key, ylabel, best_direction) in zip(axes, metrics):
        train_vals = history[f"train_{key}"]
        val_vals   = history[f"val_{key}"]

        ax.plot(epochs, train_vals, "b-o", markersize=4, label="Train",      linewidth=2)
        ax.plot(epochs, val_vals,   "r-o", markersize=4, label="Validation", linewidth=2)

        best_fn  = min if best_direction == "lower" else max
        best_idx = val_vals.index(best_fn(val_vals))
        ax.axvline(x=best_idx + 1, color="green", linestyle="--",
                alpha=0.6, label=f"Best val (ep {best_idx+1})")
        ax.plot(best_idx + 1, val_vals[best_idx], "g*", markersize=12)

        ax.set_xlabel("Epoch",    fontsize=11)
        ax.set_ylabel(ylabel,     fontsize=11)
        ax.set_title(ylabel,      fontsize=12, fontweight="bold")
        ax.legend(fontsize=9)
        ax.grid(True, alpha=0.3)
        ax.xaxis.set_major_locator(ticker.MaxNLocator(integer=True))

    plt.tight_layout()

    if save:
        path = get_results_path(dataset_name, "training_curves.png")
        plt.savefig(path, dpi=150, bbox_inches="tight")
        print(f"[Saved] Training curves → {path}")

    plt.show()


# ─────────────────────────────────────────────────────────────────────────────

# PART D — Per-Class F1 Bar Chart

# ─────────────────────────────────────────────────────────────────────────────

def plot_per_class_f1(metrics, dataset_name="iemocap", save=True):
    """
    Bar chart showing F1 score per emotion class.
    Visually shows which emotions the model recognizes well vs. struggles with.
    """
    emotions = list(metrics["per_class_f1"].keys())
    f1_scores = list(metrics["per_class_f1"].values())
    colors = ["#e74c3c", "#f39c12", "#3498db", "#2ecc71"]


    fig, ax = plt.subplots(figsize=(7, 4))
    bars = ax.bar(emotions, f1_scores, color=colors, edgecolor="white", width=0.5)

    for bar, val in zip(bars, f1_scores):
        ax.text(
            bar.get_x() + bar.get_width() / 2,
            bar.get_height() + 0.01,
            f"{val:.3f}",
            ha="center", va="bottom", fontsize=11, fontweight="bold"
        )

    ax.set_ylim(0, 1.05)
    ax.set_xlabel("Emotion", fontsize=12)
    ax.set_ylabel("F1 Score", fontsize=12)
    ax.set_title("Per-Class F1 Score on Test Set", fontsize=13, fontweight="bold")
    ax.axhline(y=sum(f1_scores)/len(f1_scores), color="gray",
            linestyle="--", alpha=0.7, label="Mean F1")
    ax.legend()
    ax.grid(axis="y", alpha=0.3)

    plt.tight_layout()

    if save:
        path = get_results_path(dataset_name, "per_class_f1.png")
        plt.savefig(path, dpi=150, bbox_inches="tight")
        print(f"[Saved] Per-class F1 → {path}")

    plt.show()


# ─────────────────────────────────────────────────────────────────────────────

# PART E — Save Results Summary

# ─────────────────────────────────────────────────────────────────────────────

def save_results(metrics, history, dataset_name="iemocap"):
    """
    Save metrics and training history to text and JSON files.
    These can be included directly in the final report.
    """
    summary_path = get_results_path(dataset_name, "results_summary.txt")
    with open(summary_path, "w") as f:
        f.write("Wav2Vec2 IEMOCAP Results\n")
        f.write("=" * 55 + "\n\n")
        f.write(f"Accuracy    : {metrics['accuracy']:.4f}\n")
        f.write(f"UAR         : {metrics['uar']:.4f}\n")
        f.write(f"Weighted F1 : {metrics['weighted_f1']:.4f}\n\n")
        f.write("Per-class F1:\n")
        for emo, val in metrics["per_class_f1"].items():
            f.write(f"  {emo:8s}: {val:.4f}\n")
        f.write(f"\nFull Classification Report:\n{metrics['classification_report']}\n")
        f.write(f"\nTraining Epochs Run: {len(history['train_loss'])}\n")
        f.write(f"Best Val UAR       : {max(history['val_uar']):.4f}\n")
    print(f"[Saved] Results summary → {summary_path}")

    json_path = get_results_path(dataset_name, "metrics.json")
    with open(json_path, "w") as f:
        exportable = {k: v for k, v in metrics.items() if k != "classification_report"}
        exportable["history"] = history
        json.dump(exportable, f, indent=2)
    print(f"[Saved] Metrics JSON   → {json_path}")


# ─────────────────────────────────────────────────────────────────────────────

# PART F — Main evaluation function (called from main.py)

# ─────────────────────────────────────────────────────────────────────────────

def run_full_evaluation(model, test_loader, criterion, device, history, dataset_name="iemocap"):
    """
    Run full evaluation: compute metrics, plot all figures, save results.

    ```
    Args:
        model       : trained Wav2Vec2EmotionClassifier (best checkpoint loaded)
        test_loader : test DataLoader (Person 1)
        criterion   : CrossEntropyLoss (Person 3)
        device      : torch.device
        history     : training history dict (Person 3)

    Returns:
        metrics (dict)
    """
    from .train import evaluate_split

    print("\nRunning full evaluation on test set...")
    test_loss, test_acc, test_uar, all_preds, all_labels = evaluate_split(
        model, test_loader, criterion, device
    )

    metrics = compute_metrics(all_labels, all_preds)
    print_metrics(metrics, split_name="Test")

    plot_confusion_matrix(all_labels, all_preds, dataset_name, save=True)
    plot_training_curves(history, dataset_name, save=True)
    plot_per_class_f1(metrics, dataset_name, save=True)

    save_results(metrics, history, dataset_name)

    return metrics


# ─────────────────────────────────────────────────────────────────────────────

# Quick sanity test — run this file directly with dummy data

# ─────────────────────────────────────────────────────────────────────────────

if __name__ == "__main__":
    import random
    random.seed(42)


    n = 200
    all_labels = [random.randint(0, 3) for _ in range(n)]
    all_preds  = [random.randint(0, 3) for _ in range(n)]

    metrics = compute_metrics(all_labels, all_preds)
    print_metrics(metrics, "Dummy")
    plot_confusion_matrix(all_labels, all_preds, save=False)
    plot_per_class_f1(metrics, save=False)

    dummy_history = {
        "train_loss": [1.4, 1.2, 1.0, 0.9, 0.85],
        "val_loss":   [1.5, 1.3, 1.1, 1.0, 0.95],
        "train_acc":  [0.40, 0.50, 0.60, 0.65, 0.68],
        "val_acc":    [0.38, 0.48, 0.57, 0.60, 0.64],
        "train_uar":  [0.35, 0.45, 0.55, 0.60, 0.63],
        "val_uar":    [0.33, 0.43, 0.52, 0.57, 0.60],
    }
    plot_training_curves(dummy_history, save=False)
    print("\n[Person 4] Evaluation module verified successfully!")