"""患者単位 K-fold 交差検証の実行（SPEC_02「交差検証の実行」）.

1. manifest 全体で患者単位・9 グループ層化の K-fold 分割を作る（`data.severities` に関わらず同じ）
2. 各 fold で
   1. 学習側から `train.inner_val_ratio` の患者を検証用に取る（患者単位・9 グループ層化）
   2. `data.severities` を学習用と検証用にだけ適用する
   3. クラス不均衡の対策（`train.class_weighting`・`train.sampler`）を学習用データのクラス別枚数
      から求める（SPEC_02「クラス不均衡」→「対策の設定」）
   4. 検証 macro F1 が最良のエポックの重みを採用する（早期終了）
   5. held_out（**全程度**）で 1 回だけ評価する
3. 全 fold の held_out の予測を呼び出し側でまとめて集計する

ログには件数・グループ別件数・損失・指標・時間だけを出し，患者 ID・ファイル名は出さない．
"""

from __future__ import annotations

import logging
import time
from collections.abc import Mapping, Sequence
from dataclasses import dataclass, field
from pathlib import Path

import numpy as np
import torch
from torch.utils.data import DataLoader, WeightedRandomSampler

from tias_judge_model.data.class_balance import (
    balanced_sample_weights,
    class_draw_probabilities,
    compute_class_weights,
    count_classes,
)
from tias_judge_model.data.classes import CLASS_SLUGS, label_to_id
from tias_judge_model.data.dataset import TongueImageDataset
from tias_judge_model.data.images import load_letterboxed_images
from tias_judge_model.data.manifest import ManifestRecord
from tias_judge_model.data.split import (
    Fold,
    count_by_group,
    filter_by_severity,
    make_kfold_splits,
    select_records,
    split_inner_val,
)
from tias_judge_model.models.factory import count_parameters, create_model
from tias_judge_model.schema import ExperimentConfig, TrainConfig
from tias_judge_model.training.loop import (
    EpochRecord,
    device_name,
    predict_logits,
    seed_everything,
    select_device,
    train_model,
)

logger = logging.getLogger(__name__)


@dataclass
class ClassBalance:
    """1 fold の学習に使ったクラス不均衡の対策（SPEC_02「クラス不均衡」→「対策の設定」）.

    Attributes:
        class_weighting: 損失のクラス重みの方式（`none`・`inverse`・`sqrt_inverse`）．
        sampler: 学習データの抜き出し方（`random`・`balanced`）．
        class_counts: 学習用データ（`data.severities` 適用後）のクラス別枚数（クラス ID 順）．
        class_weights: 損失のクラス重み（クラス ID 順）．`class_weighting` が `none` なら None．
        class_draw_probabilities: 1 回の抜き出しで各クラスが選ばれる確率（クラス ID 順）．
            `sampler` が `random` なら None．
    """

    class_weighting: str
    sampler: str
    class_counts: tuple[int, ...]
    class_weights: tuple[float, ...] | None = None
    class_draw_probabilities: tuple[float, ...] | None = None


@dataclass
class FoldResult:
    """1 fold の学習・評価の結果.

    `train_patients`・`val_patients`・`held_out_patients` と `held_out_indices` は個票に当たるため，
    サーバ上の `split.json`・`predictions.csv` にだけ書き出す．

    Attributes:
        index: fold 番号（0 始まり）．
        seed: この fold の学習に使った seed．
        train_patients: 学習用の患者 ID（除外前．昇順）．
        val_patients: 検証用の患者 ID（除外前．昇順）．
        held_out_patients: held_out の患者 ID（昇順）．
        group_counts: 分割名（`train`・`val`・`held_out`）→ 層化グループ → 画像枚数
            （`train`・`val` は `data.severities` 適用後，`held_out` は全程度）．
        history: エポックごとの記録．
        best_epoch: 最良エポック（1 始まり）．
        best_val_macro_f1: 最良エポックの検証 macro F1．
        train_seconds: 学習の所要秒数．
        eval_seconds: held_out の評価の所要秒数．
        held_out_indices: `(M,)` の held_out 画像の manifest 上の位置（0 始まり）．
        held_out_probs: `(M, C)` の held_out の予測確率（softmax）．
        class_balance: この fold の学習に使ったクラス不均衡の対策．
    """

    index: int
    seed: int
    train_patients: tuple[str, ...]
    val_patients: tuple[str, ...]
    held_out_patients: tuple[str, ...]
    group_counts: dict[str, dict[str, int]]
    history: list[EpochRecord]
    best_epoch: int
    best_val_macro_f1: float
    train_seconds: float
    eval_seconds: float
    held_out_indices: np.ndarray
    held_out_probs: np.ndarray
    class_balance: ClassBalance | None = None

    @property
    def held_out_pred(self) -> np.ndarray:
        """`(M,)` の held_out の予測クラス ID."""
        return self.held_out_probs.argmax(axis=1)


@dataclass
class CrossValidationResult:
    """交差検証全体の結果.

    Attributes:
        folds: fold ごとの結果．
        data_loading_seconds: 画像の読み込み（レターボックス化）の所要秒数．
        device: デバイスの表示名（GPU 名または `cpu`）．
        num_parameters: モデルのパラメータ数．
    """

    folds: list[FoldResult] = field(default_factory=list)
    data_loading_seconds: float = 0.0
    device: str = "cpu"
    num_parameters: int = 0


def fold_seed(base_seed: int, fold_index: int) -> int:
    """fold ごとの seed（内側の検証分割と学習の乱数）を分割の seed から決める.

    Args:
        base_seed: 設定の `split.seed`．
        fold_index: fold 番号（0 始まり）．

    Returns:
        `base_seed + fold_index`．
    """
    return base_seed + fold_index


def run_cross_validation(
    config: ExperimentConfig, records: Sequence[ManifestRecord], checkpoint_dir: Path
) -> CrossValidationResult:
    """K-fold 交差検証で学習・評価する（各 fold の最良の重みを `checkpoint_dir` に保存する）.

    Args:
        config: 実験設定．
        records: manifest の全レコード（検証済み．除外前）．
        checkpoint_dir: 各 fold の最良の重み `fold<k>.pt` の保存先．

    Returns:
        交差検証の結果．

    Raises:
        ValueError: 分割できない，画像を読めない，学習用・検証用が空になる，対策を指定したのに
            学習用データに 1 枚も無いクラスがある等の入力の問題．
    """
    records = list(records)
    device = select_device()
    result = CrossValidationResult(device=device_name(device))
    logger.info(
        "交差検証を開始する: 画像 %d 件・%d fold・モデル %s（pretrained=%s）・程度 %s・デバイス %s",
        len(records),
        config.split.n_folds,
        config.model.name,
        config.model.pretrained,
        list(config.data.severities),
        result.device,
    )

    folds = make_kfold_splits(records, config.split.n_folds, config.split.seed)

    started = time.monotonic()
    images = load_letterboxed_images(records, config.data.images_dir, config.train.image_size)
    result.data_loading_seconds = time.monotonic() - started
    logger.info("画像の読み込みに %.1f 秒かかった", result.data_loading_seconds)
    labels = np.array([label_to_id(r.label) for r in records], dtype=np.int64)
    position = {r.image_path: i for i, r in enumerate(records)}

    checkpoint_dir.mkdir(parents=True, exist_ok=True)
    n_folds = len(folds)
    for fold in folds:
        prefix = f"fold{fold.index}（{fold.index + 1}/{n_folds}）"
        seed = fold_seed(config.split.seed, fold.index)
        prepared = prepare_fold(fold, records, config, position, seed=seed, log_prefix=prefix)
        train_idx, val_idx, held_idx = prepared.train_idx, prepared.val_idx, prepared.held_idx

        balance, class_weights, sample_weights = make_class_balance(labels[train_idx], config.train)
        log_class_balance(prefix, balance)

        seed_everything(seed)
        model = create_model(config.model.name, pretrained=config.model.pretrained).to(device)
        result.num_parameters = count_parameters(model)
        train_loader = _make_loader(
            images[train_idx],
            labels[train_idx],
            config.train,
            device,
            train=True,
            seed=seed,
            sample_weights=sample_weights,
        )
        val_loader = _make_loader(
            images[val_idx], labels[val_idx], config.train, device, train=False, seed=seed
        )
        held_loader = _make_loader(
            images[held_idx], labels[held_idx], config.train, device, train=False, seed=seed
        )

        trained = train_model(
            model,
            train_loader,
            val_loader,
            config.train,
            device,
            log_prefix=prefix,
            class_weights=None if class_weights is None else torch.from_numpy(class_weights),
        )
        torch.save(
            {
                "fold": fold.index,
                "model_name": config.model.name,
                "best_epoch": trained.best_epoch,
                "state_dict": trained.best_state,
            },
            checkpoint_dir / f"fold{fold.index}.pt",
        )

        eval_started = time.monotonic()
        use_amp = config.train.amp and device.type == "cuda"
        logits, _ = predict_logits(model, held_loader, device, amp=use_amp)
        probs = torch.softmax(torch.from_numpy(logits), dim=1).numpy()
        eval_seconds = time.monotonic() - eval_started
        held_acc = float(np.mean(probs.argmax(axis=1) == labels[held_idx]))
        logger.info(
            "%s: 最良エポック %d（val_macro_f1 %.4f）・held_out accuracy %.4f"
            "（学習 %.1f 秒・評価 %.1f 秒）",
            prefix,
            trained.best_epoch,
            trained.best_val_macro_f1,
            held_acc,
            trained.seconds,
            eval_seconds,
        )

        result.folds.append(
            FoldResult(
                index=fold.index,
                seed=seed,
                train_patients=prepared.train_patients,
                val_patients=prepared.val_patients,
                held_out_patients=fold.held_out,
                group_counts=prepared.group_counts,
                history=trained.history,
                best_epoch=trained.best_epoch,
                best_val_macro_f1=trained.best_val_macro_f1,
                train_seconds=trained.seconds,
                eval_seconds=eval_seconds,
                held_out_indices=held_idx,
                held_out_probs=probs,
                class_balance=balance,
            )
        )
        del model, trained
        if device.type == "cuda":
            torch.cuda.empty_cache()
    return result


@dataclass
class PreparedFold:
    """1 fold の学習用・検証用・held_out の分け方（CNN と色特徴分類器で共通）.

    Attributes:
        train_patients: 学習用の患者 ID（除外前．昇順）．
        val_patients: 検証用の患者 ID（除外前．昇順）．
        group_counts: 分割名 → 層化グループ → 画像枚数（`train`・`val` は `data.severities`
            適用後，`held_out` は全程度）．
        train_idx: `(N_train,)` の int64．学習用画像の manifest 上の位置．
        val_idx: `(N_val,)` の int64．検証用画像の manifest 上の位置．
        held_idx: `(M,)` の int64．held_out 画像の manifest 上の位置．
    """

    train_patients: tuple[str, ...]
    val_patients: tuple[str, ...]
    group_counts: dict[str, dict[str, int]]
    train_idx: np.ndarray
    val_idx: np.ndarray
    held_idx: np.ndarray


def prepare_fold(
    fold: Fold,
    records: Sequence[ManifestRecord],
    config: ExperimentConfig,
    position: Mapping[str, int],
    *,
    seed: int,
    log_prefix: str,
) -> PreparedFold:
    """fold の学習側から内側の検証用を取り，`data.severities` を学習用と検証用にだけ適用する.

    SPEC_02「交差検証の実行」の手順 2.1・2.2．held_out には `data.severities` を適用しない
    （全程度で評価する）．件数とグループ別枚数をログに出す（患者 ID は出さない）．

    Args:
        fold: `make_kfold_splits` の fold．
        records: manifest の全レコード（除外前）．
        config: 実験設定（`train.inner_val_ratio`・`data.severities` を使う）．
        position: 画像パス → manifest 上の位置（0 始まり）．
        seed: 内側の検証分割の seed（`fold_seed`）．
        log_prefix: ログの接頭辞（fold の表示）．

    Returns:
        学習用・検証用・held_out の分け方．

    Raises:
        ValueError: 内側の検証分割ができない，または学習用・検証用・held_out のいずれかが
            0 件になる場合．
    """
    train_side = select_records(records, fold.train)
    train_patients, val_patients = split_inner_val(train_side, config.train.inner_val_ratio, seed)
    train_records = filter_by_severity(
        select_records(train_side, train_patients), config.data.severities
    )
    val_records = filter_by_severity(
        select_records(train_side, val_patients), config.data.severities
    )
    # held_out には data.severities を適用しない（全程度で評価する）．
    held_records = select_records(records, fold.held_out)
    if not train_records or not val_records or not held_records:
        raise ValueError(
            f"{log_prefix}: 学習用 {len(train_records)} 件・検証用 {len(val_records)} 件・"
            f"held_out {len(held_records)} 件のいずれかが 0 件である"
        )
    group_counts = {
        "train": count_by_group(train_records),
        "val": count_by_group(val_records),
        "held_out": count_by_group(held_records),
    }
    logger.info(
        "%s: 画像 学習 %d・検証 %d・held_out %d（患者 学習 %d・検証 %d・held_out %d）",
        log_prefix,
        len(train_records),
        len(val_records),
        len(held_records),
        len(train_patients),
        len(val_patients),
        len(fold.held_out),
    )
    for name, counts in group_counts.items():
        logger.info("%s: %s のグループ別枚数 %s", log_prefix, name, counts)
    return PreparedFold(
        train_patients=train_patients,
        val_patients=val_patients,
        group_counts=group_counts,
        train_idx=_positions(train_records, position),
        val_idx=_positions(val_records, position),
        held_idx=_positions(held_records, position),
    )


def make_class_balance(
    train_labels: np.ndarray, config: TrainConfig
) -> tuple[ClassBalance, np.ndarray | None, np.ndarray | None]:
    """学習用データのラベルからクラス不均衡の対策（損失のクラス重み・抜き出し重み）を求める.

    Args:
        train_labels: `(N,)` の整数配列（この fold の学習用データのクラス ID．
            `data.severities` 適用後）．
        config: 学習の設定（`class_weighting`・`sampler` を使う）．

    Returns:
        (記録用の対策, `(NUM_CLASSES,)` の float64 の損失のクラス重み（`none` なら None），
        `(N,)` の float64 の画像ごとの抜き出し重み（`random` なら None））．

    Raises:
        ValueError: 対策を指定したのに学習用データに 1 枚も無いクラスがある場合．
    """
    counts = count_classes(train_labels)
    class_weights: np.ndarray | None = None
    sample_weights: np.ndarray | None = None
    draw_probs: np.ndarray | None = None
    if config.class_weighting != "none":
        class_weights = compute_class_weights(train_labels, config.class_weighting)
    if config.sampler == "balanced":
        sample_weights = balanced_sample_weights(train_labels)
        draw_probs = class_draw_probabilities(train_labels, sample_weights)
    balance = ClassBalance(
        class_weighting=config.class_weighting,
        sampler=config.sampler,
        class_counts=tuple(int(n) for n in counts),
        class_weights=None if class_weights is None else tuple(float(w) for w in class_weights),
        class_draw_probabilities=(
            None if draw_probs is None else tuple(float(p) for p in draw_probs)
        ),
    )
    return balance, class_weights, sample_weights


def log_class_balance(prefix: str, balance: ClassBalance) -> None:
    """fold の学習に使うクラス不均衡の対策（クラス別枚数・重み・抜き出し確率）をログに出す.

    Args:
        prefix: ログの接頭辞（fold の表示）．
        balance: クラス不均衡の対策．
    """

    def by_class(values: Sequence[float]) -> dict[str, float]:
        return {slug: round(v, 4) for slug, v in zip(CLASS_SLUGS, values, strict=True)}

    counts = dict(zip(CLASS_SLUGS, balance.class_counts, strict=True))
    logger.info(
        "%s: 対策 class_weighting=%s・sampler=%s（学習のクラス別枚数 %s）",
        prefix,
        balance.class_weighting,
        balance.sampler,
        counts,
    )
    if balance.class_weights is not None:
        logger.info("%s: 損失のクラス重み %s", prefix, by_class(balance.class_weights))
    if balance.class_draw_probabilities is not None:
        logger.info(
            "%s: 1 回の抜き出しでクラスが選ばれる確率 %s",
            prefix,
            by_class(balance.class_draw_probabilities),
        )


def _positions(records: Sequence[ManifestRecord], position: Mapping[str, int]) -> np.ndarray:
    # `(N,)` の int64．各レコードの manifest 上の位置（0 始まり）．
    return np.array([position[r.image_path] for r in records], dtype=np.int64)


def _make_loader(
    images: np.ndarray,
    labels: np.ndarray,
    config: TrainConfig,
    device: torch.device,
    *,
    train: bool,
    seed: int,
    sample_weights: np.ndarray | None = None,
) -> DataLoader:
    dataset = TongueImageDataset(images, labels, augment=train)
    generator = torch.Generator()
    generator.manual_seed(seed)
    sampler: WeightedRandomSampler | None = None
    if sample_weights is not None:
        # バランス・サンプリング: 重複ありで，1 エポックに学習画像数と同じ枚数を抜き出す．
        # 抜き出しも seed から再現できるよう，DataLoader と同じ generator を渡す．
        sampler = WeightedRandomSampler(
            torch.as_tensor(sample_weights, dtype=torch.float64),
            num_samples=len(dataset),
            replacement=True,
            generator=generator,
        )
    return DataLoader(
        dataset,
        batch_size=config.batch_size,
        shuffle=train and sampler is None,
        sampler=sampler,
        # BatchNorm を含むモデルで最後のバッチが 1 件になると学習が止まるため，学習時は端数を捨てる
        # （件数がバッチサイズ以下のときは捨てない）．
        drop_last=train and len(dataset) > config.batch_size,
        num_workers=config.num_workers,
        pin_memory=device.type == "cuda",
        persistent_workers=config.num_workers > 0,
        generator=generator,
    )