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

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

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

from __future__ import annotations

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

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

from tias_judge_model.data.classes import 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 (
    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 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）．
    """

    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

    @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: 分割できない，画像を読めない，学習用・検証用が空になる等の入力の問題．
    """
    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)
        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"{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）",
            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", prefix, name, counts)

        train_idx = np.array([position[r.image_path] for r in train_records], dtype=np.int64)
        val_idx = np.array([position[r.image_path] for r in val_records], dtype=np.int64)
        held_idx = np.array([position[r.image_path] for r in held_records], dtype=np.int64)

        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
        )
        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
        )
        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=train_patients,
                val_patients=val_patients,
                held_out_patients=fold.held_out,
                group_counts=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,
            )
        )
        del model, trained
        if device.type == "cuda":
            torch.cuda.empty_cache()
    return result


def _make_loader(
    images: np.ndarray,
    labels: np.ndarray,
    config: TrainConfig,
    device: torch.device,
    *,
    train: bool,
    seed: int,
) -> DataLoader:
    dataset = TongueImageDataset(images, labels, augment=train)
    generator = torch.Generator()
    generator.manual_seed(seed)
    return DataLoader(
        dataset,
        batch_size=config.batch_size,
        shuffle=train,
        # 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,
    )