"""pydantic による設定スキーマと設定 YAML の読み込み.

学習・モデル関連の設定（`model`・`train`）は省略でき，省略時は SPEC_02「学習の設定」の
既定値で補う（学習しないコマンド `tias-dataset` の設定ファイルをそのまま読めるようにするため）．
"""

from __future__ import annotations

import os
from collections.abc import Mapping
from pathlib import Path
from typing import Any, Literal

import yaml
from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator

from tias_judge_model.data.classes import SEVERITIES

# データのルートを上書きする環境変数．サーバと手元で YAML を書き換えずに済ませるため．
DATA_ROOT_ENV_VAR = "TIAS_DATA_ROOT"
# 実験 ID に入れる git の短縮ハッシュ．サーバには .git を送らないため，tias-remote run が
# 手元で求めてこの環境変数で渡す（SPEC_01「実験の出力」・SPEC_03「run」）．
GIT_COMMIT_ENV_VAR = "TIAS_GIT_COMMIT"


class _StrictModel(BaseModel):
    """未知のキー（綴り間違い）を拒否し，生成後の変更を禁止する基底モデル."""

    model_config = ConfigDict(extra="forbid", frozen=True)


class DataConfig(_StrictModel):
    """データの場所に関する設定.

    Attributes:
        root: データのルートディレクトリ（`images/` と manifest を含む）．`~` は展開される．
            相対パスはカレントディレクトリ基準で解釈する．
        manifest: manifest CSV のパス．相対パスは `root` からの相対として解釈する．
        severities: 学習・評価に使う程度（`strong`・`mild`・`normal` の部分集合）．
            分割の後に適用する（SPEC_02「程度（ややあり）の扱い」）．値は `SEVERITIES` の
            順（strong → mild → normal）にそろえて保持する．
    """

    root: Path
    manifest: Path = Path("manifest.csv")
    severities: tuple[str, ...] = SEVERITIES

    @field_validator("root", "manifest", mode="after")
    @classmethod
    def _expand_user(cls, value: Path) -> Path:
        return value.expanduser()

    @field_validator("severities", mode="after")
    @classmethod
    def _check_severities(cls, value: tuple[str, ...]) -> tuple[str, ...]:
        if not value:
            raise ValueError("data.severities が空である（strong・mild・normal から 1 つ以上）")
        unknown = [s for s in value if s not in SEVERITIES]
        if unknown:
            raise ValueError(
                f"data.severities に未知の値がある: {', '.join(map(repr, unknown))}"
                f"（有効値: {', '.join(SEVERITIES)}）"
            )
        if len(set(value)) != len(value):
            raise ValueError("data.severities に重複がある")
        # 記載順の違いで config.yaml の記録が揺れないよう，定義順にそろえる．
        return tuple(s for s in SEVERITIES if s in value)

    @property
    def images_dir(self) -> Path:
        """画像ディレクトリ（`<root>/images`）."""
        return self.root / "images"

    @property
    def manifest_path(self) -> Path:
        """manifest CSV の解決済みパス."""
        return self.manifest if self.manifest.is_absolute() else self.root / self.manifest


class SplitConfig(_StrictModel):
    """患者単位分割の設定.

    Attributes:
        seed: 分割の乱数 seed．
        val_ratio: 単一分割で val に回す患者の比率．
        test_ratio: 単一分割で test に回す患者の比率．
        n_folds: 交差検証の fold 数．
    """

    seed: int = 42
    val_ratio: float = Field(default=0.15, gt=0.0, lt=1.0)
    test_ratio: float = Field(default=0.15, gt=0.0, lt=1.0)
    n_folds: int = Field(default=5, ge=2)

    @model_validator(mode="after")
    def _check_ratio_sum(self) -> SplitConfig:
        # train が残らない比率は分割として成立しないため拒否する．
        if self.val_ratio + self.test_ratio >= 1.0:
            raise ValueError(
                "val_ratio と test_ratio の合計は 1 未満でなければならない"
                f"（現在 {self.val_ratio} + {self.test_ratio}）"
            )
        return self


class ModelConfig(_StrictModel):
    """モデルの設定（SPEC_02「学習の設定」）.

    Attributes:
        name: timm のモデル名（例: `convnext_tiny`・`resnet50`・`vit_small_patch16_224`）．
        pretrained: ImageNet 事前学習の重みを使うか．
    """

    name: str = Field(default="convnext_tiny", min_length=1)
    pretrained: bool = True


class TrainConfig(_StrictModel):
    """学習の設定（SPEC_02「学習の設定」）.

    Attributes:
        image_size: 入力画像の一辺（長辺を合わせるレターボックス）．
        batch_size: バッチサイズ．
        max_epochs: 最大エポック数．
        patience: 検証 macro F1 がこのエポック数改善しなければ打ち切る．
        lr: AdamW の学習率．
        weight_decay: AdamW の weight decay．
        warmup_epochs: 線形 warmup のエポック数（その後 cosine で減衰）．
        inner_val_ratio: 交差検証の各 fold で，学習側から早期終了用の検証に回す患者の比率．
        num_workers: DataLoader のワーカー数．
        amp: CUDA のときに混合精度を使うか（CPU では無視する）．
        class_weighting: 損失のクラス重み（`none`・`inverse`・`sqrt_inverse`．SPEC_02
            「クラス不均衡」→「対策の設定」）．`none` なら重みなし（従来どおり）．
        sampler: 学習データの抜き出し方（`random`・`balanced`．同上）．`random` なら
            従来どおりのシャッフル．
    """

    image_size: int = Field(default=224, ge=16)
    batch_size: int = Field(default=32, ge=1)
    max_epochs: int = Field(default=40, ge=1)
    patience: int = Field(default=8, ge=1)
    lr: float = Field(default=1e-4, gt=0.0)
    weight_decay: float = Field(default=0.05, ge=0.0)
    warmup_epochs: int = Field(default=2, ge=0)
    inner_val_ratio: float = Field(default=0.15, gt=0.0, lt=1.0)
    num_workers: int = Field(default=4, ge=0)
    amp: bool = True
    class_weighting: Literal["none", "inverse", "sqrt_inverse"] = "none"
    sampler: Literal["random", "balanced"] = "random"


def _check_candidates(values: tuple[Any, ...], key: str) -> tuple[Any, ...]:
    # 候補が空だとハイパーパラメータを選べず，重複は同じ設定を 2 度学習するだけなので拒否する．
    if not values:
        raise ValueError(f"{key} が空である（候補を 1 つ以上）")
    if len(set(values)) != len(values):
        raise ValueError(f"{key} に重複がある")
    return values


class LogisticConfig(_StrictModel):
    """色特徴＋ロジスティック回帰（StandardScaler → 多クラス LogisticRegression）の設定.

    Attributes:
        c_values: 正則化の強さの逆数 C の候補（各 fold の内側の検証 macro F1 で 1 つ選ぶ）．
            既定は 1e-2〜1e2 を 10 倍刻み（特徴は標準化するので，この幅で弱〜強の正則化を覆う）．
        max_iter: lbfgs の最大反復回数．
    """

    c_values: tuple[float, ...] = (0.01, 0.1, 1.0, 10.0, 100.0)
    max_iter: int = Field(default=5000, ge=1)

    @field_validator("c_values", mode="after")
    @classmethod
    def _check_c_values(cls, value: tuple[float, ...]) -> tuple[float, ...]:
        if any(c <= 0 for c in value):
            raise ValueError("color.logistic.c_values は正の値でなければならない")
        return _check_candidates(value, "color.logistic.c_values")


class HistGradientBoostingConfig(_StrictModel):
    """色特徴＋勾配ブースティング（HistGradientBoostingClassifier）の設定.

    候補は `learning_rates` × `max_depths` の全組み合わせ（各 fold の内側の検証 macro F1 で
    1 つ選ぶ）．学習データが 1000 枚未満と少ないため，候補は小さく保つ．

    Attributes:
        learning_rates: 学習率の候補．
        max_depths: 木の深さの上限の候補（`null` は上限なし．葉の数は `max_leaf_nodes` で制限）．
        max_iter: ブースティングの反復回数（早期終了は使わない．内側の検証は候補選びに使うため）．
        max_leaf_nodes: 1 本の木の葉の数の上限．
        min_samples_leaf: 葉に必要な最小サンプル数（少数クラスへの過学習を抑える）．
        l2_regularization: L2 正則化の強さ．
    """

    learning_rates: tuple[float, ...] = (0.05, 0.1)
    max_depths: tuple[int | None, ...] = (3, None)
    max_iter: int = Field(default=200, ge=1)
    max_leaf_nodes: int = Field(default=31, ge=2)
    min_samples_leaf: int = Field(default=20, ge=1)
    l2_regularization: float = Field(default=0.0, ge=0.0)

    @field_validator("learning_rates", mode="after")
    @classmethod
    def _check_learning_rates(cls, value: tuple[float, ...]) -> tuple[float, ...]:
        if any(lr <= 0 for lr in value):
            raise ValueError("color.hist_gb.learning_rates は正の値でなければならない")
        return _check_candidates(value, "color.hist_gb.learning_rates")

    @field_validator("max_depths", mode="after")
    @classmethod
    def _check_max_depths(cls, value: tuple[int | None, ...]) -> tuple[int | None, ...]:
        if any(d is not None and d < 1 for d in value):
            raise ValueError("color.hist_gb.max_depths は 1 以上（または null）でなければならない")
        return _check_candidates(value, "color.hist_gb.max_depths")


class ColorClassifierConfig(_StrictModel):
    """色特徴量による分類器の設定（`tias-train color-cv`）.

    分割・内側の検証の比率・クラス重み（`train.inner_val_ratio`・`train.class_weighting`）は
    CNN の交差検証と共通の設定を使う．

    Attributes:
        classifier: 分類器（`logistic`: StandardScaler ＋ 多クラス LogisticRegression，
            `hist_gb`: HistGradientBoostingClassifier）．
        max_side: 特徴抽出の前に，長辺がこれを超える画像を最近傍法で縮小する（処理時間の短縮．
            最近傍法は画素の色を混ぜない）．`null` なら縮小しない．
        logistic: `logistic` のハイパーパラメータ．
        hist_gb: `hist_gb` のハイパーパラメータ．
    """

    classifier: Literal["logistic", "hist_gb"] = "logistic"
    max_side: int | None = Field(default=512, ge=16)
    logistic: LogisticConfig = Field(default_factory=LogisticConfig)
    hist_gb: HistGradientBoostingConfig = Field(default_factory=HistGradientBoostingConfig)


class ExperimentConfig(_StrictModel):
    """実験設定のルート.

    Attributes:
        data: データの場所．
        split: 分割の設定．
        model: モデルの設定（省略時は既定値）．
        train: 学習の設定（省略時は既定値）．
        color: 色特徴量による分類器の設定（`tias-train color-cv` で使う．CNN の設定では省略）．
    """

    data: DataConfig
    split: SplitConfig = Field(default_factory=SplitConfig)
    model: ModelConfig = Field(default_factory=ModelConfig)
    train: TrainConfig = Field(default_factory=TrainConfig)
    color: ColorClassifierConfig | None = None


def load_config(path: Path, environ: Mapping[str, str] | None = None) -> ExperimentConfig:
    """設定 YAML を読み込み，検証済みの設定を返す.

    環境変数 `TIAS_DATA_ROOT` が空でない値で設定されていれば，YAML の `data.root` より優先する．

    Args:
        path: 設定 YAML のパス．
        environ: 参照する環境変数．None なら `os.environ` を使う．

    Returns:
        検証済みの設定．

    Raises:
        ValueError: YAML のトップレベルがマッピングでない場合．
        pydantic.ValidationError: 設定値が不正な場合．
    """
    env = os.environ if environ is None else environ
    with path.open(encoding="utf-8") as f:
        raw: Any = yaml.safe_load(f)
    if raw is None:
        raw = {}
    if not isinstance(raw, dict):
        raise ValueError(f"設定 YAML のトップレベルはマッピングでなければならない: {path}")

    env_root = env.get(DATA_ROOT_ENV_VAR, "")
    if env_root:
        data = raw.get("data") or {}
        if not isinstance(data, dict):
            raise ValueError(f"設定 YAML の data はマッピングでなければならない: {path}")
        raw = {**raw, "data": {**data, "root": env_root}}

    return ExperimentConfig.model_validate(raw)