"""`tias-train` コマンド: 患者単位 K-fold 交差検証で学習・評価する（SPEC_01「CLI」）.

サブコマンド:

- `cv --config PATH [--out-root runs]`: 5-fold 交差検証で学習・評価し，`runs/<実験ID>/` に書き出す
  （SPEC_02「交差検証の実行」）．最後に実験 ID を標準出力へ出す（`tias-remote pull <実験ID>` 用）
- `color-cv --config PATH [--out-root runs]`: 舌領域の色特徴量＋古典的分類器（設定の `color` 節）で
  `cv` と同じ分割・同じ評価の交差検証を行う．出力は `cv` と同じ規則（ただし学習曲線・
  `history.csv`・重みは出さない）．GPU は使わない

実データを読むためサーバで実行する（手元では合成データの設定 `configs/cv_synthetic.yaml`・
`configs/color_synthetic.yaml`）．
回収対象（`config.yaml`・`metrics.json`・`history.csv`・`confusion_matrix.csv`・`log.txt`・
`figures/*.png`）には患者 ID・ファイル名・画像単位の値を出さない．個票（`split.json`・
`predictions.csv`）はサーバ上にだけ置く．
"""

from __future__ import annotations

import argparse
import csv
import json
import logging
import os
import sys
import time
from collections.abc import Callable, Mapping, Sequence
from datetime import datetime
from pathlib import Path
from typing import Any

import numpy as np
import sklearn
import torch
import yaml
from pydantic import ValidationError

from tias_judge_model.cli.dataset import (
    CONFIG_FILENAME,
    FIGURES_DIRNAME,
    LOG_FILENAME,
    make_run_id,
    resolve_git_commit,
    resolved_config_document,
)
from tias_judge_model.data.classes import CLASS_SLUGS, label_to_id
from tias_judge_model.data.manifest import ManifestRecord, read_manifest
from tias_judge_model.evaluation.figures import (
    LEARNING_CURVES_FILENAME,
    confusion_matrix_filename,
    plot_confusion_matrix,
    plot_learning_curves,
)
from tias_judge_model.evaluation.metrics import (
    SUBSETS,
    aggregate_folds,
    mean_std,
    subset_metrics,
    sum_confusion_matrices,
)
from tias_judge_model.schema import ExperimentConfig, load_config
from tias_judge_model.training.color_cross_validation import (
    ColorCrossValidationResult,
    candidate_params,
    check_color_config,
    run_color_cross_validation,
)
from tias_judge_model.training.cross_validation import (
    ClassBalance,
    CrossValidationResult,
    FoldResult,
    fold_seed,
    run_cross_validation,
)

logger = logging.getLogger(__name__)

METRICS_FILENAME = "metrics.json"
HISTORY_FILENAME = "history.csv"
CONFUSION_MATRIX_FILENAME = "confusion_matrix.csv"
SPLIT_FILENAME = "split.json"
PREDICTIONS_FILENAME = "predictions.csv"
CHECKPOINTS_DIRNAME = "checkpoints"
# 混同行列（全 fold 合計）を書き出す部分集合（SPEC_02「評価の出し方」）．
CONFUSION_SUBSETS: tuple[str, ...] = ("common", "all")

_LOG_FORMAT = "%(asctime)s %(levelname)s %(name)s: %(message)s"

_EPILOG = """\
使用例（サーバ上で実行する．手元では configs/cv_synthetic.yaml で合成データを使う）:
  uv run tias-train cv --config configs/cv_with_mild.yaml
  uv run tias-remote run --name cv-with-mild -- tias-train cv --config configs/cv_with_mild.yaml
  uv run tias-train color-cv --config configs/color_synthetic.yaml
  uv run tias-remote run --name color1 --cpu -- tias-train color-cv --config <configs/color/*.yaml>

環境変数:
  TIAS_DATA_ROOT   data.root を上書き（tias-remote run が各サーバの data_dir を渡す）
  TIAS_GIT_COMMIT  実験 ID に入れる git の短縮ハッシュ（tias-remote run が渡す．無ければ nogit）
"""


# ---------------------------------------------------------------------------
# cv・color-cv
# ---------------------------------------------------------------------------

# サブコマンドの本体（設定・実験ディレクトリと，キーワード引数 run_id・git_commit・started）．
_Runner = Callable[..., None]


def _cmd_run(
    args: argparse.Namespace,
    runner: _Runner,
    check: Callable[[ExperimentConfig], object] | None = None,
) -> int:
    started = time.monotonic()
    config = load_config(args.config)
    if check is not None:
        # 設定の不備は実験ディレクトリを作る前に止める（空の実験ディレクトリを残さないため）．
        check(config)
    git_commit = resolve_git_commit(os.environ)
    run_id = make_run_id(datetime.now(), args.config.stem, git_commit)
    run_dir = args.out_root / run_id
    if run_dir.exists():
        logger.error("実験ディレクトリが既にある: %s（時間をおいて再実行する）", run_dir)
        return 1
    run_dir.mkdir(parents=True)

    file_handler = logging.FileHandler(run_dir / LOG_FILENAME, encoding="utf-8")
    file_handler.setFormatter(logging.Formatter(_LOG_FORMAT))
    root = logging.getLogger()
    root.addHandler(file_handler)
    try:
        runner(config, run_dir, run_id=run_id, git_commit=git_commit, started=started)
    except (ValueError, OSError) as exc:
        # log.txt にも失敗の理由を残すため，ハンドラを外す前にここで記録する．
        logger.error("%s", exc)
        return 1
    finally:
        root.removeHandler(file_handler)
        file_handler.close()
    sys.stdout.write(f"{run_id}\n")
    sys.stdout.flush()
    return 0


def _run_cv(
    config: ExperimentConfig, run_dir: Path, *, run_id: str, git_commit: str, started: float
) -> None:
    logger.info("交差検証の実験 ID: %s", run_id)
    _write_yaml(
        run_dir / CONFIG_FILENAME,
        _cv_config_document(config, run_id=run_id, git_commit=git_commit, command="train cv"),
    )

    records = read_manifest(config.data.manifest_path)
    result = run_cross_validation(config, records, run_dir / CHECKPOINTS_DIRNAME)

    history = history_rows(result.folds)
    _write_csv(run_dir / HISTORY_FILENAME, history)
    plot_learning_curves(history, run_dir / FIGURES_DIRNAME / LEARNING_CURVES_FILENAME)

    metrics = build_metrics(config, records, result, run_id=run_id, git_commit=git_commit)
    _write_evaluation_outputs(run_dir, config, records, result, metrics, started=started)


def _run_color_cv(
    config: ExperimentConfig, run_dir: Path, *, run_id: str, git_commit: str, started: float
) -> None:
    logger.info("色特徴分類器の交差検証の実験 ID: %s", run_id)
    _write_yaml(
        run_dir / CONFIG_FILENAME,
        _cv_config_document(config, run_id=run_id, git_commit=git_commit, command="train color-cv"),
    )

    records = read_manifest(config.data.manifest_path)
    result = run_color_cross_validation(config, records)

    # エポックが無いため history.csv と学習曲線は出さない．重みも保存しない．
    metrics = build_color_metrics(config, records, result, run_id=run_id, git_commit=git_commit)
    _write_evaluation_outputs(run_dir, config, records, result, metrics, started=started)


def _cv_config_document(
    config: ExperimentConfig, *, run_id: str, git_commit: str, command: str
) -> dict[str, Any]:
    document = resolved_config_document(
        config, run_id=run_id, git_commit=git_commit, command=command
    )
    document["seeds"] = {
        "split": config.split.seed,
        "folds": [fold_seed(config.split.seed, i) for i in range(config.split.n_folds)],
    }
    return document


def _write_evaluation_outputs(
    run_dir: Path,
    config: ExperimentConfig,
    records: Sequence[ManifestRecord],
    result: CrossValidationResult,
    metrics: dict[str, Any],
    *,
    started: float,
) -> None:
    # cv・color-cv に共通の書き出し（分割・個票・混同行列・図・metrics.json）とログの要約．
    _write_split(run_dir / SPLIT_FILENAME, config, result)
    _write_predictions(run_dir / PREDICTIONS_FILENAME, records, result)
    _write_confusion_csv(run_dir / CONFUSION_MATRIX_FILENAME, metrics["confusion_matrix_sum"])
    figures_dir = run_dir / FIGURES_DIRNAME
    for subset in CONFUSION_SUBSETS:
        matrix = metrics["confusion_matrix_sum"].get(subset)
        if matrix is not None:
            plot_confusion_matrix(
                matrix,
                f"{subset} (sum of {len(result.folds)} folds held-out)",
                figures_dir / confusion_matrix_filename(subset),
            )

    total_seconds = time.monotonic() - started
    metrics["timing"]["total_seconds"] = round(total_seconds, 2)
    _write_json(run_dir / METRICS_FILENAME, metrics)

    for subset in SUBSETS:
        summary = metrics["mean_std"].get(subset)
        if summary:
            logger.info(
                "%s: macro F1 %.4f ± %.4f・accuracy %.4f ± %.4f",
                subset,
                summary["macro_f1"]["mean"],
                summary["macro_f1"]["std"],
                summary["accuracy"]["mean"],
                summary["accuracy"]["std"],
            )
    breakdown = (metrics["pooled"].get("mild") or {}).get("prediction_breakdown")
    if breakdown:
        overall = breakdown["overall"]
        logger.info(
            "mild の予測の内訳（全 fold 合計 %d 件）: 淡紅 %.3f・元のクラス %.3f・その他 %.3f",
            overall["n"],
            overall["light_red"],
            overall["original"],
            overall["other"],
        )
    logger.info(
        "出力先: %s（全体の所要 %.1f 秒・うち画像の読み込み %.1f 秒）",
        run_dir,
        total_seconds,
        result.data_loading_seconds,
    )


def history_rows(folds: Sequence[FoldResult]) -> list[dict[str, Any]]:
    """`history.csv` の行（fold・エポックごとの損失・検証 macro F1・学習率・所要秒数）を作る.

    Args:
        folds: fold ごとの結果．

    Returns:
        行の dict のリスト．
    """
    rows: list[dict[str, Any]] = []
    for fold in folds:
        for h in fold.history:
            rows.append(
                {
                    "fold": fold.index,
                    "epoch": h.epoch,
                    "train_loss": round(h.train_loss, 6),
                    "val_loss": round(h.val_loss, 6),
                    "val_macro_f1": round(h.val_macro_f1, 6),
                    "lr": h.lr,
                    "epoch_seconds": round(h.seconds, 3),
                }
            )
    return rows


def build_metrics(
    config: ExperimentConfig,
    records: Sequence[ManifestRecord],
    result: CrossValidationResult,
    *,
    run_id: str,
    git_commit: str,
) -> dict[str, Any]:
    """`metrics.json` の内容（fold ごと・fold 平均 ± 標準偏差・全 fold 合算の指標と所要時間）を作る.

    集計値だけを含み，患者 ID・ファイル名・画像単位の値は含まない．

    Args:
        config: 実験設定．
        records: manifest の全レコード．
        result: 交差検証の結果．
        run_id: 実験 ID．
        git_commit: git の短縮ハッシュ（または `nogit`）．

    Returns:
        JSON に書き出せる dict（`timing.total_seconds` は呼び出し側で埋める）．
    """
    labels = np.array([label_to_id(r.label) for r in records], dtype=np.int64)
    severities = np.array([r.severity for r in records])

    fold_entries: list[dict[str, Any]] = []
    all_true: list[np.ndarray] = []
    all_pred: list[np.ndarray] = []
    all_sev: list[np.ndarray] = []
    for fold in result.folds:
        idx = fold.held_out_indices
        y_true, y_pred, sev = labels[idx], fold.held_out_pred, severities[idx]
        all_true.append(y_true)
        all_pred.append(y_pred)
        all_sev.append(sev)
        seconds = [h.seconds for h in fold.history]
        fold_entries.append(
            {
                "fold": fold.index,
                "seed": fold.seed,
                "best_epoch": fold.best_epoch,
                "best_val_macro_f1": fold.best_val_macro_f1,
                "counts": {
                    "patients": {
                        "train": len(fold.train_patients),
                        "val": len(fold.val_patients),
                        "held_out": len(fold.held_out_patients),
                    },
                    "images_by_group": fold.group_counts,
                },
                "timing": {
                    "train_seconds": round(fold.train_seconds, 2),
                    "eval_seconds": round(fold.eval_seconds, 2),
                    "epochs_run": len(fold.history),
                    "mean_epoch_seconds": round(float(np.mean(seconds)), 3) if seconds else None,
                },
                "subsets": subset_metrics(y_true, y_pred, sev),
                "class_balance": class_balance_entry(fold.class_balance),
            }
        )

    pooled = subset_metrics(
        np.concatenate(all_true), np.concatenate(all_pred), np.concatenate(all_sev)
    )
    mean_std_by_subset = {
        subset: aggregate_folds([f["subsets"][subset] for f in fold_entries]) for subset in SUBSETS
    }
    confusion_sum = {
        subset: (
            sum_confusion_matrices(
                [
                    f["subsets"][subset]["confusion_matrix"]
                    for f in fold_entries
                    if f["subsets"][subset]
                ]
            )
            if pooled.get(subset)
            else None
        )
        for subset in CONFUSION_SUBSETS
    }
    return {
        "run_id": run_id,
        "git_commit": git_commit,
        "classes": list(CLASS_SLUGS),
        "severities_trained": list(config.data.severities),
        "imbalance": {
            "class_weighting": config.train.class_weighting,
            "sampler": config.train.sampler,
        },
        "subsets": {name: list(sevs) for name, sevs in SUBSETS.items()},
        "n_folds": len(result.folds),
        "best_epoch": mean_std([f.best_epoch for f in result.folds]),
        "mean_std": mean_std_by_subset,
        "pooled": pooled,
        "confusion_matrix_sum": confusion_sum,
        "folds": fold_entries,
        "timing": {
            "total_seconds": None,
            "data_loading_seconds": round(result.data_loading_seconds, 2),
            "train_seconds_total": round(sum(f.train_seconds for f in result.folds), 2),
            "eval_seconds_total": round(sum(f.eval_seconds for f in result.folds), 2),
        },
        "environment": {
            "device": result.device,
            "model_name": config.model.name,
            "pretrained": config.model.pretrained,
            "num_parameters": result.num_parameters,
            "image_size": config.train.image_size,
            "torch_version": torch.__version__,
        },
    }


def build_color_metrics(
    config: ExperimentConfig,
    records: Sequence[ManifestRecord],
    result: ColorCrossValidationResult,
    *,
    run_id: str,
    git_commit: str,
) -> dict[str, Any]:
    """色特徴分類器の `metrics.json` の内容を作る（`build_metrics` ＋ 選んだハイパーパラメータ）.

    `build_metrics` と同じ指標・集計に次の変更を加える（集計値だけを含み，患者 ID・ファイル名・
    画像単位の値は含まない）．

    - `environment`: モデル名を `color_<classifier>` にし，特徴量の次元数・縮小の上限・
      scikit-learn の版を記録する（CNN 固有のパラメータ数・入力サイズ・事前学習は持たない）
    - `best_epoch`（トップレベル・fold ごと）: エポックが無いため None
    - fold ごとに `selected_params`（選んだハイパーパラメータ）と `candidates`（候補ごとの
      検証 macro F1・学習秒数）を加える
    - `color_classifier`: 分類器・候補の一覧・特徴量の名前
    - `timing.feature_extraction_seconds`: 画像の読み込みと特徴抽出の秒数
      （`data_loading_seconds` と同じ値）

    Args:
        config: 実験設定（`color` 節を含む）．
        records: manifest の全レコード．
        result: 色特徴分類器の交差検証の結果．
        run_id: 実験 ID．
        git_commit: git の短縮ハッシュ（または `nogit`）．

    Returns:
        JSON に書き出せる dict（`timing.total_seconds` は呼び出し側で埋める）．

    Raises:
        ValueError: 設定に `color` 節が無い場合．
    """
    color = check_color_config(config)
    metrics = build_metrics(config, records, result, run_id=run_id, git_commit=git_commit)
    metrics["best_epoch"] = None
    selections = {s.fold: s for s in result.selections}
    for entry in metrics["folds"]:
        selection = selections[entry["fold"]]
        entry["best_epoch"] = None
        entry["selected_params"] = selection.selected_params
        entry["candidates"] = [
            {
                "params": c.params,
                "val_macro_f1": c.val_macro_f1,
                "fit_seconds": round(c.fit_seconds, 3),
            }
            for c in selection.candidates
        ]
    metrics["color_classifier"] = {
        "classifier": color.classifier,
        "candidates": candidate_params(color),
        "num_features": len(result.feature_names),
        "feature_names": list(result.feature_names),
    }
    metrics["timing"]["feature_extraction_seconds"] = round(result.data_loading_seconds, 2)
    metrics["environment"] = {
        "device": result.device,
        "model_name": result.model_name,
        "num_features": len(result.feature_names),
        "max_side": color.max_side,
        "sklearn_version": sklearn.__version__,
    }
    return metrics


def class_balance_entry(balance: ClassBalance | None) -> dict[str, Any] | None:
    """fold の `class_balance`（学習のクラス別枚数・クラス重み・抜き出し確率）を作る.

    Args:
        balance: fold の学習に使ったクラス不均衡の対策（None なら記録なし）．

    Returns:
        クラス slug をキーとする dict を含む dict．`class_weights` は `class_weighting` が
        `none` のとき，`class_draw_probabilities` は `sampler` が `random` のとき None．
    """
    if balance is None:
        return None

    def by_class(values: Sequence[float] | None) -> dict[str, float] | None:
        if values is None:
            return None
        return {slug: round(float(v), 6) for slug, v in zip(CLASS_SLUGS, values, strict=True)}

    return {
        "class_weighting": balance.class_weighting,
        "sampler": balance.sampler,
        "class_counts": dict(zip(CLASS_SLUGS, balance.class_counts, strict=True)),
        "class_weights": by_class(balance.class_weights),
        "class_draw_probabilities": by_class(balance.class_draw_probabilities),
    }


def _write_split(path: Path, config: ExperimentConfig, result: CrossValidationResult) -> None:
    # 患者 ID を含む（回収しない）．
    _write_json(
        path,
        {
            "seed": config.split.seed,
            "n_folds": config.split.n_folds,
            "inner_val_ratio": config.train.inner_val_ratio,
            "severities_trained": list(config.data.severities),
            "folds": [
                {
                    "fold": f.index,
                    "seed": f.seed,
                    "train": list(f.train_patients),
                    "val": list(f.val_patients),
                    "held_out": list(f.held_out_patients),
                }
                for f in result.folds
            ],
        },
    )


def _write_predictions(
    path: Path, records: Sequence[ManifestRecord], result: CrossValidationResult
) -> None:
    # 画像ごとの予測（個票．患者 ID・ファイル名を含む．回収しない）．
    rows: list[dict[str, Any]] = []
    for fold in result.folds:
        for i, probs in zip(fold.held_out_indices, fold.held_out_probs, strict=True):
            record = records[int(i)]
            row: dict[str, Any] = {
                "fold": fold.index,
                "patient_id": record.patient_id,
                "image_path": record.image_path,
                "label": record.label,
                "severity": record.severity,
                "pred": CLASS_SLUGS[int(np.argmax(probs))],
            }
            row.update(
                {
                    f"prob_{slug}": round(float(p), 6)
                    for slug, p in zip(CLASS_SLUGS, probs, strict=True)
                }
            )
            rows.append(row)
    _write_csv(path, rows)


def _write_confusion_csv(path: Path, matrices: Mapping[str, list[list[int]] | None]) -> None:
    rows: list[dict[str, Any]] = []
    for subset, matrix in matrices.items():
        if matrix is None:
            continue
        for slug, counts in zip(CLASS_SLUGS, matrix, strict=True):
            row: dict[str, Any] = {"subset": subset, "true": slug}
            row.update({f"pred_{p}": c for p, c in zip(CLASS_SLUGS, counts, strict=True)})
            rows.append(row)
    _write_csv(path, rows)


def _write_csv(path: Path, rows: Sequence[Mapping[str, Any]]) -> None:
    with path.open("w", encoding="utf-8", newline="") as f:
        if not rows:
            return
        writer = csv.DictWriter(f, fieldnames=list(rows[0]), lineterminator="\n")
        writer.writeheader()
        writer.writerows(rows)


def _write_json(path: Path, data: Mapping[str, Any]) -> None:
    with path.open("w", encoding="utf-8", newline="\n") as f:
        json.dump(data, f, ensure_ascii=False, indent=2)
        f.write("\n")


def _write_yaml(path: Path, data: Mapping[str, Any]) -> None:
    with path.open("w", encoding="utf-8", newline="\n") as f:
        yaml.safe_dump(dict(data), f, allow_unicode=True, sort_keys=False)


# ---------------------------------------------------------------------------
# エントリポイント
# ---------------------------------------------------------------------------


def _parse_args(argv: Sequence[str] | None) -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        prog="tias-train",
        description="患者単位 K-fold 交差検証で学習・評価する",
        epilog=_EPILOG,
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    sub = parser.add_subparsers(dest="subcommand", required=True)
    p_cv = sub.add_parser(
        "cv",
        help="K-fold 交差検証で学習・評価して runs/<実験ID>/ に書き出す",
        description=(
            "manifest 全体で患者単位・クラス × 程度の層化の K-fold 分割を作り，各 fold で学習し"
            "held_out（全程度）で評価する．最後に実験 ID を標準出力へ出す．"
        ),
    )
    p_color = sub.add_parser(
        "color-cv",
        help="色特徴量＋古典的分類器で cv と同じ分割・評価の交差検証をする",
        description=(
            "舌領域の色特徴量（L*a*b* 系の統計量・ヒストグラム）と設定の color 節の分類器で，"
            "cv と同じ分割・同じ評価の交差検証を行う．各 fold でハイパーパラメータの候補を内側の"
            "検証 macro F1 で選ぶ．GPU は使わない．最後に実験 ID を標準出力へ出す．"
        ),
    )
    for p in (p_cv, p_color):
        p.add_argument("--config", type=Path, required=True, help="実験設定 YAML．")
        p.add_argument(
            "--out-root",
            type=Path,
            default=Path("runs"),
            help="実験ディレクトリの親（既定: runs）．",
        )
    return parser.parse_args(argv)


def _configure_logging() -> None:
    # ログは標準エラーへ（標準出力は実験 ID の受け渡しに使うため汚さない）．
    handler = logging.StreamHandler(sys.stderr)
    handler.setFormatter(logging.Formatter(_LOG_FORMAT))
    logging.basicConfig(level=logging.INFO, handlers=[handler])


def main(argv: Sequence[str] | None = None) -> int:
    """tias-train コマンドのメインエントリポイント.

    Args:
        argv: 引数リスト．None のとき sys.argv[1:] を使用．

    Returns:
        終了コード（成功 0・入力の問題 1）．
    """
    args = _parse_args(argv)
    _configure_logging()
    try:
        if args.subcommand == "color-cv":
            return _cmd_run(args, _run_color_cv, check=check_color_config)
        return _cmd_run(args, _run_cv)
    except ValidationError as exc:
        logger.error("設定が不正である:\n%s", exc)
        return 1
    except (ValueError, OSError) as exc:
        logger.error("%s", exc)
        return 1


if __name__ == "__main__":
    sys.exit(main())