"""評価指標のテスト（SPEC_02「評価指標」「交差検証の実行」の「評価の出し方」）.

- macro F1・accuracy・クラス別 precision / recall / F1・混同行列
  （クラス ID 順固定．SPEC_01「クラス定義」）
- 部分集合 `all`（全程度）・`common`（strong と normal）・`mild`
- mild の予測の内訳（淡紅・元のクラス・その他と予測された割合．全体とクラス別）
- fold 間の平均 ± 標準偏差・混同行列の全 fold 合計

期待値は手計算した既知の予測・正解から求める．
"""

from __future__ import annotations

import math

import numpy as np
import pytest

from tias_judge_model.data.classes import CLASS_SLUGS, label_to_id
from tias_judge_model.evaluation.metrics import (
    aggregate_folds,
    classification_metrics,
    macro_f1,
    mean_std,
    mild_prediction_breakdown,
    subset_metrics,
    sum_confusion_matrices,
)

PALE, LIGHT_RED, RED, DARK_RED, PURPLE = range(5)

# 正解 [淡白, 淡白, 淡紅, 淡紅, 紅, 紅]・予測 [淡白, 淡紅, 淡紅, 淡紅, 紅, 淡白]
Y_TRUE = np.array([PALE, PALE, LIGHT_RED, LIGHT_RED, RED, RED])
Y_PRED = np.array([PALE, LIGHT_RED, LIGHT_RED, LIGHT_RED, RED, PALE])
# 手計算: 淡白 P=1/2 R=1/2 F1=1/2・淡紅 P=2/3 R=1 F1=4/5・紅 P=1 R=1/2 F1=2/3
EXPECTED = {
    "pale": (0.5, 0.5, 0.5, 2),
    "light_red": (2 / 3, 1.0, 0.8, 2),
    "red": (1.0, 0.5, 2 / 3, 2),
}
EXPECTED_CM = [
    [1, 1, 0, 0, 0],
    [0, 2, 0, 0, 0],
    [1, 0, 1, 0, 0],
    [0, 0, 0, 0, 0],
    [0, 0, 0, 0, 0],
]


class TestClassOrder:
    def test_class_ids_are_fixed_as_in_spec(self) -> None:
        # SPEC_01「クラス定義」: 0 淡白・1 淡紅・2 紅・3 暗赤紅・4 紫
        assert CLASS_SLUGS == ("pale", "light_red", "red", "dark_red", "purple")
        assert [label_to_id(s) for s in CLASS_SLUGS] == [0, 1, 2, 3, 4]


class TestClassificationMetrics:
    def test_accuracy(self) -> None:
        m = classification_metrics(Y_TRUE, Y_PRED)
        assert m["accuracy"] == pytest.approx(4 / 6)
        assert m["n"] == 6

    def test_per_class_precision_recall_f1(self) -> None:
        m = classification_metrics(Y_TRUE, Y_PRED)
        for slug, (p, r, f1, support) in EXPECTED.items():
            got = m["per_class"][slug]
            assert got["precision"] == pytest.approx(p)
            assert got["recall"] == pytest.approx(r)
            assert got["f1"] == pytest.approx(f1)
            assert got["support"] == support

    def test_macro_f1_averages_classes_present_in_truth(self) -> None:
        expected = (0.5 + 0.8 + 2 / 3) / 3
        m = classification_metrics(Y_TRUE, Y_PRED)
        assert m["macro_f1"] == pytest.approx(expected)
        assert macro_f1(Y_TRUE, Y_PRED) == pytest.approx(expected)

    def test_macro_f1_perfect_prediction_is_one(self) -> None:
        y = np.array([0, 1, 2, 3, 4, 4])
        assert macro_f1(y, y) == pytest.approx(1.0)

    def test_confusion_matrix_rows_true_cols_pred_in_class_id_order(self) -> None:
        m = classification_metrics(Y_TRUE, Y_PRED)
        assert m["confusion_matrix"] == EXPECTED_CM

    def test_confusion_matrix_is_always_5x5(self) -> None:
        # 正解・予測に現れないクラスがあっても軸がぶれない
        m = classification_metrics(np.array([PURPLE]), np.array([PURPLE]))
        cm = np.array(m["confusion_matrix"])
        assert cm.shape == (5, 5)
        assert cm[PURPLE, PURPLE] == 1
        assert cm.sum() == 1

    def test_input_order_does_not_change_metrics(self) -> None:
        rng = np.random.default_rng(0)
        perm = rng.permutation(len(Y_TRUE))
        a = classification_metrics(Y_TRUE, Y_PRED)
        b = classification_metrics(Y_TRUE[perm], Y_PRED[perm])
        assert a == b


class TestMildPredictionBreakdown:
    # mild の正解 [紅, 紅, 紅, 暗赤紅, 紫, 淡白]・予測 [淡紅, 紅, 暗赤紅, 暗赤紅, 淡紅, 淡白]
    TRUE = np.array([RED, RED, RED, DARK_RED, PURPLE, PALE])
    PRED = np.array([LIGHT_RED, RED, DARK_RED, DARK_RED, LIGHT_RED, PALE])

    def test_overall_fractions(self) -> None:
        b = mild_prediction_breakdown(self.TRUE, self.PRED)["overall"]
        assert b["n"] == 6
        assert b["light_red"] == pytest.approx(2 / 6)
        assert b["original"] == pytest.approx(3 / 6)
        assert b["other"] == pytest.approx(1 / 6)

    def test_fractions_sum_to_one(self) -> None:
        b = mild_prediction_breakdown(self.TRUE, self.PRED)
        for entry in [b["overall"], *b["per_class"].values()]:
            assert entry["light_red"] + entry["original"] + entry["other"] == pytest.approx(1.0)

    def test_per_class_fractions(self) -> None:
        per_class = mild_prediction_breakdown(self.TRUE, self.PRED)["per_class"]
        red = per_class["red"]
        assert red["n"] == 3
        assert red["light_red"] == pytest.approx(1 / 3)
        assert red["original"] == pytest.approx(1 / 3)
        assert red["other"] == pytest.approx(1 / 3)
        assert per_class["dark_red"]["original"] == pytest.approx(1.0)
        assert per_class["purple"]["light_red"] == pytest.approx(1.0)
        assert per_class["pale"]["original"] == pytest.approx(1.0)


class TestSubsetMetrics:
    # 程度つきの held_out: strong 2・normal 2・mild 2
    TRUE = np.array([RED, RED, LIGHT_RED, LIGHT_RED, RED, PURPLE])
    PRED = np.array([RED, PALE, LIGHT_RED, LIGHT_RED, LIGHT_RED, PURPLE])
    SEV = ["strong", "strong", "normal", "normal", "mild", "mild"]

    def test_subset_sizes(self) -> None:
        s = subset_metrics(self.TRUE, self.PRED, self.SEV)
        assert s["all"]["n"] == 6
        assert s["common"]["n"] == 4
        assert s["mild"]["n"] == 2

    def test_common_is_strong_and_normal_only(self) -> None:
        s = subset_metrics(self.TRUE, self.PRED, self.SEV)
        expected = classification_metrics(self.TRUE[:4], self.PRED[:4])
        assert s["common"]["accuracy"] == pytest.approx(expected["accuracy"])
        assert s["common"]["confusion_matrix"] == expected["confusion_matrix"]

    def test_all_uses_every_image(self) -> None:
        s = subset_metrics(self.TRUE, self.PRED, self.SEV)
        assert s["all"]["accuracy"] == pytest.approx(4 / 6)
        assert np.array(s["all"]["confusion_matrix"]).sum() == 6

    def test_mild_has_prediction_breakdown(self) -> None:
        s = subset_metrics(self.TRUE, self.PRED, self.SEV)
        overall = s["mild"]["prediction_breakdown"]["overall"]
        assert overall["n"] == 2
        assert overall["light_red"] == pytest.approx(0.5)
        assert overall["original"] == pytest.approx(0.5)
        assert overall["other"] == pytest.approx(0.0)
        assert s["mild"]["accuracy"] == pytest.approx(0.5)

    def test_empty_mild_subset_does_not_crash(self) -> None:
        s = subset_metrics(self.TRUE[:4], self.PRED[:4], self.SEV[:4])
        assert s["common"]["n"] == 4
        assert s["all"]["n"] == 4
        assert not s["mild"]


class TestFoldAggregation:
    def test_mean_std_uses_sample_std(self) -> None:
        r = mean_std([0.5, 0.7, 0.9])
        assert r["mean"] == pytest.approx(0.7)
        assert r["std"] == pytest.approx(0.2)  # ddof=1

    def test_mean_std_single_value(self) -> None:
        r = mean_std([0.4])
        assert r["mean"] == pytest.approx(0.4)
        assert r["std"] == pytest.approx(0.0)

    def test_aggregate_folds_averages_macro_f1_and_accuracy(self) -> None:
        f1 = classification_metrics(np.array([0, 1]), np.array([0, 1]))  # acc 1.0
        f2 = classification_metrics(np.array([0, 1]), np.array([1, 1]))  # acc 0.5
        agg = aggregate_folds([f1, f2])
        assert agg["accuracy"]["mean"] == pytest.approx(0.75)
        assert agg["accuracy"]["std"] == pytest.approx(math.sqrt(0.125))
        assert agg["macro_f1"]["mean"] == pytest.approx((1.0 + f2["macro_f1"]) / 2)
        assert agg["per_class"]["light_red"]["recall"]["mean"] == pytest.approx(1.0)

    def test_sum_confusion_matrices(self) -> None:
        a = classification_metrics(Y_TRUE, Y_PRED)["confusion_matrix"]
        b = classification_metrics(np.array([PURPLE, DARK_RED]), np.array([PURPLE, RED]))[
            "confusion_matrix"
        ]
        total = np.array(sum_confusion_matrices([a, b]))
        assert total.shape == (5, 5)
        assert total.sum() == 8
        assert total[PURPLE, PURPLE] == 1
        assert total[DARK_RED, RED] == 1
        assert total[LIGHT_RED, LIGHT_RED] == 2