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

サブコマンド:

- `cv --config PATH [--out-root runs]`: 5-fold 交差検証で学習・評価し，`runs/<実験ID>/` に書き出す
  （SPEC_02「交差検証の実行」）．最後に実験 ID を標準出力へ出す（`tias-remote pull <実験ID>` 用）

実データを読むためサーバで実行する（手元では合成データの設定 `configs/cv_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 Mapping, Sequence
from datetime import datetime
from pathlib import Path
from typing import Any

import numpy as np
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.cross_validation import (
    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

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


# ---------------------------------------------------------------------------
# cv
# ---------------------------------------------------------------------------


def _cmd_cv(args: argparse.Namespace) -> int:
    started = time.monotonic()
    config = load_config(args.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:
        _run_cv(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)
    document = resolved_config_document(
        config, run_id=run_id, git_commit=git_commit, command="train cv"
    )
    document["seeds"] = {
        "split": config.split.seed,
        "folds": [fold_seed(config.split.seed, i) for i in range(config.split.n_folds)],
    }
    _write_yaml(run_dir / CONFIG_FILENAME, document)

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

    _write_split(run_dir / SPLIT_FILENAME, config, result)
    _write_predictions(run_dir / PREDICTIONS_FILENAME, records, result)
    history = history_rows(result.folds)
    _write_csv(run_dir / HISTORY_FILENAME, history)

    metrics = build_metrics(config, records, result, run_id=run_id, git_commit=git_commit)
    _write_confusion_csv(run_dir / CONFUSION_MATRIX_FILENAME, metrics["confusion_matrix_sum"])
    figures_dir = run_dir / FIGURES_DIRNAME
    plot_learning_curves(history, figures_dir / LEARNING_CURVES_FILENAME)
    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),
            }
        )

    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),
        "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 _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_cv.add_argument("--config", type=Path, required=True, help="実験設定 YAML．")
    p_cv.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:
        return _cmd_cv(args)
    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())