# -*- coding: utf-8 -*-
"""濃度モデルのキャリブレーション補正を学習する。

full_density_eval.csv（全件のモデル推定 vs 人間ラベル）を読み、
5分割CVで複数の補正手法を比較。ベースラインを実際に上回る手法だけを採用し、
calibration.json（予測値→補正値のマップ）を保存する。
density_resolver はこの calibration.json を自動で読み込んで適用する。
"""

from __future__ import annotations

import argparse
import csv
import json
from collections import Counter
from pathlib import Path

import numpy as np

ROOT = Path(__file__).resolve().parents[1]


def load_pairs(csv_path: Path, conf: str):
    out = []
    for r in csv.DictReader(open(csv_path, encoding="utf-8-sig")):
        if r["status"] != "ok" or not r["model_value"] or not r["human"]:
            continue
        if conf != "all" and r["confidence"] != conf:
            continue
        out.append((int(r["model_value"]), int(r["human"])))
    return np.array(out)


def metrics(pred, true):
    d = np.abs(pred - true)
    return dict(exact=float((d == 0).mean()), pm1=float((d <= 1).mean()), mae=float(d.mean()))


def fit_mode(P, T):
    m = {int(p): Counter(T[P == p]).most_common(1)[0][0] for p in np.unique(P)}
    return m


def fit_median(P, T):
    return {int(p): int(round(np.median(T[P == p]))) for p in np.unique(P)}


def fit_linear(P, T):
    a, b = np.polyfit(P, T, 1)
    return {int(p): int(np.clip(round(a * p + b), 10, 14)) for p in range(10, 15)}


def apply_map(m, x):
    return np.array([m.get(int(v), int(v)) for v in x])


def cv(data, fit_fn, k=5, seed=42):
    rng = np.random.RandomState(seed)
    idx = rng.permutation(len(data))
    folds = np.array_split(idx, k)
    res = []
    for i in range(k):
        te = folds[i]
        tr = np.concatenate([folds[j] for j in range(k) if j != i])
        m = fit_fn(data[tr, 0], data[tr, 1])
        res.append(metrics(apply_map(m, data[te, 0]), data[te, 1]))
    return {k2: float(np.mean([r[k2] for r in res])) for k2 in res[0]}


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--csv", default=str(ROOT / "local_batch_api" / "full_density_eval.csv"))
    ap.add_argument("--out", default=str(ROOT / "local_batch_api" / "calibration.json"))
    ap.add_argument("--fit-on", default="all", choices=["all", "high"],
                    help="補正マップを学習する対象（既定all）")
    ap.add_argument("--select-metric", default="mae", choices=["mae", "exact", "pm1"])
    ap.add_argument("--min-gain", type=float, default=0.005, help="採用に必要な改善量")
    args = ap.parse_args()

    csv_path = Path(args.csv)
    methods = {"mode": fit_mode, "median": fit_median, "linear": fit_linear}

    print("================ キャリブレーション検証 (5-fold CV) ================")
    results = {}
    for name in ["high", "all"]:
        data = load_pairs(csv_path, name)
        base = metrics(data[:, 0], data[:, 1])
        print(f"\n[{name}] n={len(data)}  baseline: exact={base['exact']:.3f} pm1={base['pm1']:.3f} mae={base['mae']:.3f}")
        results[name] = {"baseline": base, "n": len(data), "methods": {}}
        for mname, fn in methods.items():
            cvm = cv(data, fn)
            results[name]["methods"][mname] = cvm
            print(f"   {mname:8} CV: exact={cvm['exact']:.3f} pm1={cvm['pm1']:.3f} mae={cvm['mae']:.3f}")

    # 採用判定: fit-on 対象で select-metric を最も改善する手法。改善が閾値未満なら identity。
    data = load_pairs(csv_path, args.fit_on)
    base = metrics(data[:, 0], data[:, 1])
    better = lambda cand, ref: (ref - cand) if args.select_metric == "mae" else (cand - ref)
    best_name, best_gain, best_map = None, 0.0, None
    for mname, fn in methods.items():
        cvm = results[args.fit_on]["methods"][mname]
        gain = better(cvm[args.select_metric], base[args.select_metric])
        if gain > best_gain:
            best_gain, best_name = gain, mname
    if best_name and best_gain >= args.min_gain:
        best_map = methods[best_name](data[:, 0], data[:, 1])
        full_map = {str(k): int(v) for k, v in best_map.items()}
        payload = {"enabled": True, "method": best_name, "fit_on": args.fit_on,
                   "select_metric": args.select_metric, "cv_gain": round(best_gain, 4),
                   "map": full_map, "baseline": base,
                   "cv": results[args.fit_on]["methods"][best_name]}
        Path(args.out).write_text(json.dumps(payload, ensure_ascii=False, indent=2), encoding="utf-8")
        print(f"\n採用: {best_name} (改善 {args.select_metric} {best_gain:+.3f}) -> {args.out}")
        print(f"補正マップ: {full_map}")
    else:
        payload = {"enabled": False, "reason": "CVで有意な改善が出なかったため補正なし（恒等写像）",
                   "baseline": base, "cv": results[args.fit_on]["methods"]}
        Path(args.out).write_text(json.dumps(payload, ensure_ascii=False, indent=2), encoding="utf-8")
        print(f"\n採用なし: ベースラインを上回る補正が無かった -> {args.out} (enabled=false)")
        print("（モデルは既に校正済み。補正は適用されません）")


if __name__ == "__main__":
    raise SystemExit(main())