"""Train + freeze the v2 judgement model exactly as validated in data/ticksim/PREREG_v2.md.

usage: python -m judge.model            -> writes data/judge/model_v2.joblib (+ .json metadata)
"""
from __future__ import annotations

import json
from pathlib import Path

import joblib
import numpy as np
import pandas as pd
from sklearn.ensemble import HistGradientBoostingClassifier

from judge.engine import FEATS, SETUP_I

ROOT = Path(__file__).resolve().parent.parent
TICKSIM = ROOT / "data" / "ticksim"
OUT = ROOT / "data" / "judge"
PARAMS = dict(max_iter=200, learning_rate=0.05, max_leaf_nodes=15, min_samples_leaf=40, l2_regularization=1.0,
              random_state=0)


def _prep(df: pd.DataFrame) -> pd.DataFrame:
    daily = pd.read_pickle(ROOT / "data" / "ross" / "daily_all.pkl")
    med = daily.sort_values("date").groupby("symbol").v.transform(lambda s: s.shift(1).rolling(20, min_periods=5).median())
    daily = daily.assign(med20=med)[["symbol", "date", "gap", "med20"]]
    if "gap" not in df:
        df = df.merge(daily, on=["symbol", "date"], how="left")
        df["cumvol_vs_med"] = df.cum_vol / df.med20
    df["setup_i"] = df.setup.map(SETUP_I)
    df["gap_known"] = np.where(df.rth == 1, df.gap, np.nan)
    return df


def train() -> dict:
    s1 = _prep(pd.read_pickle(TICKSIM / "ml_data.pkl"))
    s2 = _prep(pd.read_pickle(TICKSIM / "ml_fresh_data.pkl"))
    tr = pd.concat([s1, s2], ignore_index=True)
    tr = tr[tr.date < "2026-08-01"]
    m = HistGradientBoostingClassifier(**PARAMS).fit(tr[FEATS], (tr.r > 0).astype(int))
    cutoff = float(np.quantile(m.predict_proba(tr[tr.rth == 1][FEATS])[:, 1], 0.8))
    meta = dict(version="v2", trained_on="samples 1+2, signals dated < 2026-08-01", n_train=int(len(tr)),
                cutoff=cutoff, features=FEATS, params=PARAMS, validation="PREREG_v2 PASS: sample-3 Aug-Sep RTH "
                "top20% n=332 +1.11%/trade (day-block lo +0.54%)")
    return dict(model=m, **meta)


def save(art: dict) -> Path:
    OUT.mkdir(parents=True, exist_ok=True)
    p = OUT / "model_v2.joblib"
    joblib.dump(art, p)
    (OUT / "model_v2.json").write_text(json.dumps({k: v for k, v in art.items() if k != "model"}, indent=2, default=str))
    return p


def load(path: Path | None = None) -> dict:
    return joblib.load(path or OUT / "model_v2.joblib")


def score(art: dict, feats: dict) -> float:
    x = pd.DataFrame([{k: feats.get(k, np.nan) for k in art["features"]}])
    return float(art["model"].predict_proba(x)[:, 1][0])


if __name__ == "__main__":
    a = train()
    print("saved", save(a), "cutoff", round(a["cutoff"], 4), "n_train", a["n_train"])