"""患者単位の分割（単一分割・K-fold 交差検証）と分割後の集計.

分割は常に `patient_id` 単位で行い，同一患者が複数の分割に跨らないようにする（SPEC_02）．
1 患者は 1 つの label しか持たない（manifest の検証で保証）ため，患者の一覧を label で
層化して分割すれば「患者単位」と「クラス層化」を同時に満たせる．
"""

from __future__ import annotations

import logging
import math
from collections import Counter
from collections.abc import Collection, Iterable, Sequence
from dataclasses import dataclass

import numpy as np
from sklearn.model_selection import StratifiedKFold

from tias_judge_model.data.classes import CLASS_SLUGS
from tias_judge_model.data.manifest import ManifestRecord

logger = logging.getLogger(__name__)


@dataclass(frozen=True)
class SplitResult:
    """単一分割（train / val / test）の結果.

    各フィールドは患者 ID の昇順タプル．`runs/` への保存はサーバ上でのみ行うこと
    （患者 ID を含むため回収対象外．SPEC_01「回収ホワイトリスト」）．

    Attributes:
        train: train に属する患者 ID．
        val: val に属する患者 ID．
        test: test に属する患者 ID．
    """

    train: tuple[str, ...]
    val: tuple[str, ...]
    test: tuple[str, ...]

    def to_dict(self) -> dict[str, list[str]]:
        """保存用に `{"train": [...], "val": [...], "test": [...]}` を返す."""
        return {"train": list(self.train), "val": list(self.val), "test": list(self.test)}


@dataclass(frozen=True)
class Fold:
    """K-fold 交差検証の 1 fold.

    Attributes:
        index: fold 番号（0 始まり）．
        train: 学習に使う患者 ID（昇順）．
        held_out: この fold で評価に使う患者 ID（昇順）．
    """

    index: int
    train: tuple[str, ...]
    held_out: tuple[str, ...]

    def to_dict(self) -> dict[str, int | list[str]]:
        """保存用に `{"index": i, "train": [...], "held_out": [...]}` を返す."""
        return {"index": self.index, "train": list(self.train), "held_out": list(self.held_out)}


def patient_labels(records: Iterable[ManifestRecord]) -> dict[str, str]:
    """患者 ID → label の対応を作る.

    Args:
        records: manifest レコード．

    Returns:
        患者 ID をキー，label を値とする辞書．

    Raises:
        ValueError: 同一患者に異なる label が付いている場合．
    """
    mapping: dict[str, str] = {}
    for r in records:
        existing = mapping.setdefault(r.patient_id, r.label)
        if existing != r.label:
            raise ValueError("同一患者に異なる label が付いている（manifest の検証を通すこと）")
    return mapping


def split_train_val_test(
    records: Sequence[ManifestRecord], val_ratio: float, test_ratio: float, seed: int
) -> SplitResult:
    """患者単位・クラス層化で train / val / test に分割する.

    クラスごとに患者 ID を昇順に並べてから seed で並べ替え，先頭から test・val を
    `round(患者数 × 比率)` 人ずつ取り，残りを train とする．同じ入力と seed なら
    入力レコードの順序によらず同じ結果になる．

    Args:
        records: manifest レコード（検証済み）．
        val_ratio: val の比率（0 より大きく 1 未満）．
        test_ratio: test の比率（0 より大きく 1 未満）．
        seed: 乱数 seed．

    Returns:
        分割結果．

    Raises:
        ValueError: 比率が不正，レコードが空，同一患者に異なる label がある，
            またはあるクラスで train に患者が残らない場合．
    """
    if not (0.0 < val_ratio < 1.0 and 0.0 < test_ratio < 1.0 and val_ratio + test_ratio < 1.0):
        raise ValueError(
            f"分割比率が不正である（val_ratio={val_ratio}, test_ratio={test_ratio}）．"
            "それぞれ 0 より大きく，合計が 1 未満であること"
        )
    by_class = _patients_by_class(records)
    rng = np.random.default_rng(seed)

    train: list[str] = []
    val: list[str] = []
    test: list[str] = []
    for label in CLASS_SLUGS:
        patients = by_class.get(label)
        if not patients:
            continue
        shuffled = [patients[i] for i in rng.permutation(len(patients))]
        n_test = _round_half_up(len(patients) * test_ratio)
        n_val = _round_half_up(len(patients) * val_ratio)
        if n_test + n_val >= len(patients):
            raise ValueError(
                f"クラス {label} の患者 {len(patients)} 人では train に患者が残らない"
                f"（val {n_val} 人・test {n_test} 人）"
            )
        if n_test == 0 or n_val == 0:
            logger.warning(
                "クラス %s は患者 %d 人のため val %d 人・test %d 人になった",
                label,
                len(patients),
                n_val,
                n_test,
            )
        test.extend(shuffled[:n_test])
        val.extend(shuffled[n_test : n_test + n_val])
        train.extend(shuffled[n_test + n_val :])

    result = SplitResult(
        train=tuple(sorted(train)), val=tuple(sorted(val)), test=tuple(sorted(test))
    )
    logger.info(
        "単一分割（seed=%d）: 患者 train %d・val %d・test %d",
        seed,
        len(result.train),
        len(result.val),
        len(result.test),
    )
    return result


def make_kfold_splits(records: Sequence[ManifestRecord], n_folds: int, seed: int) -> list[Fold]:
    """患者単位・クラス層化の K-fold 交差検証の分割を作る.

    患者の一覧（患者 ID 昇順）を label で層化して `StratifiedKFold` にかける．
    各患者はちょうど 1 つの fold の `held_out` に入る．

    Args:
        records: manifest レコード（検証済み）．
        n_folds: fold 数（2 以上）．
        seed: 乱数 seed．

    Returns:
        fold のリスト（長さ `n_folds`）．

    Raises:
        ValueError: fold 数が不正，レコードが空，同一患者に異なる label がある，
            またはあるクラスの患者数が fold 数より少ない場合．
    """
    if n_folds < 2:
        raise ValueError(f"fold 数は 2 以上でなければならない（現在 {n_folds}）")
    by_class = _patients_by_class(records)
    for label, patients in by_class.items():
        # 患者数が fold 数未満のクラスは，held_out に現れない fold ができ評価が歪むため止める．
        if len(patients) < n_folds:
            raise ValueError(
                f"クラス {label} の患者数 {len(patients)} 人が fold 数 {n_folds} より少ない"
            )

    labels_of = patient_labels(records)
    patient_ids = sorted(labels_of)
    y = np.array([labels_of[p] for p in patient_ids])
    splitter = StratifiedKFold(n_splits=n_folds, shuffle=True, random_state=seed)

    folds: list[Fold] = []
    for index, (train_idx, held_idx) in enumerate(splitter.split(np.zeros(len(y)), y)):
        folds.append(
            Fold(
                index=index,
                train=tuple(sorted(patient_ids[i] for i in train_idx)),
                held_out=tuple(sorted(patient_ids[i] for i in held_idx)),
            )
        )
    logger.info(
        "%d-fold 分割（seed=%d）: held_out の患者数 %s",
        n_folds,
        seed,
        [len(f.held_out) for f in folds],
    )
    return folds


def select_records(
    records: Iterable[ManifestRecord], patient_ids: Collection[str]
) -> list[ManifestRecord]:
    """指定した患者のレコードだけを取り出す（入力の順序を保つ）.

    Args:
        records: manifest レコード．
        patient_ids: 取り出す患者 ID．

    Returns:
        該当するレコード．
    """
    wanted = set(patient_ids)
    return [r for r in records if r.patient_id in wanted]


def count_by_class(records: Iterable[ManifestRecord]) -> dict[str, int]:
    """クラスごとの画像枚数を数える.

    Args:
        records: manifest レコード．

    Returns:
        クラス slug → 枚数（全クラスをクラス ID 順に含み，無いクラスは 0）．
    """
    counts = Counter(r.label for r in records)
    return {label: counts.get(label, 0) for label in CLASS_SLUGS}


def count_by_class_severity(records: Iterable[ManifestRecord]) -> dict[tuple[str, str], int]:
    """(クラス, 程度) ごとの画像枚数を数える（分割後の severity の偏り確認用）.

    Args:
        records: manifest レコード．

    Returns:
        (label, severity) → 枚数（出現した組み合わせのみ．クラス ID 順・程度名順）．
    """
    counts = Counter((r.label, r.severity) for r in records)
    order = {label: i for i, label in enumerate(CLASS_SLUGS)}
    return dict(sorted(counts.items(), key=lambda kv: (order[kv[0][0]], kv[0][1])))


def _patients_by_class(records: Sequence[ManifestRecord]) -> dict[str, list[str]]:
    if not records:
        raise ValueError("分割対象のレコードが 0 件である")
    by_class: dict[str, list[str]] = {}
    for patient_id, label in patient_labels(records).items():
        if label not in CLASS_SLUGS:
            raise ValueError(f"未知の label がある: {label!r}（manifest の検証を通すこと）")
        by_class.setdefault(label, []).append(patient_id)
    # 入力順に依存させず決定的にするため，患者 ID でソートしておく．
    return {label: sorted(patients) for label, patients in by_class.items()}


def _round_half_up(value: float) -> int:
    # 組み込み round は偶数丸めで直感に反するため四捨五入を使う．
    return math.floor(value + 0.5)