"""Post-hoc isotonic recalibration for a checkpoint (used by the paid API, not the Laya format).

Laya checkpoints carry one temperature per question type, which cannot fix the shape of the
miscalibration we measured (over-confident at the top only). The API applies this monotone
map after the model, so AUC is unchanged by construction; only the probabilities move.
Fitted on calibration hosts (never trained on), scored on the held-out test set.
Output: <ckpt>/isotonic_paid7d.json  ({"x": [...], "y": [...]} for np.interp).
"""
from __future__ import annotations

import argparse
import json
from pathlib import Path

import numpy as np
import pandas as pd
import torch
from sklearn.isotonic import IsotonicRegression

from common import SPLITS, metrics, split
from evaluate import load_ckpt, score
from train_laya import CALIB_HOSTS

ap = argparse.ArgumentParser()
ap.add_argument("--ckpt", type=Path, required=True)
args = ap.parse_args()

device = torch.device("cuda")
model, tok, cfg = load_ckpt(args.ckpt, device)
variant = cfg["x402_variant"]
df = pd.read_parquet("data/rows.parquet")
cal = df[df.anchor_idx.isin(SPLITS[7]["train_anchors"]) & df.host_bucket.isin(CALIB_HOSTS) & (df.paid_7d >= 0)]
_, _, te = split(df, 7)

p_cal = score(model, tok, cfg, cal, variant, "paid_7d", device)
p_te = np.load(args.ckpt / "test_pred_paid7d.npy")
iso = IsotonicRegression(out_of_bounds="clip", y_min=0.001, y_max=0.999).fit(p_cal, cal.paid_7d)
p_iso = iso.predict(p_te)

y = te.paid_7d.to_numpy()
before, after = metrics(y, p_te), metrics(y, p_iso)
print("test before:", before)
print("test after :", after)
out = {"fit_on": f"{len(cal):,} calibration-host rows", "x": iso.X_thresholds_.round(5).tolist(),
       "y": iso.y_thresholds_.round(5).tolist(), "test_before": before, "test_after": after}
(args.ckpt / "isotonic_paid7d.json").write_text(json.dumps(out, indent=1))
np.save(args.ckpt / "test_pred_paid7d_isotonic.npy", p_iso)