"""Fine-tune Laya on x402 paid-demand labels (single GPU port of Laya's own RLCD recipe).

The training loop is a faithful single-GPU port of the upstream notebook
(notebooks/laya_finetune_typed_decisions_2xT4_kaggle.ipynb, laya 0.3.6):
RLCD policy-gradient on a strictly proper scoring rule + soft cross-entropy
guidance, then per-type temperature calibration on a HELD-OUT host slice.
The output directory loads with `laya.load(path)`.

Variants (the only thing that differs is the STATE the model sees):
  text : the public listing only - url, description, price, network.
         Works on a listing with no history, including one not yet listed.
  full : the listing + the Bazaar's own telemetry at the anchor day.
"""
from __future__ import annotations

import argparse
import json
import math
import os
import random
import time
from pathlib import Path

import numpy as np
import pandas as pd
import torch
from huggingface_hub import snapshot_download
from safetensors.torch import load_file, save_file
from transformers import AutoTokenizer

from laya.agent import _fix_tokenizer_config
from laya.common import QTYPES, build_model, build_sequence, proper_reward, render_options

from common import SPLITS, TEST_HOST_BUCKET

BASE = "convaiinnovations/laya"
CALIB_HOSTS = range(70, 80)  # train-side hosts held out for temperature fit + model selection

QUESTIONS = {
    "paid_7d": {"type": "noul", "instructions":
                "Will this x402 pay-per-call API receive at least one paid call in the next 7 days?"},
    "paid_14d": {"type": "noul", "instructions":
                 "Will this x402 pay-per-call API receive at least one paid call in the next 14 days?"},
}


def fmt_num(x, nd=1):
    if x is None or (isinstance(x, float) and math.isnan(x)):
        return "unknown"
    return f"{x:.{nd}f}" if isinstance(x, float) else str(x)


def state_of(row: pd.Series, variant: str) -> dict:
    """The model's input. Keep in sync with app.py / predict.py (published)."""
    s = {
        "resource": row.resource,
        "description": row.description[:600],
        "price_usd": f"{row.price_usd:.6g}",
        "network": row.network,
    }
    if variant == "full":
        s.update({
            "paid_calls_30d": int(row.calls_30d),
            "unique_payers_30d": int(row.payers_30d),
            "days_since_last_paid_call": fmt_num(row.days_since_last_call),
            "days_since_listing_updated": fmt_num(row.days_since_last_update),
            "host_listing_count": int(row.host_listings),
            "host_share_called_in_last_7d": fmt_num(row.host_active_rate_7d, 2),
        })
    return s


def make_items(df: pd.DataFrame, variant: str, tok, cfg, soft: bool) -> list[dict]:
    """soft=True: one item per resource with target = fraction of anchors labelled 1
    (text variant - the text is identical across anchors, so the label rate IS the target).
    soft=False: one item per row."""
    items = []
    if soft:
        groups = df.groupby("resource", sort=False)
        rows = [(g.iloc[-1], {q: g[q][g[q] >= 0].mean() for q in QUESTIONS}) for _, g in groups]
    else:
        rows = [(r, {q: r[q] for q in QUESTIONS}) for _, r in df.iterrows()]
    for row, tgt in rows:
        state = json.dumps(state_of(row, variant), ensure_ascii=False)
        for q, qd in QUESTIONS.items():
            y = tgt[q]
            if y is None or (isinstance(y, float) and math.isnan(y)) or y < 0:
                continue
            k = len(render_options({"t": "noul", "crit": {}}))
            seq, markers = build_sequence(tok, state, {"t": "noul", "ins": qd["instructions"], "crit": {}},
                                          cfg["max_len"], cfg["head_max_len"])
            if len(markers) != k:
                continue
            items.append({"ids": seq, "markers": markers, "qtype": QTYPES["noul"],
                          "target": [1.0 - float(y), float(y)], "q": q})
    return items


def collate(items, pad_id):
    n, L = len(items), max(len(it["ids"]) for it in items)
    kmax = max(len(it["markers"]) for it in items)
    ids = torch.full((n, L), pad_id, dtype=torch.long)
    att = torch.zeros((n, L), dtype=torch.long)
    mpos = torch.zeros((n, kmax), dtype=torch.long)
    mmask = torch.zeros((n, kmax), dtype=torch.bool)
    target = torch.zeros((n, kmax))
    for i, it in enumerate(items):
        ids[i, :len(it["ids"])] = torch.tensor(it["ids"])
        att[i, :len(it["ids"])] = 1
        k = len(it["markers"])
        mpos[i, :k] = torch.tensor(it["markers"])
        mmask[i, :k] = True
        target[i, :k] = torch.tensor(it["target"])
    return dict(input_ids=ids, attention_mask=att, marker_pos=mpos, marker_mask=mmask,
                target=target, qtype=torch.tensor([it["qtype"] for it in items]))


@torch.no_grad()
def logits_of(model, items, pad_id, device, bs=32):
    model.eval()
    out = []
    order = sorted(range(len(items)), key=lambda i: len(items[i]["ids"]))  # length-bucketed
    res = [None] * len(items)
    for b in range(0, len(order), bs):
        idx = order[b:b + bs]
        batch = collate([items[i] for i in idx], pad_id)
        with torch.autocast("cuda", dtype=torch.bfloat16):
            lg, _ = model(*(batch[k].to(device) for k in
                            ("input_ids", "attention_mask", "marker_pos", "marker_mask", "qtype")))
        lg = lg.float().cpu().numpy()
        for j, i in enumerate(idx):
            res[i] = lg[j, :len(items[i]["markers"])]
    return res


def p_true(logits, temp=1.0):
    z = np.stack(logits) / temp
    z = z - z.max(1, keepdims=True)
    e = np.exp(z)
    return e[:, 1] / e.sum(1)


def fit_temp(logits, targets):
    Z = torch.tensor(np.stack(logits))
    T = torch.tensor(np.array(targets), dtype=torch.float32)
    log_t = torch.zeros(1, requires_grad=True)
    opt = torch.optim.LBFGS([log_t], lr=0.1, max_iter=100)

    def closure():
        opt.zero_grad()
        loss = -(T * torch.log_softmax(Z / log_t.exp(), -1)).sum(-1).mean()
        loss.backward()
        return loss
    opt.step(closure)
    return float(torch.clamp(log_t.exp(), 0.1, 10.0))


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--variant", choices=["text", "full"], required=True)
    ap.add_argument("--epochs", type=int, default=3)
    ap.add_argument("--micro-batch", type=int, default=16)
    ap.add_argument("--grad-accum", type=int, default=4)
    ap.add_argument("--lr-encoder", type=float, default=2.5e-5)
    ap.add_argument("--lr-head", type=float, default=1e-4)
    ap.add_argument("--anchors-per-resource", type=int, default=3, help="full variant sampling")
    ap.add_argument("--max-len", type=int, default=384)
    ap.add_argument("--out", type=Path, default=None)
    ap.add_argument("--seed", type=int, default=20260923)
    args = ap.parse_args()
    out = args.out or Path(f"runs/laya-{args.variant}")
    random.seed(args.seed); np.random.seed(args.seed); torch.manual_seed(args.seed)
    device = torch.device("cuda")

    model_dir = snapshot_download(BASE, allow_patterns=["*.json", "*.safetensors", "tokenizer/*", "encoder/*"],
                                  ignore_patterns=["multilingual/*", "typed-decisions/*"])
    _fix_tokenizer_config(model_dir)
    cfg = json.loads(Path(model_dir, "rl_agent_config.json").read_text())
    cfg["max_len"] = args.max_len  # states are short; 384 leaves ~190 tokens for the listing
    tok = AutoTokenizer.from_pretrained(os.path.join(model_dir, "tokenizer"))

    df = pd.read_parquet("data/rows.parquet")
    s7 = SPLITS[7]
    train_side = df[df.host_bucket < TEST_HOST_BUCKET]
    tr = train_side[train_side.anchor_idx.isin(s7["train_anchors"]) & ~train_side.host_bucket.isin(CALIB_HOSTS)]
    ca = train_side[train_side.anchor_idx.isin(s7["train_anchors"]) & train_side.host_bucket.isin(CALIB_HOSTS)]
    # paid_14d labels exist for anchors 0..15; restrict so its window also ends before the test period
    for d in (tr, ca):
        d.loc[d.anchor_idx > max(SPLITS[14]["train_anchors"]), "paid_14d"] = -1

    if args.variant == "full":
        tr = tr.sample(frac=1.0, random_state=args.seed).groupby("resource").head(args.anchors_per_resource)
        ca = ca.sample(frac=1.0, random_state=0).groupby("resource").head(1)
    soft = args.variant == "text"
    t0 = time.time()
    train_items = make_items(tr, args.variant, tok, cfg, soft)
    calib_items = make_items(ca, args.variant, tok, cfg, soft=False if args.variant == "full" else True)
    print(f"items: train {len(train_items):,} calib {len(calib_items):,} "
          f"(built in {time.time()-t0:.0f}s; mean len {np.mean([len(i['ids']) for i in train_items]):.0f})")

    model = build_model(cfg, encoder_dir=os.path.join(model_dir, "encoder"))
    model.load_state_dict(load_file(os.path.join(model_dir, "model.safetensors")), strict=True)
    model.encoder.config.reference_compile = False  # same as laya.Agent: eager path, no torch.compile hang
    model.encoder.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False})
    model.to(device)

    enc = [p for n, p in model.named_parameters() if n.startswith("encoder.")]
    head = [p for n, p in model.named_parameters() if not n.startswith("encoder.")]
    opt = torch.optim.AdamW([{"params": enc, "lr": args.lr_encoder}, {"params": head, "lr": args.lr_head}],
                            weight_decay=0.01)
    mb, ga, G = args.micro_batch, args.grad_accum, 4
    total = max(1, (len(train_items) // (mb * ga)) * args.epochs)
    warm = max(1, total // 20)
    sched = torch.optim.lr_scheduler.LambdaLR(
        opt, lambda s: min(1.0, (s + 1) / warm) * 0.5 * (1 + math.cos(math.pi * min(1.0, s / total))))

    log = []
    best = (-1.0, None)
    for ep in range(args.epochs):
        model.train()
        sigma = 0.4 + (0.1 - 0.4) * (ep / max(1, args.epochs - 1))
        random.shuffle(train_items)
        # length-bucketed batches: shuffle, then sort within windows of 50 batches
        win = mb * 50
        batches = []
        for w in range(0, len(train_items), win):
            chunk = sorted(train_items[w:w + win], key=lambda it: len(it["ids"]))
            batches += [chunk[i:i + mb] for i in range(0, len(chunk), mb)]
        random.shuffle(batches)
        te = time.time(); run = 0.0
        opt.zero_grad(set_to_none=True)
        for bi, chunk in enumerate(batches):
            batch = collate(chunk, tok.pad_token_id)
            with torch.autocast("cuda", dtype=torch.bfloat16):
                logits, act = model(*(batch[k].to(device) for k in
                                      ("input_ids", "attention_mask", "marker_pos", "marker_mask", "qtype")))
            logits = logits.float()
            mask = batch["marker_mask"].to(device)
            k = mask.sum(-1, keepdim=True).float()
            target = batch["target"].to(device)
            qt = batch["qtype"].to(device)
            eps = torch.randn((G,) + logits.shape, device=device) * sigma * mask
            eps = (eps - eps.sum(-1, keepdim=True) / k) * mask
            z = logits.detach().unsqueeze(0) + eps
            q = torch.softmax(z.masked_fill(~mask, -1e4), -1)
            with torch.no_grad():
                r = proper_reward(q, target.unsqueeze(0), qt, mask, w_sph=0.75, w_rps=1.0)
                adv = (r - r.mean(0, keepdim=True)) / ((r - r.mean(0, keepdim=True)).std() + 1e-6)
            logp = -(((z - logits.unsqueeze(0)) ** 2) * mask).sum(-1) / (2 * sigma ** 2)
            loss_rl = -(adv * logp).mean()
            loss_ce = -(target * torch.log_softmax(logits.masked_fill(~mask, -1e4), -1)).sum(-1).mean()
            loss = (loss_rl + loss_ce) / ga + 0.0 * act.sum()
            loss.backward()
            run += loss_ce.item()
            if (bi + 1) % ga == 0 or bi == len(batches) - 1:
                torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
                opt.step(); sched.step(); opt.zero_grad(set_to_none=True)
            if (bi + 1) % 200 == 0:
                print(f"  ep {ep+1} batch {bi+1}/{len(batches)} ce {run/200:.4f} "
                      f"{(time.time()-te)/(bi+1)*1000:.0f}ms/batch mem {torch.cuda.max_memory_allocated()/1e9:.1f}GB",
                      flush=True)
                run = 0.0
        # model selection on held-out calib hosts (never trained on)
        cl = logits_of(model, calib_items, tok.pad_token_id, device)
        from sklearn.metrics import roc_auc_score
        idx7 = [i for i, it in enumerate(calib_items) if it["q"] == "paid_7d"]
        y7 = np.array([calib_items[i]["target"][1] for i in idx7])
        auc = roc_auc_score((y7 >= 0.5).astype(int), p_true([cl[i] for i in idx7]))
        log.append(dict(epoch=ep + 1, calib_auc_paid7d=round(float(auc), 4), secs=round(time.time() - te)))
        print(f"=== epoch {ep+1}: calib AUC paid_7d {auc:.4f}  ({time.time()-te:.0f}s)", flush=True)
        if auc > best[0]:
            best = (auc, {k: v.detach().to(torch.float16).cpu().clone() for k, v in model.state_dict().items()})

    model.load_state_dict({k: v.float() for k, v in best[1].items()})
    cl = logits_of(model, calib_items, tok.pad_token_id, device)
    temp = fit_temp(cl, [it["target"] for it in calib_items])
    print(f"best calib AUC {best[0]:.4f}; fitted noul temperature {temp:.3f}")

    out.mkdir(parents=True, exist_ok=True)
    save_file({k: v.contiguous() for k, v in best[1].items()}, str(out / "model.safetensors"))
    model.encoder.config.save_pretrained(out / "encoder")
    tok.save_pretrained(out / "tokenizer")
    cfg.update(fine_tuned=True, model_name=f"x402-demand-{args.variant}", base_model=BASE,
               temperature=[1.0, 1.0, temp], x402_variant=args.variant)
    cfg.pop("temperature_by_options", None)
    (out / "rl_agent_config.json").write_text(json.dumps(cfg, indent=2))
    (out / "train_log.json").write_text(json.dumps(dict(args={k: str(v) for k, v in vars(args).items()},
                                                        epochs=log, best_calib_auc=best[0], temperature=temp,
                                                        n_train_items=len(train_items),
                                                        n_calib_items=len(calib_items)), indent=1))
    print("saved", out)


if __name__ == "__main__":
    main()