"""
Evaluation Metrics for AES
Computes QWK, Accuracy, F1-Macro, Spearman, MAE — standard metrics for AES research.
"""
import numpy as np
from sklearn.metrics import (
    cohen_kappa_score,
    accuracy_score,
    f1_score,
    mean_absolute_error,
    classification_report,
    confusion_matrix,
)
from scipy.stats import spearmanr


def compute_all_metrics(y_true, y_pred):
    """
    Compute all evaluation metrics for AES.
    
    Args:
        y_true: Ground truth scores (1-5)
        y_pred: Predicted scores (1-5)
    
    Returns:
        dict with all metrics
    """
    y_true = np.array(y_true)
    y_pred = np.array(y_pred)

    return {
        "qwk": cohen_kappa_score(y_true, y_pred, weights="quadratic"),
        "accuracy": accuracy_score(y_true, y_pred),
        "f1_macro": f1_score(y_true, y_pred, average="macro", zero_division=0),
        "spearman": spearmanr(y_true, y_pred).correlation,
        "mae": mean_absolute_error(y_true, y_pred),
    }


def print_evaluation_report(y_true, y_pred, label_names=None):
    """
    Print a full evaluation report including confusion matrix.
    
    Args:
        y_true: Ground truth scores
        y_pred: Predicted scores
        label_names: Optional label names for the report
    """
    if label_names is None:
        label_names = [f"Score {i}" for i in range(1, 6)]

    metrics = compute_all_metrics(y_true, y_pred)

    print("\n" + "=" * 60)
    print("  EVALUATION REPORT — AES-Feedback")
    print("=" * 60)

    print("\n  Primary Metrics:")
    print(f"    {'QWK':>20}: {metrics['qwk']:.4f}")
    print(f"    {'Accuracy':>20}: {metrics['accuracy']:.4f}")
    print(f"    {'F1-Macro':>20}: {metrics['f1_macro']:.4f}")
    print(f"    {'Spearman Corr.':>20}: {metrics['spearman']:.4f}")
    print(f"    {'MAE':>20}: {metrics['mae']:.4f}")

    print("\n  Classification Report:")
    print(
        classification_report(
            y_true, y_pred, target_names=label_names, zero_division=0
        )
    )

    print("  Confusion Matrix:")
    cm = confusion_matrix(y_true, y_pred)
    # Pretty print
    header = "       " + "  ".join([f"P={i}" for i in range(1, 6)])
    print(f"  {header}")
    for i, row in enumerate(cm):
        row_str = "  ".join([f"{v:4d}" for v in row])
        print(f"  T={i+1}  {row_str}")

    print("=" * 60)
    return metrics