"""クラス不均衡の対策の計算のテスト（SPEC_02「クラス不均衡」→「対策の設定」）.

- `inverse` は w_c ∝ N / (K・n_c)，`sqrt_inverse` はその平方根に比例する
- 重みは Σ_c (n_c / N)・w_c = 1 に正規化される（学習データでの平均が 1）
- クラス ID 順（淡白・淡紅・紅・暗赤紅・紫）で，少ないクラスほど重みが大きい
- 学習データに 1 枚も無いクラスがあれば停止する（メッセージに個人情報を出さない）
- バランス・サンプリングの抜き出し重みは 1 / n_c に比例し，各クラスがほぼ等確率で選ばれる
- 対策の組み立て（`make_class_balance`）は渡された学習ラベルだけから枚数を数える
"""

from __future__ import annotations

import numpy as np
import pytest

from tias_judge_model.data.class_balance import (
    CLASS_WEIGHTING_METHODS,
    balanced_sample_weights,
    class_draw_probabilities,
    compute_class_weights,
    count_classes,
)
from tias_judge_model.data.classes import CLASS_SLUGS, NUM_CLASSES
from tias_judge_model.schema import TrainConfig
from tias_judge_model.training.cross_validation import make_class_balance

# 実データに近い不均衡（最大 / 最小 ≒ 3.7 倍）．クラス ID 順．
COUNTS = np.array([104, 380, 250, 200, 151])


def _labels(counts: np.ndarray, *, shuffle_seed: int | None = 0) -> np.ndarray:
    labels = np.repeat(np.arange(NUM_CLASSES), counts)
    if shuffle_seed is not None:
        np.random.default_rng(shuffle_seed).shuffle(labels)
    return labels


def _expected_weights(counts: np.ndarray, method: str) -> np.ndarray:
    n = counts.astype(np.float64)
    total = n.sum()
    raw = total / (NUM_CLASSES * n)
    if method == "sqrt_inverse":
        raw = np.sqrt(raw)
    return raw / np.sum(n / total * raw)


def test_class_order_is_fixed() -> None:
    assert CLASS_SLUGS == ("pale", "light_red", "red", "dark_red", "purple")


def test_weighting_methods_are_the_spec_values() -> None:
    assert set(CLASS_WEIGHTING_METHODS) == {"inverse", "sqrt_inverse"}


def test_count_classes_in_class_id_order() -> None:
    counts = count_classes(_labels(COUNTS))
    assert counts.tolist() == COUNTS.tolist()


class TestComputeClassWeights:
    @pytest.mark.parametrize("method", ["inverse", "sqrt_inverse"])
    def test_matches_formula(self, method: str) -> None:
        w = compute_class_weights(_labels(COUNTS), method)
        assert w.shape == (NUM_CLASSES,)
        np.testing.assert_allclose(w, _expected_weights(COUNTS, method), rtol=1e-12)

    def test_inverse_proportional_to_inverse_count(self) -> None:
        w = compute_class_weights(_labels(COUNTS), "inverse")
        # w_c・n_c が全クラスで一定（w_c ∝ 1 / n_c）
        prod = w * COUNTS
        np.testing.assert_allclose(prod, prod[0], rtol=1e-12)

    def test_sqrt_inverse_is_proportional_to_sqrt_of_inverse(self) -> None:
        inv = compute_class_weights(_labels(COUNTS), "inverse")
        sq = compute_class_weights(_labels(COUNTS), "sqrt_inverse")
        ratio = sq / np.sqrt(inv)
        np.testing.assert_allclose(ratio, ratio[0], rtol=1e-12)

    @pytest.mark.parametrize("method", ["inverse", "sqrt_inverse"])
    def test_normalized_so_training_mean_is_one(self, method: str) -> None:
        w = compute_class_weights(_labels(COUNTS), method)
        n = COUNTS.astype(np.float64)
        assert float(np.sum(n / n.sum() * w)) == pytest.approx(1.0, abs=1e-12)
        # 画像ごとの重みの平均としても 1
        labels = _labels(COUNTS)
        assert float(w[labels].mean()) == pytest.approx(1.0, abs=1e-12)

    def test_inverse_known_values(self) -> None:
        # n = [1, 3] ではなく 5 クラスで手計算: n = [10, 20, 20, 25, 25]（N = 100）
        counts = np.array([10, 20, 20, 25, 25])
        w = compute_class_weights(_labels(counts), "inverse")
        # 正規化後は w_c = 1 / (K・p_c)（Σ p_c・w_c = K・(1/K) = 1 がそのまま成り立つ）
        expected = 100 / (5 * counts)
        np.testing.assert_allclose(w, expected, rtol=1e-12)
        assert w.tolist() == pytest.approx([2.0, 1.0, 1.0, 0.8, 0.8])

    @pytest.mark.parametrize("method", ["inverse", "sqrt_inverse"])
    def test_rarer_class_gets_larger_weight(self, method: str) -> None:
        w = compute_class_weights(_labels(COUNTS), method)
        order_by_count = np.argsort(COUNTS)  # 少ない順
        assert np.all(np.diff(w[order_by_count]) < 0)
        assert int(np.argmax(w)) == 0  # pale（104 枚）が最小
        assert int(np.argmin(w)) == 1  # light_red（380 枚）が最大

    @pytest.mark.parametrize("method", ["inverse", "sqrt_inverse"])
    def test_equal_counts_give_unit_weights(self, method: str) -> None:
        w = compute_class_weights(_labels(np.array([7] * NUM_CLASSES)), method)
        np.testing.assert_allclose(w, 1.0, rtol=1e-12)

    def test_sqrt_inverse_is_milder_than_inverse(self) -> None:
        inv = compute_class_weights(_labels(COUNTS), "inverse")
        sq = compute_class_weights(_labels(COUNTS), "sqrt_inverse")
        assert sq.max() / sq.min() < inv.max() / inv.min()
        assert sq.max() / sq.min() == pytest.approx(np.sqrt(inv.max() / inv.min()))

    @pytest.mark.parametrize("method", ["inverse", "sqrt_inverse"])
    def test_does_not_depend_on_label_order(self, method: str) -> None:
        a = compute_class_weights(_labels(COUNTS, shuffle_seed=None), method)
        b = compute_class_weights(_labels(COUNTS, shuffle_seed=1), method)
        np.testing.assert_allclose(a, b, rtol=1e-12)

    @pytest.mark.parametrize("method", ["inverse", "sqrt_inverse"])
    @pytest.mark.parametrize("missing", range(NUM_CLASSES))
    def test_missing_class_raises(self, method: str, missing: int) -> None:
        counts = COUNTS.copy()
        counts[missing] = 0
        with pytest.raises(ValueError):
            compute_class_weights(_labels(counts), method)

    @pytest.mark.parametrize("method", ["none", "balanced", "INVERSE", ""])
    def test_unknown_method_raises(self, method: str) -> None:
        with pytest.raises(ValueError):
            compute_class_weights(_labels(COUNTS), method)


class TestBalancedSampling:
    def test_sample_weight_proportional_to_inverse_class_count(self) -> None:
        labels = _labels(COUNTS)
        sw = balanced_sample_weights(labels)
        assert sw.shape == labels.shape
        # 各画像の重み × その画像のクラスの枚数 が一定（∝ 1 / n_c）
        prod = sw * COUNTS[labels]
        np.testing.assert_allclose(prod, prod[0], rtol=1e-12)

    def test_each_class_equally_likely_per_draw(self) -> None:
        labels = _labels(COUNTS)
        probs = class_draw_probabilities(labels, balanced_sample_weights(labels))
        assert probs.shape == (NUM_CLASSES,)
        np.testing.assert_allclose(probs, 1.0 / NUM_CLASSES, rtol=1e-12)

    def test_uniform_weights_give_class_frequencies(self) -> None:
        labels = _labels(COUNTS)
        probs = class_draw_probabilities(labels, np.ones(len(labels)))
        np.testing.assert_allclose(probs, COUNTS / COUNTS.sum(), rtol=1e-12)

    @pytest.mark.parametrize("missing", range(NUM_CLASSES))
    def test_missing_class_raises(self, missing: int) -> None:
        counts = COUNTS.copy()
        counts[missing] = 0
        with pytest.raises(ValueError):
            balanced_sample_weights(_labels(counts))


def test_missing_class_message_has_no_personal_info() -> None:
    # メッセージにファイル名・患者 ID が入らないこと
    counts = COUNTS.copy()
    counts[4] = 0
    with pytest.raises(ValueError) as exc:
        compute_class_weights(_labels(counts), "inverse")
    msg = str(exc.value)
    assert msg
    assert ".png" not in msg
    assert "patient" not in msg.lower()


class TestMakeClassBalance:
    def _config(self, **train: str) -> TrainConfig:
        return TrainConfig.model_validate(train)

    def test_defaults_have_no_countermeasure(self) -> None:
        labels = _labels(COUNTS)
        balance, class_weights, sample_weights = make_class_balance(labels, self._config())
        assert class_weights is None
        assert sample_weights is None
        assert balance.class_weighting == "none"
        assert balance.sampler == "random"
        assert list(balance.class_counts) == COUNTS.tolist()
        assert balance.class_weights is None
        assert balance.class_draw_probabilities is None

    def test_defaults_do_not_stop_on_missing_class(self) -> None:
        # 対策なしは従来どおりの学習（重みが要らないため停止しない）
        counts = COUNTS.copy()
        counts[4] = 0
        balance, cw, sw = make_class_balance(_labels(counts), self._config())
        assert cw is None and sw is None
        assert balance.class_counts[4] == 0

    @pytest.mark.parametrize("method", ["inverse", "sqrt_inverse"])
    def test_class_weighting(self, method: str) -> None:
        labels = _labels(COUNTS)
        balance, cw, sw = make_class_balance(labels, self._config(class_weighting=method))
        assert sw is None
        assert cw is not None
        np.testing.assert_allclose(cw, _expected_weights(COUNTS, method), rtol=1e-12)
        assert balance.class_weighting == method
        np.testing.assert_allclose(balance.class_weights, cw, rtol=1e-12)
        assert balance.class_draw_probabilities is None

    def test_balanced_sampler(self) -> None:
        labels = _labels(COUNTS)
        balance, cw, sw = make_class_balance(labels, self._config(sampler="balanced"))
        assert cw is None
        assert sw is not None and sw.shape == labels.shape
        assert balance.sampler == "balanced"
        assert balance.class_weights is None
        np.testing.assert_allclose(balance.class_draw_probabilities, 1 / NUM_CLASSES, rtol=1e-12)

    @pytest.mark.parametrize(
        "train", [{"class_weighting": "inverse"}, {"sampler": "balanced"}], ids=str
    )
    def test_missing_class_with_countermeasure_raises(self, train: dict) -> None:
        counts = COUNTS.copy()
        counts[0] = 0
        with pytest.raises(ValueError):
            make_class_balance(_labels(counts), self._config(**train))