"""Score trained Laya checkpoints on the held-out test set (future anchors x unseen hosts).

Also answers the go/no-go question for the text model: does its score add lift on top
of the telemetry LightGBM? The stacker is trained ONLY on the calibration hosts, which
the Laya model never saw, and is compared against a LightGBM trained on those same rows
without the text score - so the comparison isolates the text signal.
"""
from __future__ import annotations

import argparse
import json
import os
from pathlib import Path

import lightgbm as lgb
import numpy as np
import pandas as pd
import torch
from safetensors.torch import load_file
from transformers import AutoTokenizer

from laya.common import build_model, build_sequence

from common import NUMERIC, SPLITS, TEST_HOST_BUCKET, report, split
from train_laya import CALIB_HOSTS, QUESTIONS, collate, logits_of, p_true, state_of


def load_ckpt(path: Path, device):
    cfg = json.loads((path / "rl_agent_config.json").read_text())
    tok = AutoTokenizer.from_pretrained(path / "tokenizer")
    model = build_model(cfg, encoder_dir=str(path / "encoder"))
    model.load_state_dict({k: v.float() for k, v in load_file(str(path / "model.safetensors")).items()})
    model.encoder.config.reference_compile = False
    return model.to(device), tok, cfg


def score(model, tok, cfg, df: pd.DataFrame, variant: str, q: str, device) -> np.ndarray:
    # text variant: identical state per resource -> score each resource once
    key = "resource" if variant == "text" else None
    src = df.drop_duplicates("resource") if key else df
    items = []
    for _, row in src.iterrows():
        seq, markers = build_sequence(tok, json.dumps(state_of(row, variant), ensure_ascii=False),
                                      {"t": "noul", "ins": QUESTIONS[q]["instructions"], "crit": {}},
                                      cfg["max_len"], cfg["head_max_len"])
        items.append({"ids": seq, "markers": markers, "qtype": 2, "target": [0, 0]})
    p = p_true(logits_of(model, items, tok.pad_token_id, device, bs=64), cfg["temperature"][2])
    if key:
        m = dict(zip(src.resource, p))
        return df.resource.map(m).to_numpy()
    return p


def main():
    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")
    df["network_cat"] = df.network.astype("category")
    out = {"checkpoint": str(args.ckpt), "variant": variant}

    for h in (7, 14):
        _, _, te = split(df, h)
        p = score(model, tok, cfg, te, variant, f"paid_{h}d", device)
        out[f"paid_{h}d"] = report(te, p, h)
        print(f"paid_{h}d laya-{variant}:", json.dumps(out[f"paid_{h}d"]))
        if h == 7:
            np.save(args.ckpt / "test_pred_paid7d.npy", p)

    # ---- stacking test (paid_7d) ----
    lab = "paid_7d"
    s = SPLITS[7]
    cal = df[df.anchor_idx.isin(s["train_anchors"]) & df.host_bucket.isin(CALIB_HOSTS) & (df[lab] >= 0)].copy()
    _, _, te = split(df, 7)
    te = te.copy()
    cal["laya"] = score(model, tok, cfg, cal, variant, lab, device)
    te["laya"] = np.load(args.ckpt / "test_pred_paid7d.npy")
    feats = NUMERIC + ["network_cat"]
    res = {}
    preds = {}
    for name, fs in (("gbm_calib_only", feats), ("gbm_calib+laya", feats + ["laya"])):
        # hold out calib hosts 78-79 as early-stopping set
        es = cal.host_bucket >= 78
        m = lgb.LGBMClassifier(n_estimators=600, learning_rate=0.03, num_leaves=15, min_child_samples=50,
                               subsample=0.8, subsample_freq=1, colsample_bytree=0.8, verbose=-1)
        m.fit(cal.loc[~es, fs], cal.loc[~es, lab], eval_set=[(cal.loc[es, fs], cal.loc[es, lab])],
              callbacks=[lgb.early_stopping(50, verbose=False)])
        preds[name] = m.predict_proba(te[fs])[:, 1]
        res[name] = report(te, preds[name], 7)
        print(f"stack {name}:", json.dumps(res[name]))
    # paired bootstrap over HOSTS (rows within a host are correlated), 1000 resamples
    from sklearn.metrics import roc_auc_score
    rng = np.random.default_rng(0)
    hosts = te.host.unique()
    y = te[lab].to_numpy()
    segs = {"all": np.ones(len(te), bool), "new_listings": (te.is_new_listing == 1).to_numpy()}
    by_host = {h: np.flatnonzero(te.host.to_numpy() == h) for h in hosts}
    ci = {}
    for seg, mask in segs.items():
        d = []
        for _ in range(1000):
            idx = np.concatenate([by_host[h] for h in rng.choice(hosts, len(hosts))])
            idx = idx[mask[idx]]
            if len(np.unique(y[idx])) < 2:
                continue
            d.append(roc_auc_score(y[idx], preds["gbm_calib+laya"][idx]) - roc_auc_score(y[idx], preds["gbm_calib_only"][idx]))
        ci[seg] = dict(delta_auc_median=round(float(np.median(d)), 4), ci95=[round(float(np.percentile(d, 2.5)), 4),
                       round(float(np.percentile(d, 97.5)), 4)], p_delta_le_0=round(float(np.mean(np.array(d) <= 0)), 4))
        print(f"stack lift {seg}: median dAUC {ci[seg]['delta_auc_median']} 95% CI {ci[seg]['ci95']} P(<=0) {ci[seg]['p_delta_le_0']}")
    res["bootstrap_host_resampled"] = ci
    out["stacking_paid_7d"] = res
    (args.ckpt / "eval.json").write_text(json.dumps(out, indent=1))


if __name__ == "__main__":
    main()