"""1 つの学習（1 fold 分）のループと推論（SPEC_02「モデル」「学習の設定」）.

AdamW ＋ 線形 warmup → cosine 減衰（ステップ単位）．CUDA のときだけ AMP を使う．
エポックごとに検証の macro F1 を測り，最良エポックの重みを保持し，`patience` エポック改善が
無ければ打ち切る．ログには集計値（損失・F1・時間）だけを出す．
"""

from __future__ import annotations

import logging
import math
import random
import time
from collections.abc import Callable
from dataclasses import dataclass, field

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

from tias_judge_model.data.dataset import normalize_batch
from tias_judge_model.evaluation.metrics import macro_f1
from tias_judge_model.schema import TrainConfig

logger = logging.getLogger(__name__)


@dataclass
class EpochRecord:
    """1 エポックの記録（`history.csv` の 1 行）.

    Attributes:
        epoch: エポック番号（1 始まり）．
        train_loss: 学習損失の平均．
        val_loss: 検証損失の平均．
        val_macro_f1: 検証の macro F1．
        lr: エポック終了時の学習率．
        seconds: このエポックの所要秒数（学習 ＋ 検証）．
    """

    epoch: int
    train_loss: float
    val_loss: float
    val_macro_f1: float
    lr: float
    seconds: float


@dataclass
class TrainResult:
    """1 回の学習の結果.

    Attributes:
        best_state: 最良エポックの重み（CPU 上の state_dict）．
        best_epoch: 最良エポック番号（1 始まり）．
        best_val_macro_f1: 最良エポックの検証 macro F1．
        history: エポックごとの記録．
        seconds: 学習全体の所要秒数．
    """

    best_state: dict[str, torch.Tensor]
    best_epoch: int
    best_val_macro_f1: float
    history: list[EpochRecord] = field(default_factory=list)
    seconds: float = 0.0


def seed_everything(seed: int) -> None:
    """`random`・`numpy`・`torch`（CUDA を含む）の乱数 seed を固定する.

    Args:
        seed: 乱数 seed．
    """
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)


def select_device() -> torch.device:
    """CUDA があれば `cuda`（`CUDA_VISIBLE_DEVICES` で割り当てられた GPU），無ければ CPU を返す."""
    return torch.device("cuda" if torch.cuda.is_available() else "cpu")


def device_name(device: torch.device) -> str:
    """デバイスの表示名（GPU 名または `cpu`）を返す."""
    if device.type == "cuda":
        return torch.cuda.get_device_name(device)
    return "cpu"


def make_scheduler(
    optimizer: torch.optim.Optimizer, *, warmup_steps: int, total_steps: int
) -> torch.optim.lr_scheduler.LambdaLR:
    """線形 warmup の後に cosine で 0 へ減衰する学習率スケジューラ（ステップ単位）を作る.

    Args:
        optimizer: 最適化器．
        warmup_steps: warmup のステップ数．
        total_steps: 全ステップ数．

    Returns:
        ステップごとに `step()` を呼ぶスケジューラ．
    """

    def factor(step: int) -> float:
        if warmup_steps > 0 and step < warmup_steps:
            return (step + 1) / warmup_steps
        progress = (step - warmup_steps) / max(1, total_steps - warmup_steps)
        return 0.5 * (1.0 + math.cos(math.pi * min(1.0, progress)))

    return torch.optim.lr_scheduler.LambdaLR(optimizer, factor)


def train_model(
    model: nn.Module,
    train_loader: DataLoader,
    val_loader: DataLoader,
    config: TrainConfig,
    device: torch.device,
    *,
    log_prefix: str = "",
    on_epoch_end: Callable[[EpochRecord], None] | None = None,
    class_weights: torch.Tensor | None = None,
) -> TrainResult:
    """モデルを学習し，検証 macro F1 が最良のエポックの重みを返す（早期終了あり）.

    Args:
        model: 学習するモデル（`device` 上に置いておく）．
        train_loader: 学習用（`(B, 3, H, W)` の uint8 画像と `(B,)` のクラス ID を返す）．
        val_loader: 検証用（同上）．
        config: 学習の設定．
        device: デバイス．
        log_prefix: ログの先頭に付ける文字列（例: `fold 1/5`）．
        on_epoch_end: エポックごとに呼ぶ関数（任意）．
        class_weights: `(C,)` の float の損失のクラス重み（クラス ID 順．学習データでの平均が
            1 になるよう正規化済みのもの）．None なら重みなし．**学習の損失にだけかけ**，
            検証の損失は重みなしで計算する（SPEC_02「クラス不均衡」→「対策の設定」）．

    Returns:
        学習の結果．`model` の重みは最良エポックのものに戻しておく．
    """
    started = time.monotonic()
    use_amp = config.amp and device.type == "cuda"
    optimizer = torch.optim.AdamW(
        model.parameters(), lr=config.lr, weight_decay=config.weight_decay
    )
    steps_per_epoch = max(1, len(train_loader))
    scheduler = make_scheduler(
        optimizer,
        warmup_steps=config.warmup_epochs * steps_per_epoch,
        total_steps=config.max_epochs * steps_per_epoch,
    )
    scaler = torch.amp.GradScaler("cuda", enabled=use_amp)
    criterion = nn.CrossEntropyLoss()
    train_criterion = _make_train_criterion(class_weights, device)

    best_state: dict[str, torch.Tensor] | None = None
    best_epoch = 0
    best_f1 = -1.0
    history: list[EpochRecord] = []
    for epoch in range(1, config.max_epochs + 1):
        epoch_started = time.monotonic()
        model.train()
        loss_sum = 0.0
        count = 0
        for images, labels in train_loader:
            inputs = normalize_batch(images.to(device, non_blocking=True))
            targets = labels.to(device, non_blocking=True)
            optimizer.zero_grad(set_to_none=True)
            with torch.autocast(device_type=device.type, dtype=torch.float16, enabled=use_amp):
                loss = train_criterion(model(inputs), targets)
            scaler.scale(loss).backward()
            scaler.step(optimizer)
            scaler.update()
            scheduler.step()
            loss_sum += float(loss.detach()) * len(targets)
            count += len(targets)

        val_logits, val_labels = predict_logits(model, val_loader, device, amp=use_amp)
        val_loss = float(
            criterion(torch.from_numpy(val_logits), torch.from_numpy(val_labels)).item()
        )
        val_f1 = macro_f1(val_labels, val_logits.argmax(axis=1))
        record = EpochRecord(
            epoch=epoch,
            train_loss=loss_sum / max(1, count),
            val_loss=val_loss,
            val_macro_f1=val_f1,
            lr=float(optimizer.param_groups[0]["lr"]),
            seconds=time.monotonic() - epoch_started,
        )
        history.append(record)
        if on_epoch_end is not None:
            on_epoch_end(record)

        improved = val_f1 > best_f1
        if improved:
            best_f1 = val_f1
            best_epoch = epoch
            best_state = {k: v.detach().to("cpu", copy=True) for k, v in model.state_dict().items()}
        logger.info(
            "%s epoch %d/%d: train_loss %.4f・val_loss %.4f・val_macro_f1 %.4f・lr %.2e"
            "（%.1f 秒・経過 %.1f 秒）%s",
            log_prefix,
            epoch,
            config.max_epochs,
            record.train_loss,
            record.val_loss,
            record.val_macro_f1,
            record.lr,
            record.seconds,
            time.monotonic() - started,
            " *best" if improved else "",
        )
        if epoch - best_epoch >= config.patience:
            logger.info(
                "%s 検証 macro F1 が %d エポック改善しないため打ち切る", log_prefix, config.patience
            )
            break

    if best_state is None:
        raise RuntimeError("学習が 1 エポックも行われなかった")
    model.load_state_dict(best_state)
    return TrainResult(
        best_state=best_state,
        best_epoch=best_epoch,
        best_val_macro_f1=best_f1,
        history=history,
        seconds=time.monotonic() - started,
    )


def _make_train_criterion(
    class_weights: torch.Tensor | None, device: torch.device
) -> Callable[[torch.Tensor, torch.Tensor], torch.Tensor]:
    # 重みなしは従来と同じ損失（結果を変えないため，同じ nn.CrossEntropyLoss を使う）．
    if class_weights is None:
        return nn.CrossEntropyLoss()
    weighted = nn.CrossEntropyLoss(
        weight=class_weights.to(device=device, dtype=torch.float32), reduction="none"
    )

    def criterion(logits: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
        # nn.CrossEntropyLoss(weight=...) の既定の mean はバッチ内の重みの和で割るため，重みの
        # 正規化（Σ (n_c / N)・w_c = 1）が効かない．件数で割り，損失の大きさを対策なしと揃える．
        return weighted(logits, targets).mean()

    return criterion


@torch.no_grad()
def predict_logits(
    model: nn.Module, loader: DataLoader, device: torch.device, *, amp: bool
) -> tuple[np.ndarray, np.ndarray]:
    """データ全体のロジットと正解を返す（推論モード）.

    Args:
        model: モデル．
        loader: `(B, 3, H, W)` の uint8 画像と `(B,)` のクラス ID を返す DataLoader
            （shuffle なし）．
        device: デバイス．
        amp: 混合精度を使うか（CUDA のときだけ有効にすること）．

    Returns:
        (`(N, C)` の float32 ロジット, `(N,)` の int64 正解クラス ID)．並びは loader の順．
    """
    model.eval()
    logits: list[np.ndarray] = []
    labels: list[np.ndarray] = []
    for images, targets in loader:
        inputs = normalize_batch(images.to(device, non_blocking=True))
        with torch.autocast(device_type=device.type, dtype=torch.float16, enabled=amp):
            out = model(inputs)
        logits.append(out.float().cpu().numpy())
        labels.append(targets.numpy().astype(np.int64))
    if not logits:
        return np.zeros((0, 0), dtype=np.float32), np.zeros((0,), dtype=np.int64)
    return np.concatenate(logits), np.concatenate(labels)