"""交差検証でのクラス不均衡の対策のテスト（SPEC_02「クラス不均衡」→「対策の設定」）.

- クラスの枚数 n_c は各 fold の**学習用データ（`data.severities` 適用後）**から数える
  （検証用・held_out は使わない）．「ややあり除外」なら mild を除いた枚数で重みが決まる
- バランス・サンプリングは各画像を 1 / n_c に比例する確率で重複ありで抜き出し，1 回の抜き出しで
  各クラスがほぼ等確率．1 エポックの枚数は学習画像数と同じ．同じ seed で同じ抜き出し順
- 既定（`random`）はシャッフルで，各画像を 1 エポックに 1 回ずつ使う
"""

from __future__ import annotations

import importlib.util
import sys
from collections import Counter
from pathlib import Path
from types import ModuleType

import numpy as np
import pytest
import torch

from tias_judge_model.data.class_balance import balanced_sample_weights, compute_class_weights
from tias_judge_model.data.classes import CLASS_SLUGS, NUM_CLASSES
from tias_judge_model.data.manifest import read_manifest
from tias_judge_model.schema import ExperimentConfig, TrainConfig
from tias_judge_model.training.cross_validation import (
    CrossValidationResult,
    _make_loader,
    run_cross_validation,
)

SCRIPT_PATH = Path(__file__).resolve().parents[2] / "scripts" / "make_synthetic_dataset.py"
N_FOLDS = 5
CPU = torch.device("cpu")

# 不均衡な学習データ（クラス ID 順）
COUNTS = np.array([4, 24, 12, 8, 6])  # N = 54


def _labels() -> np.ndarray:
    labels = np.repeat(np.arange(NUM_CLASSES), COUNTS)
    np.random.default_rng(0).shuffle(labels)
    return labels


def _drawn_labels(loader) -> list[int]:
    return [int(y) for _, ys in loader for y in ys]


def _balanced_loader(seed: int, labels: np.ndarray):

    images = np.zeros((len(labels), 16, 16, 3), dtype=np.uint8)
    config = TrainConfig.model_validate(
        {"batch_size": 128, "num_workers": 0, "sampler": "balanced", "image_size": 16}
    )
    return _make_loader(
        images,
        labels,
        config,
        CPU,
        train=True,
        seed=seed,
        sample_weights=balanced_sample_weights(labels),
    )


class TestBalancedLoader:
    def test_epoch_size_equals_number_of_training_images(self) -> None:
        labels = _labels()
        drawn = _drawn_labels(_balanced_loader(0, labels))
        assert len(drawn) == len(labels)

    def test_draws_with_replacement(self) -> None:
        # 4 枚しか無い pale も約 54 / 5 ≒ 11 回抜き出される → 重複が必要
        labels = _labels()
        loader = _balanced_loader(0, labels)
        counts = Counter()
        for _ in range(5):
            counts.update(_drawn_labels(loader))
        assert counts[0] > 5 * COUNTS[0]

    def test_each_class_drawn_about_equally(self) -> None:
        labels = _labels()
        loader = _balanced_loader(0, labels)
        drawn: list[int] = []
        for _ in range(40):
            drawn.extend(_drawn_labels(loader))
        freq = np.bincount(drawn, minlength=NUM_CLASSES) / len(drawn)
        # 抜き出し 2160 回．一様（0.2）から ±0.04 以内（二項分布の標準偏差 ≒ 0.0086）
        np.testing.assert_allclose(freq, 1 / NUM_CLASSES, atol=0.04)
        # 元の比率（light_red 0.44・pale 0.07）とは明確に違う
        assert freq[1] < 0.3
        assert freq[0] > 0.12

    def test_same_seed_gives_same_draw_order(self) -> None:
        labels = _labels()
        a = [_drawn_labels(loader) for loader in [_balanced_loader(5, labels)] for _ in range(3)]
        b = [_drawn_labels(loader) for loader in [_balanced_loader(5, labels)] for _ in range(3)]
        assert a == b

    def test_different_seed_gives_different_draw_order(self) -> None:
        labels = _labels()
        assert _drawn_labels(_balanced_loader(1, labels)) != _drawn_labels(
            _balanced_loader(2, labels)
        )


def test_random_sampler_uses_each_image_once_per_epoch() -> None:
    labels = _labels()
    images = np.zeros((len(labels), 16, 16, 3), dtype=np.uint8)
    config = TrainConfig.model_validate({"batch_size": 128, "num_workers": 0, "image_size": 16})
    loader = _make_loader(images, labels, config, CPU, train=True, seed=0)
    drawn = _drawn_labels(loader)
    assert sorted(drawn) == sorted(labels.tolist())


# --- 交差検証を実際に回す（slow） ---


def _load_script() -> ModuleType:
    spec = importlib.util.spec_from_file_location("make_synthetic_dataset_cv_imb", SCRIPT_PATH)
    assert spec is not None and spec.loader is not None
    module = importlib.util.module_from_spec(spec)
    sys.modules[spec.name] = module
    spec.loader.exec_module(module)
    return module


def _config(root: Path, severities: list[str], **train: str) -> ExperimentConfig:
    return ExperimentConfig.model_validate(
        {
            "data": {"root": str(root), "severities": severities},
            "split": {"seed": 7, "n_folds": N_FOLDS},
            "model": {"name": "resnet18", "pretrained": False},
            "train": {
                "image_size": 32,
                "batch_size": 8,
                "max_epochs": 1,
                "warmup_epochs": 0,
                "inner_val_ratio": 0.2,
                "num_workers": 0,
                "amp": False,
                **train,
            },
        }
    )


@pytest.fixture(scope="module")
def imbalance_runs(
    tmp_path_factory: pytest.TempPathFactory,
) -> tuple[CrossValidationResult, CrossValidationResult]:
    root = tmp_path_factory.mktemp("cv_imb") / "dataset"
    _load_script().generate(root, per_dir=6, seed=0)
    records = read_manifest(root / "manifest.csv")
    weighted_without_mild = run_cross_validation(
        _config(root, ["strong", "normal"], class_weighting="inverse"),
        records,
        root.parent / "ck_w",
    )
    balanced_with_mild = run_cross_validation(
        _config(root, ["strong", "mild", "normal"], sampler="balanced"),
        records,
        root.parent / "ck_b",
    )
    return weighted_without_mild, balanced_with_mild


def _train_counts_by_class(group_counts: dict[str, int]) -> list[int]:
    out = [0] * NUM_CLASSES
    for group, n in group_counts.items():
        label = next(s for s in sorted(CLASS_SLUGS, key=len, reverse=True) if group.startswith(s))
        out[CLASS_SLUGS.index(label)] += n
    return out


@pytest.mark.slow
def test_class_counts_come_from_training_data_after_severities(imbalance_runs) -> None:
    weighted, balanced = imbalance_runs
    for result, mild_expected in ((weighted, False), (balanced, True)):
        for f in result.folds:
            b = f.class_balance
            assert b is not None
            train_groups = f.group_counts["train"]
            assert list(b.class_counts) == _train_counts_by_class(train_groups)
            has_mild = sum(n for g, n in train_groups.items() if g.endswith("_mild")) > 0
            assert has_mild is mild_expected
            # 検証用・held_out を含めた枚数ではない
            with_val = [
                a + c
                for a, c in zip(
                    _train_counts_by_class(train_groups),
                    _train_counts_by_class(f.group_counts["val"]),
                    strict=True,
                )
            ]
            assert list(b.class_counts) != with_val


@pytest.mark.slow
def test_weights_follow_formula_on_fold_training_counts(imbalance_runs) -> None:
    weighted, _ = imbalance_runs
    for f in weighted.folds:
        b = f.class_balance
        assert b.class_weighting == "inverse"
        assert b.sampler == "random"
        assert b.class_draw_probabilities is None
        labels = np.repeat(np.arange(NUM_CLASSES), b.class_counts)
        np.testing.assert_allclose(
            b.class_weights, compute_class_weights(labels, "inverse"), rtol=1e-9
        )
        n = np.array(b.class_counts, dtype=float)
        assert float(np.sum(n / n.sum() * np.array(b.class_weights))) == pytest.approx(1.0)


@pytest.mark.slow
def test_without_mild_weights_use_counts_excluding_mild(imbalance_runs) -> None:
    # 分割は data.severities に関わらず同じなので，ややあり込みの run の同じ fold の学習用枚数から
    # mild を引いたものが，ややあり除外の学習用枚数になる
    weighted, balanced = imbalance_runs
    for fw, fb in zip(weighted.folds, balanced.folds, strict=True):
        inclusive = _train_counts_by_class(fb.group_counts["train"])
        mild = _train_counts_by_class(
            {g: n for g, n in fb.group_counts["train"].items() if g.endswith("_mild")}
        )
        excluding = [a - m for a, m in zip(inclusive, mild, strict=True)]
        assert list(fw.class_balance.class_counts) == excluding
        assert list(fb.class_balance.class_counts) == inclusive
        # 合成データは各ディレクトリ同数なので，mild 込みなら淡紅（normal のみ）の重みが最大
        # になるが，
        # mild を除くと淡紅も他と同程度になる．重みは mild を除いた枚数の式に従う（上のテスト）
        w_inclusive = compute_class_weights(np.repeat(np.arange(NUM_CLASSES), inclusive), "inverse")
        assert int(np.argmax(w_inclusive)) == CLASS_SLUGS.index("light_red")
        assert not np.allclose(fw.class_balance.class_weights, w_inclusive)


@pytest.mark.slow
def test_balanced_sampler_records_equal_draw_probabilities(imbalance_runs) -> None:
    _, balanced = imbalance_runs
    for f in balanced.folds:
        b = f.class_balance
        assert b.sampler == "balanced"
        assert b.class_weighting == "none"
        assert b.class_weights is None
        np.testing.assert_allclose(b.class_draw_probabilities, 1 / NUM_CLASSES, rtol=1e-9)


@pytest.mark.slow
def test_training_runs_with_countermeasures(imbalance_runs) -> None:
    for result in imbalance_runs:
        assert len(result.folds) == N_FOLDS
        for f in result.folds:
            assert np.isfinite(f.history[0].train_loss)
            assert np.isfinite(f.history[0].val_loss)