"""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"


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

    Attributes:
        data: データの場所．
        split: 分割の設定．
        model: モデルの設定（省略時は既定値）．
        train: 学習の設定（省略時は既定値）．
    """

    data: DataConfig
    split: SplitConfig = Field(default_factory=SplitConfig)
    model: ModelConfig = Field(default_factory=ModelConfig)
    train: TrainConfig = Field(default_factory=TrainConfig)


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)