"""患者単位の分割（単一分割・K-fold 交差検証）・程度による絞り込み・分割後の集計.

分割は常に `patient_id` 単位で行い，同一患者が複数の分割に跨らないようにする（SPEC_02）．
層化は「クラス × 程度」（`<label>_<severity>` の 9 グループ）で行う．1 患者は 1 つの label
しか持たない（manifest の検証で保証）ため，患者ごとに層化グループを 1 つ決めて患者の一覧を
分ければ「患者単位」と「層化」を同時に満たせる．

程度（`data.severities`）による除外は分割の**後**に適用する（`split_records`）．分割は常に
manifest 全体で行うことで，「ややあり除外」の各分割が「ややあり込み」の各分割の部分集合になり，
2 つの条件のモデルを同じ患者の test で比べられる（SPEC_02「程度（ややあり）の扱い」）．
"""

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 ALL_DIR_NAMES, CLASS_SLUGS, SEVERITIES
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)}

    def parts(self) -> dict[str, tuple[str, ...]]:
        """分割名 → 患者 ID（`train`・`val`・`test` の順）を返す."""
        return {"train": self.train, "val": self.val, "test": 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 parts(self) -> dict[str, tuple[str, ...]]:
        """分割名 → 患者 ID（`train`・`held_out` の順）を返す."""
        return {"train": self.train, "held_out": 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 record_group(record: ManifestRecord) -> str:
    """レコードの層化グループ名（`<label>_<severity>`）を返す（組み合わせの検証はしない）.

    Args:
        record: manifest レコード．

    Returns:
        グループ名（例: `"red_mild"`）．
    """
    return f"{record.label}_{record.severity}"


def patient_groups(records: Iterable[ManifestRecord]) -> dict[str, str]:
    """患者 ID → 層化グループ（`<label>_<severity>`）の対応を作る.

    同じ患者に複数の severity がある場合（1 患者複数枚）は，最も強いもの
    （`strong` > `mild` > `normal`）をその患者のグループとする（SPEC_02「データ分割」）．

    Args:
        records: manifest レコード．

    Returns:
        患者 ID をキー，層化グループ名を値とする辞書．

    Raises:
        ValueError: 同一患者に異なる label が付いている，または未知の severity がある場合．
    """
    records = list(records)
    labels_of = patient_labels(records)
    # SEVERITIES は強い順（strong → mild → normal）に並んでいる．
    rank = {severity: i for i, severity in enumerate(SEVERITIES)}
    strongest: dict[str, str] = {}
    for r in records:
        if r.severity not in rank:
            raise ValueError(f"未知の severity がある: {r.severity!r}（manifest の検証を通すこと）")
        current = strongest.get(r.patient_id)
        if current is None or rank[r.severity] < rank[current]:
            strongest[r.patient_id] = r.severity
    return {p: f"{label}_{strongest[p]}" for p, label in labels_of.items()}


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

    層化グループ（`<label>_<severity>`．`patient_groups`）ごとに患者 ID を昇順に並べてから
    seed で並べ替え，先頭から test・val を `round(患者数 × 比率)` 人ずつ取り，残りを train
    とする．同じ入力と seed なら入力レコードの順序によらず同じ結果になる．
    `data.severities` による除外はここでは行わない（分割後に `split_records` で適用する）．

    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_group = _patients_by_group(records)
    rng = np.random.default_rng(seed)

    train: list[str] = []
    val: list[str] = []
    test: list[str] = []
    # グループを固定の順（ALL_DIR_NAMES）で処理し，乱数の消費順を入力に依存させない．
    for group in ALL_DIR_NAMES:
        patients = by_group.get(group)
        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"グループ {group} の患者 {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 人になった",
                group,
                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>_<severity>`．`patient_groups`）で
    層化して `StratifiedKFold` にかける．各患者はちょうど 1 つの fold の `held_out` に入る．
    `data.severities` による除外はここでは行わない（分割後に `split_records` で適用する）．

    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_group = _patients_by_group(records)
    for group, patients in by_group.items():
        # 患者数が fold 数未満のグループは，held_out に現れない fold ができ評価が歪むため止める．
        if len(patients) < n_folds:
            raise ValueError(
                f"グループ {group} の患者数 {len(patients)} 人が fold 数 {n_folds} より少ない"
            )

    groups_of = patient_groups(records)
    patient_ids = sorted(groups_of)
    y = np.array([groups_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 filter_by_severity(
    records: Iterable[ManifestRecord], severities: Collection[str]
) -> list[ManifestRecord]:
    """指定した程度の画像だけを残す（入力の順序を保つ）.

    除外は画像（レコード）単位で行う．1 患者複数枚で程度が混在する場合は，その患者の
    該当しない画像だけが除かれる．

    Args:
        records: manifest レコード．
        severities: 残す程度（`strong`・`mild`・`normal` の空でない部分集合）．

    Returns:
        該当するレコード．

    Raises:
        ValueError: `severities` が空，または未知の値を含む場合．
    """
    wanted = _check_severities(severities)
    return [r for r in records if r.severity in wanted]


def split_records(
    records: Sequence[ManifestRecord],
    split: SplitResult | Fold,
    severities: Collection[str] | None = None,
) -> dict[str, list[ManifestRecord]]:
    """分割結果（患者 ID）を各分割のレコードに展開し，程度による除外を適用する.

    分割は manifest 全体で行い，`data.severities` はこの関数で**分割の後に**適用する
    （SPEC_02「程度（ややあり）の扱い」）．こうすると，同じ seed の分割から作った
    「ややあり除外」の各分割は「ややあり込み」の各分割の部分集合になる．

    Args:
        records: 分割に使った manifest レコード（除外前の全体）．
        split: 単一分割の結果，または K-fold の 1 fold．
        severities: 残す程度．None なら除外しない．

    Returns:
        分割名（単一分割は `train`・`val`・`test`，fold は `train`・`held_out`）→
        その分割のレコード（入力の順序を保つ）．

    Raises:
        ValueError: `severities` が空，または未知の値を含む場合．
    """
    wanted = None if severities is None else _check_severities(severities)
    result: dict[str, list[ManifestRecord]] = {}
    for name, patient_ids in split.parts().items():
        selected = select_records(records, patient_ids)
        if wanted is not None:
            selected = [r for r in selected if r.severity in wanted]
        result[name] = selected
    return result


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 count_by_group(records: Iterable[ManifestRecord]) -> dict[str, int]:
    """層化グループ（`<label>_<severity>`）ごとの画像枚数を数える.

    Args:
        records: manifest レコード．

    Returns:
        グループ名 → 枚数（有効な 9 グループをすべて `ALL_DIR_NAMES` の順に含み，
        無いものは 0）．
    """
    counts = Counter(record_group(r) for r in records)
    return {group: counts.get(group, 0) for group in ALL_DIR_NAMES}


def _check_severities(severities: Collection[str]) -> frozenset[str]:
    # 設定は schema で検証済みだが，関数を直接呼んだときに空集合で全件が消えるのを防ぐ．
    if not severities:
        raise ValueError("残す程度が空である（strong・mild・normal から 1 つ以上）")
    unknown = sorted(set(severities) - set(SEVERITIES))
    if unknown:
        raise ValueError(f"未知の程度がある: {', '.join(map(repr, unknown))}")
    return frozenset(severities)


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


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