"""Shared evaluation metrics for cognitive-level classification.

Used by both classical and deep models. Reports accuracy, macro-F1, and
Quadratic Weighted Kappa (QWK), which accounts for the ordinal label structure.
"""

import os

import matplotlib
matplotlib.use("Agg")  # headless backend, safe for servers
import matplotlib.pyplot as plt
import numpy as np
from sklearn.metrics import (
    accuracy_score,
    cohen_kappa_score,
    confusion_matrix,
    f1_score,
    precision_recall_fscore_support,
)


# Ordered class labels, lowest to highest cognitive level.
ORDERED_LABELS = ["Surface", "Mechanistic", "Critical"]


def _to_ordinal(labels):
    """Map class names to ordinal ranks."""
    index = {label: rank for rank, label in enumerate(ORDERED_LABELS)}
    return [index[label] for label in labels]


def compute_metrics(y_true, y_pred):
    """Compute accuracy, macro-F1, QWK, and per-class scores."""
    accuracy = accuracy_score(y_true, y_pred)
    macro_f1 = f1_score(y_true, y_pred, labels=ORDERED_LABELS, average="macro")
    qwk = cohen_kappa_score(
        _to_ordinal(y_true), _to_ordinal(y_pred), weights="quadratic"
    )

    precision, recall, f1, support = precision_recall_fscore_support(
        y_true, y_pred, labels=ORDERED_LABELS, zero_division=0
    )
    per_class = {}
    for i, label in enumerate(ORDERED_LABELS):
        per_class[label] = {
            "precision": float(precision[i]),
            "recall": float(recall[i]),
            "f1": float(f1[i]),
            "support": int(support[i]),
        }

    matrix = confusion_matrix(y_true, y_pred, labels=ORDERED_LABELS)

    return {
        "accuracy": float(accuracy),
        "macro_f1": float(macro_f1),
        "qwk": float(qwk),
        "per_class": per_class,
        "confusion_matrix": matrix.tolist(),
    }


def print_metrics(name, metrics):
    """Print headline metrics and per-class scores."""
    print(f"\n{name}")
    print(f"  accuracy : {metrics['accuracy']:.3f}")
    print(f"  macro-F1 : {metrics['macro_f1']:.3f}")
    print(f"  QWK      : {metrics['qwk']:.3f}")
    print(f"  {'class':12s} {'prec':>6s} {'rec':>6s} {'f1':>6s} {'n':>6s}")
    for label in ORDERED_LABELS:
        scores = metrics["per_class"][label]
        print(f"  {label:12s} {scores['precision']:6.3f} "
              f"{scores['recall']:6.3f} {scores['f1']:6.3f} {scores['support']:6d}")


def plot_confusion_matrix(metrics, title, save_path):
    """Save a confusion matrix plot to disk."""
    matrix = np.array(metrics["confusion_matrix"])
    os.makedirs(os.path.dirname(save_path), exist_ok=True)

    figure, axis = plt.subplots(figsize=(4.5, 4))
    image = axis.imshow(matrix, cmap="Blues")
    axis.set_xticks(range(len(ORDERED_LABELS)))
    axis.set_yticks(range(len(ORDERED_LABELS)))
    axis.set_xticklabels(ORDERED_LABELS, rotation=30, ha="right")
    axis.set_yticklabels(ORDERED_LABELS)
    axis.set_xlabel("Predicted")
    axis.set_ylabel("True")
    axis.set_title(title)

    threshold = matrix.max() / 2.0 if matrix.max() > 0 else 0.5
    for row in range(matrix.shape[0]):
        for col in range(matrix.shape[1]):
            axis.text(
                col, row, int(matrix[row, col]),
                ha="center", va="center",
                color="white" if matrix[row, col] > threshold else "black",
            )
    figure.colorbar(image, ax=axis, fraction=0.046, pad=0.04)
    figure.tight_layout()
    figure.savefig(save_path, dpi=150)
    plt.close(figure)
    print(f"  saved confusion matrix -> {save_path}")