"""評価指標（SPEC_02「評価指標」「交差検証の実行」）.

すべて入力 → 出力の純関数．クラスの軸は常に `classes.py` の固定 ID 順（混同行列の軸をぶらさない）．
出力は集計値だけであり，画像単位の値・患者 ID は含まない．
"""

from __future__ import annotations

from collections.abc import Mapping, Sequence
from typing import Any

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

from tias_judge_model.data.classes import CLASS_SLUGS, NUM_CLASSES, label_to_id

# 淡紅（正常）のクラス ID．mild の予測の内訳で「淡紅と予測された割合」に使う．
LIGHT_RED_ID: int = label_to_id("light_red")

# 評価の部分集合と含める程度（SPEC_02「評価の出し方」）．
SUBSETS: dict[str, tuple[str, ...]] = {
    "all": ("strong", "mild", "normal"),
    "common": ("strong", "normal"),
    "mild": ("mild",),
}

# fold 間の平均・標準偏差を取らないキー（件数・行列）．
_NON_AVERAGED_KEYS = frozenset({"n", "support", "confusion_matrix"})


def macro_f1(y_true: np.ndarray, y_pred: np.ndarray) -> float:
    """macro F1 を返す（`y_true` に現れるクラスで平均する）.

    `y_true` に現れないクラス（例: mild の部分集合の淡紅）を平均に含めると，そのクラスへの予測が
    それだけで F1 = 0 として平均を下げるため，正解側に現れるクラスだけで平均する．

    Args:
        y_true: `(N,)` の正解クラス ID．
        y_pred: `(N,)` の予測クラス ID．

    Returns:
        macro F1（N = 0 なら 0.0）．
    """
    if len(y_true) == 0:
        return 0.0
    labels = np.unique(y_true)
    return float(f1_score(y_true, y_pred, labels=labels, average="macro", zero_division=0))


def classification_metrics(y_true: np.ndarray, y_pred: np.ndarray) -> dict[str, Any]:
    """分類指標（macro F1・accuracy・クラス別 precision / recall / F1・混同行列）を返す.

    Args:
        y_true: `(N,)` の正解クラス ID（0〜4）．
        y_pred: `(N,)` の予測クラス ID（0〜4）．

    Returns:
        `{"n", "macro_f1", "accuracy", "per_class": {slug: {precision, recall, f1, support}},
        "confusion_matrix": [[...]]}`．混同行列は行が正解・列が予測でクラス ID 順（5×5）．
        クラス別は `y_true` に現れるクラスだけを含む．
    """
    y_true = np.asarray(y_true, dtype=np.int64)
    y_pred = np.asarray(y_pred, dtype=np.int64)
    all_labels = list(range(NUM_CLASSES))
    cm = confusion_matrix(y_true, y_pred, labels=all_labels)
    precision, recall, f1, support = precision_recall_fscore_support(
        y_true, y_pred, labels=all_labels, zero_division=0
    )
    present = set(np.unique(y_true).tolist())
    per_class = {
        CLASS_SLUGS[c]: {
            "precision": float(precision[c]),
            "recall": float(recall[c]),
            "f1": float(f1[c]),
            "support": int(support[c]),
        }
        for c in all_labels
        if c in present
    }
    return {
        "n": int(len(y_true)),
        "macro_f1": macro_f1(y_true, y_pred),
        "accuracy": float(np.mean(y_true == y_pred)) if len(y_true) else 0.0,
        "per_class": per_class,
        "confusion_matrix": cm.astype(int).tolist(),
    }


def mild_prediction_breakdown(y_true: np.ndarray, y_pred: np.ndarray) -> dict[str, Any]:
    """mild の画像の予測の内訳（淡紅・元のクラス・その他と予測された割合）を返す.

    淡紅と予測することを誤りと断定せず，割合として報告する（SPEC_02「程度（ややあり）の扱い」）．
    淡紅には mild が無いため，元のクラスと淡紅は重ならない．

    Args:
        y_true: `(N,)` の mild の画像の正解クラス ID．
        y_pred: `(N,)` の予測クラス ID．

    Returns:
        `{"overall": {...}, "per_class": {slug: {...}}}`．各値は
        `{"n", "light_red", "original", "other"}`（割合．N = 0 なら None）．
        クラス別は `y_true` に現れるクラスだけを含む．
    """
    y_true = np.asarray(y_true, dtype=np.int64)
    y_pred = np.asarray(y_pred, dtype=np.int64)
    per_class = {
        CLASS_SLUGS[c]: _breakdown(y_true[y_true == c], y_pred[y_true == c])
        for c in range(NUM_CLASSES)
        if np.any(y_true == c)
    }
    return {"overall": _breakdown(y_true, y_pred), "per_class": per_class}


def _breakdown(y_true: np.ndarray, y_pred: np.ndarray) -> dict[str, Any]:
    n = len(y_true)
    if n == 0:
        return {"n": 0, "light_red": None, "original": None, "other": None}
    to_light_red = y_pred == LIGHT_RED_ID
    to_original = (y_pred == y_true) & ~to_light_red
    other = ~to_light_red & ~to_original
    return {
        "n": int(n),
        "light_red": float(np.mean(to_light_red)),
        "original": float(np.mean(to_original)),
        "other": float(np.mean(other)),
    }


def subset_metrics(
    y_true: np.ndarray, y_pred: np.ndarray, severities: Sequence[str]
) -> dict[str, Any]:
    """部分集合 `all`・`common`・`mild` ごとの指標を返す（SPEC_02「評価の出し方」）.

    Args:
        y_true: `(N,)` の正解クラス ID．
        y_pred: `(N,)` の予測クラス ID．
        severities: 長さ N の各画像の程度（`strong`・`mild`・`normal`）．

    Returns:
        部分集合名 → `classification_metrics` の結果（該当 0 件なら None）．
        `mild` には予測の内訳 `prediction_breakdown` を加える．
    """
    y_true = np.asarray(y_true, dtype=np.int64)
    y_pred = np.asarray(y_pred, dtype=np.int64)
    sev = np.asarray(severities)
    result: dict[str, Any] = {}
    for name, included in SUBSETS.items():
        mask = np.isin(sev, included)
        if not np.any(mask):
            result[name] = None
            continue
        metrics = classification_metrics(y_true[mask], y_pred[mask])
        if name == "mild":
            metrics["prediction_breakdown"] = mild_prediction_breakdown(y_true[mask], y_pred[mask])
        result[name] = metrics
    return result


def mean_std(values: Sequence[float]) -> dict[str, float | int | None]:
    """値の平均と標準偏差（不偏．1 件なら 0）を返す.

    Args:
        values: fold ごとの値．

    Returns:
        `{"mean", "std", "n"}`（0 件なら mean・std は None）．
    """
    arr = np.asarray(values, dtype=np.float64)
    if arr.size == 0:
        return {"mean": None, "std": None, "n": 0}
    std = float(np.std(arr, ddof=1)) if arr.size > 1 else 0.0
    return {"mean": float(np.mean(arr)), "std": std, "n": int(arr.size)}


def aggregate_folds(fold_values: Sequence[Mapping[str, Any] | None]) -> dict[str, Any] | None:
    """fold ごとの入れ子の指標（同じ形の dict）から，数値の葉ごとに平均 ± 標準偏差を作る.

    件数（`n`・`support`）と混同行列は平均しない．ある fold に無いキー（例: その fold の
    mild に現れないクラス）は，そのキーを持つ fold だけで集計する（`n` に fold 数が入る）．

    Args:
        fold_values: fold ごとの指標（None の fold は除く）．

    Returns:
        同じ形の dict で，数値の葉を `{"mean", "std", "n"}` に置き換えたもの．
        有効な fold が無ければ None．
    """
    present = [v for v in fold_values if v is not None]
    if not present:
        return None
    keys: list[str] = []
    for v in present:
        keys.extend(k for k in v if k not in keys)
    result: dict[str, Any] = {}
    for key in keys:
        if key in _NON_AVERAGED_KEYS:
            continue
        items = [v.get(key) for v in present]
        items = [x for x in items if x is not None]
        if not items:
            continue
        if all(isinstance(x, Mapping) for x in items):
            nested = aggregate_folds(items)
            if nested:
                result[key] = nested
        elif all(isinstance(x, int | float) and not isinstance(x, bool) for x in items):
            result[key] = mean_std([float(x) for x in items])
    return result


def sum_confusion_matrices(matrices: Sequence[Sequence[Sequence[int]]]) -> list[list[int]]:
    """混同行列（同じ形）の合計を返す.

    Args:
        matrices: 混同行列のリスト（各 5×5）．

    Returns:
        合計の混同行列（5×5．空なら 0 行列）．
    """
    total = np.zeros((NUM_CLASSES, NUM_CLASSES), dtype=np.int64)
    for m in matrices:
        total += np.asarray(m, dtype=np.int64)
    return total.astype(int).tolist()