"""Metrics and helper utilities for training."""

import numpy as np
from sklearn.metrics import (
    accuracy_score,
    precision_recall_fscore_support,
    confusion_matrix,
    classification_report,
    roc_auc_score,
)


def compute_metrics(eval_pred):
    """Compute metrics for HuggingFace Trainer evaluation."""
    logits, labels = eval_pred
    preds = np.argmax(logits, axis=-1)

    # Numerically stable softmax
    logits_shifted = logits - logits.max(axis=-1, keepdims=True)
    exp_logits = np.exp(logits_shifted)
    probs = exp_logits / exp_logits.sum(axis=-1, keepdims=True)

    # Handle any remaining NaN
    probs = np.nan_to_num(probs, nan=0.5)

    precision, recall, f1, _ = precision_recall_fscore_support(
        labels, preds, average="binary", zero_division=0,
    )
    accuracy = accuracy_score(labels, preds)

    try:
        auc = roc_auc_score(labels, probs[:, 1])
    except (ValueError, IndexError):
        auc = 0.0

    return {
        "accuracy": accuracy,
        "precision": precision,
        "recall": recall,
        "f1": f1,
        "roc_auc": auc,
    }


def print_evaluation_report(labels, preds, probs=None, title="Evaluation Report"):
    """Print a detailed classification report."""
    print(f"\n{'='*60}")
    print(f"  {title}")
    print(f"{'='*60}")

    target_names = ["non-conspiracy", "conspiracy"]
    print(classification_report(labels, preds, target_names=target_names))

    print("Confusion Matrix:")
    cm = confusion_matrix(labels, preds)
    print(f"  {'':>18} Predicted")
    print(f"  {'':>18} {'non-consp':>10} {'conspiracy':>10}")
    print(f"  Actual non-consp {cm[0][0]:>10} {cm[0][1]:>10}")
    print(f"  Actual conspiracy{cm[1][0]:>10} {cm[1][1]:>10}")

    if probs is not None:
        try:
            auc = roc_auc_score(labels, probs)
            print(f"\n  ROC-AUC: {auc:.4f}")
        except ValueError:
            pass
    print()