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


# ---------------------------------------------------------------------------
# セグメンテーションの評価（tias-seg eval）
# ---------------------------------------------------------------------------

# セグメンテーションのデータ・モデルの場所を上書きする環境変数．tias-remote run が渡す
# TIAS_DATA_ROOT は判定モデルのデータ（~/data/TIASJudgeModel）を指すため，別の名前にする．
SEG_DATA_ROOT_ENV_VAR = "TIAS_SEG_DATA_ROOT"
SEG_MODEL_ROOT_ENV_VAR = "TIAS_SEG_MODEL_ROOT"


class SegEvalSetConfig(_StrictModel):
    """評価セット 1 つ（`<data.root>/<dir>/sample/*` と `<data.root>/<dir>/label/*` の組）.

    Attributes:
        name: 評価セットの名前（metrics.json のキー・図のファイル名に使う．英数字・`_`・`-`）．
        dir: `data.root` からの相対パス．
        kind: `cropped`（舌の周りで切り出し済み．ResUNet の学習時と同じ条件で画像全体に掛ける）
            または `fullsize`（撮影画像そのもの．検出の扱いを比べる）．
    """

    name: str = Field(pattern=r"^[A-Za-z0-9][A-Za-z0-9_-]{0,63}$")
    dir: Path
    kind: Literal["cropped", "fullsize"]

    @field_validator("dir", mode="after")
    @classmethod
    def _check_relative(cls, value: Path) -> Path:
        if value.is_absolute() or ".." in value.parts:
            raise ValueError("data.sets[].dir は data.root からの相対パス（.. を含まない）")
        return value


class SegDataConfig(_StrictModel):
    """セグメンテーションの評価データの設定.

    Attributes:
        root: データのルート（`~` を展開する．相対パスはカレントディレクトリ基準）．
        sample_dirname: 評価セット内の入力画像のディレクトリ名．
        label_dirname: 評価セット内の正解マスクのディレクトリ名（入力画像と同じ stem で対応）．
        label_threshold: 正解マスク（グレースケールで読む．JPEG のため中間値を含む）を前景と
            みなす下限の画素値（この値以上が前景）．
        sets: 評価セット（記載順に評価する）．
    """

    root: Path
    sample_dirname: str = Field(default="sample", min_length=1)
    label_dirname: str = Field(default="label", min_length=1)
    label_threshold: int = Field(default=128, ge=1, le=255)
    sets: tuple[SegEvalSetConfig, ...]

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

    @field_validator("sets", mode="after")
    @classmethod
    def _check_sets(cls, value: tuple[SegEvalSetConfig, ...]) -> tuple[SegEvalSetConfig, ...]:
        if not value:
            raise ValueError("data.sets が空である（評価セットを 1 つ以上）")
        names = [s.name for s in value]
        if len(set(names)) != len(names):
            raise ValueError("data.sets の name に重複がある")
        return value


class SegModelsConfig(_StrictModel):
    """ONNX モデルの場所.

    Attributes:
        root: モデルを置いたディレクトリ（`~` を展開する．相対パスはカレントディレクトリ基準）．
        resunet: ResUNet の ONNX（`root` からの相対パス）．
        detection: 検出モデルの ONNX（`root` からの相対パス．`scripts/convert_detection_to_onnx.py`
            で変換したもの）．null なら検出モデルを使わない．ファイルが無い場合も警告して省く．
    """

    root: Path
    resunet: Path = Path("ResUNet.onnx")
    detection: Path | None = None

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

    @property
    def resunet_path(self) -> Path:
        """ResUNet の ONNX の解決済みパス."""
        return self.resunet if self.resunet.is_absolute() else self.root / self.resunet

    @property
    def detection_path(self) -> Path | None:
        """検出モデルの ONNX の解決済みパス（使わないなら None）."""
        if self.detection is None:
            return None
        return self.detection if self.detection.is_absolute() else self.root / self.detection


class ResUNetConfig(_StrictModel):
    """ResUNet の入出力の設定（DeepTIAS 2.7 と同じ値が既定）.

    Attributes:
        input_size: 入力の一辺（切り出した画像をこのサイズに線形補間で縮小・拡大する）．
        foreground_channel: 出力 `(1, 2, S, S)` のうち前景として使うチャンネル（DeepTIAS は 0）．
    """

    input_size: int = Field(default=256, ge=16)
    foreground_channel: Literal[0, 1] = 0


class SegPostprocessConfig(_StrictModel):
    """ResUNet の出力の後処理の設定.

    Attributes:
        min_area: 小領域除去で取り除く領域の面積の上限（未満を除く．前景・背景とも．
            ResUNet の出力サイズ（input_size × input_size）での画素数）．
        mask_threshold: 枠のサイズに線形補間で戻したマスク（0〜255）を前景とみなす下限を超える値
            （この値より大きい画素を前景とする）．
    """

    min_area: int = Field(default=500, ge=0)
    mask_threshold: int = Field(default=127, ge=0, le=254)


class SegDetectionConfig(_StrictModel):
    """検出モデル（TF1 物体検出を ONNX に変換したもの）の設定.

    Attributes:
        score_threshold: 採用するスコアの下限（これを超える中で最高の 1 つを採用する）．
        input_color: 検出モデルへ渡す画像の色の順（`rgb`・`bgr`）．
        input_name: 入力名（`:0` の有無は問わない）．
        boxes_name: 枠（正規化 `[ymin, xmin, ymax, xmax]`）の出力名．
        scores_name: スコアの出力名．
    """

    score_threshold: float = Field(default=0.5, ge=0.0, lt=1.0)
    input_color: Literal["rgb", "bgr"] = "rgb"
    input_name: str = "image_tensor"
    boxes_name: str = "detection_boxes"
    scores_name: str = "detection_scores"


class SegEvaluationConfig(_StrictModel):
    """評価の設定.

    Attributes:
        iou_threshold: IoU がこれ未満の件数を数える．
    """

    iou_threshold: float = Field(default=0.8, gt=0.0, le=1.0)


class SegEvalConfig(_StrictModel):
    """セグメンテーションの評価（`tias-seg eval`）の設定のルート."""

    data: SegDataConfig
    models: SegModelsConfig
    resunet: ResUNetConfig = Field(default_factory=ResUNetConfig)
    postprocess: SegPostprocessConfig = Field(default_factory=SegPostprocessConfig)
    detection: SegDetectionConfig = Field(default_factory=SegDetectionConfig)
    evaluation: SegEvaluationConfig = Field(default_factory=SegEvaluationConfig)


def load_seg_eval_config(path: Path, environ: Mapping[str, str] | None = None) -> SegEvalConfig:
    """セグメンテーションの評価の設定 YAML を読み込む.

    環境変数 `TIAS_SEG_DATA_ROOT`・`TIAS_SEG_MODEL_ROOT` が空でない値で設定されていれば，
    それぞれ `data.root`・`models.root` より優先する．

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

    Returns:
        検証済みの設定．

    Raises:
        ValueError: YAML のトップレベル・`data`・`models` がマッピングでない場合．
        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}")
    for section, var in (("data", SEG_DATA_ROOT_ENV_VAR), ("models", SEG_MODEL_ROOT_ENV_VAR)):
        value = env.get(var, "")
        if value:
            node = raw.get(section) or {}
            if not isinstance(node, dict):
                raise ValueError(f"設定 YAML の {section} はマッピングでなければならない: {path}")
            raw = {**raw, section: {**node, "root": value}}
    return SegEvalConfig.model_validate(raw)