# -*- coding: utf-8 -*-
"""5段階グレード基準での学習・検証（新モデル / 現行は上書きしない）。

グレード基準（5段階）:
    0–39 / 40–59 / 60–84 / 85–99 / 100
後段は連続値回帰のため回帰自体は現行と同一。本スクリプトは
  - この基準でのグレード一致率を 5-fold OOF で算出
  - 帯ごとの分布と混同表を出力
  - 回帰モデルを models_grade5/ に別保存（grade基準をmetadataに記録）
する。LLM出力(training_dataset.csv)は書き換えない。
"""
from __future__ import annotations
import argparse
import json
import sys
from pathlib import Path

import numpy as np
import pandas as pd
from sklearn.linear_model import LinearRegression
from sklearn.model_selection import KFold
from sklearn.pipeline import Pipeline

HERE = Path(__file__).resolve().parent
ROOT = HERE.parent
sys.path.insert(0, str(HERE / "core"))
import soup_scoring as S  # noqa: E402

DEFAULT_INPUT = ROOT / "local_batch_api" / "output_full" / "training_dataset.csv"
OUTDIR = HERE / "validation_results_grade5"
MODEL_DIR = HERE / "models_grade5"
SEED, N = 42, 5

BANDS = ["0-39", "40-59", "60-84", "85-99", "100"]


def grade5(score) -> np.ndarray:
    v = np.asarray(score, dtype=float)
    out = np.full(v.shape, "0-39", dtype="<U8")
    out[v >= 40] = "40-59"
    out[v >= 60] = "60-84"
    out[v >= 85] = "85-99"
    out[v >= 100] = "100"
    return out


def oof(num, cat, y, X, kf):
    pred = np.zeros(len(y))
    for tr, te in kf.split(X):
        p = Pipeline([("prep", S._preprocessor(num, cat)), ("m", LinearRegression())])
        p.fit(X.iloc[tr], y[tr])
        pred[te] = np.clip(p.predict(X.iloc[te]), 0, 100)
    return pred


def confusion(yt, yp) -> pd.DataFrame:
    m = pd.DataFrame(0, index=BANDS, columns=BANDS, dtype=int)
    for a, b in zip(yt, yp):
        m.loc[a, b] += 1
    return m


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--input", default=str(DEFAULT_INPUT))
    args = ap.parse_args()
    OUTDIR.mkdir(parents=True, exist_ok=True)
    MODEL_DIR.mkdir(parents=True, exist_ok=True)

    raw = pd.read_csv(args.input)
    f = S.prepare(raw, require_targets=True)
    g = f[f["assessable"] == 1].reset_index(drop=True)
    fsets = S.feature_sets()
    kf = KFold(N, shuffle=True, random_state=SEED)

    print("=" * 64)
    print(" 5段階グレード基準 検証 (5-fold OOF / 線形回帰)")
    print(" 基準: 0-39 / 40-59 / 60-84 / 85-99 / 100")
    print("=" * 64)
    print(f"採点対象(assessable=1) n={len(g)}\n")

    rows = []
    meta = {"grade_scheme": BANDS, "n_total": int(len(f)), "n_scoreable": int(len(g)), "models": {}}
    for target, tname in [("human_total_score", "total"), ("human_soup_score", "soup")]:
        fs = fsets[S.TARGETS[target]]
        num, cat = list(fs.numeric), list(fs.categorical)
        y = g[target].to_numpy()
        pred = oof(num, cat, y, g[num + cat], kf)
        ae = np.abs(y - pred)
        gt, gp = grade5(y), grade5(pred)
        match = (gt == gp).mean() * 100
        rows.append({"target": tname, "n": len(g), "MAE": round(ae.mean(), 2),
                     "5段階一致%": round(match, 1), "±15点以内%": round((ae <= 15).mean() * 100, 1)})

        # 人間スコアの帯分布
        dist = pd.Series(gt).value_counts().reindex(BANDS).fillna(0).astype(int)
        print(f"[{tname}] 人間スコアの帯分布: " + " / ".join(f"{b}:{dist[b]}" for b in BANDS))
        print(f"   MAE={ae.mean():.2f}  5段階グレード一致={match:.1f}%  ±15点以内={(ae<=15).mean()*100:.1f}%")
        cm = confusion(gt, gp)
        cm.to_csv(OUTDIR / f"confusion_{tname}.csv", encoding="utf-8-sig")
        print("   混同表(行=人間, 列=予測):")
        print(cm.to_string().replace("\n", "\n   "))
        print()

        # 別モデルとして保存（回帰は現行と同一、grade基準をmetadataに記録）
        pipe = Pipeline([("prep", S._preprocessor(num, cat)), ("m", LinearRegression())])
        pipe.fit(g[num + cat], y)
        import joblib
        joblib.dump({"pipeline": pipe, "feature_set": fs.name, "numeric": num,
                     "categorical": cat, "target": target, "grade_scheme": BANDS},
                    MODEL_DIR / f"{target}.pkl")
        meta["models"][target] = {"feature_set": fs.name, "path": f"{target}.pkl", "model": "LinearRegression"}

    (MODEL_DIR / "metadata.json").write_text(json.dumps(meta, ensure_ascii=False, indent=2), encoding="utf-8")
    res = pd.DataFrame(rows)[["target", "n", "MAE", "5段階一致%", "±15点以内%"]]
    res.to_csv(OUTDIR / "grade5_validation.csv", index=False, encoding="utf-8-sig")
    print("=== サマリ ===")
    print(res.to_string(index=False))
    print("\nモデル保存:", MODEL_DIR)
    print("検証結果:", OUTDIR)


if __name__ == "__main__":
    main()