from __future__ import annotations

from typing import Any

import numpy as np
from sklearn.metrics import accuracy_score, confusion_matrix, f1_score, precision_recall_fscore_support

from tiny_router.calibration import apply_temperature_scaling
from tiny_router.constants import HEAD_LABELS


def softmax(logits: np.ndarray) -> np.ndarray:
    logits = logits - np.max(logits, axis=-1, keepdims=True)
    exp = np.exp(logits)
    return exp / np.sum(exp, axis=-1, keepdims=True)


def confidence_calibration(
    confidences: np.ndarray,
    correct: np.ndarray,
    bins: int = 10,
) -> dict[str, Any]:
    if confidences.size == 0:
        return {"ece": 0.0, "bins": []}
    edges = np.linspace(0.0, 1.0, bins + 1)
    rows = []
    ece = 0.0
    for idx in range(bins):
        lower = edges[idx]
        upper = edges[idx + 1]
        if idx == bins - 1:
            mask = (confidences >= lower) & (confidences <= upper)
        else:
            mask = (confidences >= lower) & (confidences < upper)
        count = int(mask.sum())
        if count == 0:
            continue
        avg_conf = float(confidences[mask].mean())
        avg_acc = float(correct[mask].mean())
        gap = abs(avg_conf - avg_acc)
        ece += gap * (count / confidences.size)
        rows.append(
            {
                "range": [round(float(lower), 3), round(float(upper), 3)],
                "count": count,
                "avg_confidence": round(avg_conf, 4),
                "accuracy": round(avg_acc, 4),
            }
        )
    return {"ece": round(float(ece), 6), "bins": rows}


def compute_head_metrics(
    y_true: np.ndarray,
    y_pred: np.ndarray,
    labels: list[str],
) -> dict[str, Any]:
    accuracy = accuracy_score(y_true, y_pred)
    macro_f1 = f1_score(y_true, y_pred, average="macro", zero_division=0)
    precision, recall, f1, support = precision_recall_fscore_support(
        y_true, y_pred, labels=list(range(len(labels))), zero_division=0
    )
    per_label = {}
    for idx, label in enumerate(labels):
        per_label[label] = {
            "precision": round(float(precision[idx]), 4),
            "recall": round(float(recall[idx]), 4),
            "f1": round(float(f1[idx]), 4),
            "support": int(support[idx]),
        }
    return {
        "accuracy": round(float(accuracy), 4),
        "macro_f1": round(float(macro_f1), 4),
        "per_label": per_label,
        "confusion_matrix": confusion_matrix(
            y_true, y_pred, labels=list(range(len(labels)))
        ).tolist(),
    }


def evaluate_multitask(
    logits_by_head: dict[str, np.ndarray],
    labels_by_head: dict[str, np.ndarray],
    threshold: float = 0.8,
    temperatures: dict[str, float] | None = None,
) -> dict[str, Any]:
    metrics: dict[str, Any] = {"per_head": {}}
    predictions: dict[str, np.ndarray] = {}
    confidences: dict[str, np.ndarray] = {}
    logits_by_head = apply_temperature_scaling(logits_by_head, temperatures)

    for head, labels in HEAD_LABELS.items():
        probs = softmax(logits_by_head[head])
        preds = probs.argmax(axis=-1)
        conf = probs.max(axis=-1)
        predictions[head] = preds
        confidences[head] = conf
        metrics["per_head"][head] = compute_head_metrics(labels_by_head[head], preds, labels)

    stacked_true = np.stack([labels_by_head[head] for head in HEAD_LABELS], axis=1)
    stacked_pred = np.stack([predictions[head] for head in HEAD_LABELS], axis=1)
    exact = (stacked_true == stacked_pred).all(axis=1)
    overall_confidence = np.stack([confidences[head] for head in HEAD_LABELS], axis=1).mean(axis=1)
    safe_mask = overall_confidence >= threshold
    safe_accuracy = float(exact[safe_mask].mean()) if safe_mask.any() else 0.0

    macro_average = float(
        np.mean([metrics["per_head"][head]["macro_f1"] for head in HEAD_LABELS])
    )
    metrics["overall"] = {
        "exact_match": round(float(exact.mean()), 4),
        "macro_average_f1": round(macro_average, 4),
        "automation_safe_accuracy": round(safe_accuracy, 4),
        "automation_safe_coverage": round(float(safe_mask.mean()), 4),
        "confidence_threshold": threshold,
        "confidence_calibration": confidence_calibration(overall_confidence, exact.astype(float)),
    }
    if temperatures:
        metrics["temperature_scaling"] = {
            "method": "per_head_temperature_scaling",
            "per_head": {head: round(float(temp), 6) for head, temp in temperatures.items()},
        }
    return metrics