"""損失のクラス重みのテスト（SPEC_02「クラス不均衡」→「対策の設定」）.

- 重みは**学習の損失にだけ**かかり，検証の損失は重みなし
- 正規化（Σ_c (n_c / N)・w_c = 1）により，重み付きでも損失の大きさの尺度が対策なしと揃う
  （全クラス同数なら重み付きと重みなしの損失が一致する）
- 重み付きの学習損失は Σ_i w_{y_i}・CE_i / B（件数で割る）．重みの和で割ると正規化の意味が無くなる

出力が入力に依らない固定ロジットのモデルで，損失の値そのものを手計算と比べる
（重みの勾配が 0 なので学習しても値が変わらない）．
"""

from __future__ import annotations

import sys

import numpy as np
import pytest
import torch
from torch import nn
from torch.utils.data import DataLoader, TensorDataset

from tias_judge_model.data.class_balance import compute_class_weights
from tias_judge_model.data.classes import NUM_CLASSES
from tias_judge_model.schema import TrainConfig
from tias_judge_model.training.loop import train_model

CPU = torch.device("cpu")
# 固定ロジット（クラスごとに CE が異なるように非一様にする）
LOGITS = torch.tensor([2.0, 0.5, -1.0, 0.0, 1.0])


class _ConstLogits(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.dummy = nn.Parameter(torch.zeros(NUM_CLASSES))
        self.register_buffer("logits", LOGITS.clone())

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.logits.expand(x.shape[0], -1) + 0.0 * self.dummy


def _loader(labels: np.ndarray) -> DataLoader:
    images = torch.zeros((len(labels), 3, 8, 8), dtype=torch.uint8)
    ds = TensorDataset(images, torch.as_tensor(labels, dtype=torch.int64))
    # 1 バッチに全件を入れる（バッチ平均 = 学習データ全体の平均）
    return DataLoader(ds, batch_size=len(labels), shuffle=False)


def _config() -> TrainConfig:
    return TrainConfig.model_validate(
        {
            "max_epochs": 1,
            "warmup_epochs": 0,
            "weight_decay": 0.0,
            "num_workers": 0,
            "amp": False,
            "batch_size": 1024,
        }
    )


def _ce_per_class() -> np.ndarray:
    return (-torch.log_softmax(LOGITS, dim=0)).numpy().astype(np.float64)


def _run(labels: np.ndarray, class_weights: torch.Tensor | None) -> tuple[float, float]:
    torch.manual_seed(0)
    result = train_model(
        _ConstLogits(),
        _loader(labels),
        _loader(labels),
        _config(),
        CPU,
        class_weights=class_weights,
    )
    rec = result.history[0]
    return rec.train_loss, rec.val_loss


def _labels(counts: list[int]) -> np.ndarray:
    return np.repeat(np.arange(NUM_CLASSES), counts)


IMBALANCED = [2, 12, 6, 4, 3]


def test_unweighted_losses_are_plain_cross_entropy() -> None:
    labels = _labels(IMBALANCED)
    train_loss, val_loss = _run(labels, None)
    expected = float(_ce_per_class()[labels].mean())
    assert train_loss == pytest.approx(expected, rel=1e-5)
    assert val_loss == pytest.approx(expected, rel=1e-5)


@pytest.mark.parametrize("method", ["inverse", "sqrt_inverse"])
def test_weighted_train_loss_on_full_training_data(method: str) -> None:
    labels = _labels(IMBALANCED)
    w = compute_class_weights(labels, method)
    train_loss, _ = _run(labels, torch.from_numpy(w))
    ce = _ce_per_class()[labels]
    expected = float(np.sum(w[labels] * ce) / len(labels))
    assert train_loss == pytest.approx(expected, rel=1e-5)
    # 少ないクラスが重く効くので，重みなしとは値が変わる
    assert train_loss != pytest.approx(float(ce.mean()), rel=1e-3)


def test_weighted_train_loss_divides_by_batch_count_not_weight_sum() -> None:
    # バッチ内の重みの和が件数と異なる場合（ミニバッチのクラス比が学習データ全体と違う場合）でも，
    # 損失は Σ_i w_{y_i}・CE_i / B．重みの和で割ると重みの正規化が意味を持たなくなる
    labels = _labels(IMBALANCED)
    w = np.array([3.0, 0.5, 1.0, 1.5, 2.0])
    train_loss, _ = _run(labels, torch.from_numpy(w))
    ce = _ce_per_class()[labels]
    by_count = float(np.sum(w[labels] * ce) / len(labels))
    by_weight_sum = float(np.sum(w[labels] * ce) / np.sum(w[labels]))
    assert by_count != pytest.approx(by_weight_sum, rel=1e-3)
    assert train_loss == pytest.approx(by_count, rel=1e-5)


@pytest.mark.parametrize("method", ["inverse", "sqrt_inverse"])
def test_validation_loss_is_unweighted(method: str) -> None:
    labels = _labels(IMBALANCED)
    w = compute_class_weights(labels, method)
    _, val_weighted = _run(labels, torch.from_numpy(w))
    _, val_plain = _run(labels, None)
    assert val_weighted == pytest.approx(val_plain, rel=1e-6)
    assert val_weighted == pytest.approx(float(_ce_per_class()[labels].mean()), rel=1e-5)


@pytest.mark.parametrize("method", ["inverse", "sqrt_inverse"])
def test_equal_class_counts_give_same_loss_as_unweighted(method: str) -> None:
    labels = _labels([5] * NUM_CLASSES)
    w = compute_class_weights(labels, method)
    weighted = _run(labels, torch.from_numpy(w))
    plain = _run(labels, None)
    assert weighted[0] == pytest.approx(plain[0], rel=1e-6)
    assert weighted[1] == pytest.approx(plain[1], rel=1e-6)


@pytest.mark.parametrize("method", ["inverse", "sqrt_inverse"])
def test_loss_scale_matches_unweighted_when_per_sample_loss_is_constant(
    method: str, monkeypatch: pytest.MonkeyPatch
) -> None:
    # 全画像の CE が等しい（一様ロジット）とき，正規化により重み付きの学習損失は重みなしと一致する
    monkeypatch.setattr(sys.modules[__name__], "LOGITS", torch.zeros(NUM_CLASSES), raising=True)
    labels = _labels(IMBALANCED)
    w = compute_class_weights(labels, method)
    weighted, _ = _run(labels, torch.from_numpy(w))
    plain, _ = _run(labels, None)
    assert plain == pytest.approx(float(np.log(NUM_CLASSES)), rel=1e-6)
    assert weighted == pytest.approx(plain, rel=1e-6)