"""Single-GPU (RTX 3090) RLCD fine-tune of convaiinnovations/laya on typed-decisions.

Adapted from the official Kaggle 2xT4 DDP notebook (kept hyperparameters comparable):
  - no DDP; bf16 autocast instead of fp16+GradScaler (Ampere native)
  - MICRO_BATCH 16 x GRAD_ACCUM 4 -> same effective batch of 64 sequences
  - calibration slice (400 items) is held out of training before any batching
    (upstream issue #186 fix), and BOTH per-type and per-(type,option-count)
    bucket temperatures are fitted on that held-out slice only.

Input : train_items.pt (from preprocess.py) + convaiinnovations/laya in HF cache
Output: output/model/{model.safetensors, encoder/, tokenizer/, rl_agent_config.json}
Run:  D:\\venvs\\voxcpm-cuda\\Scripts\\python.exe  d:\\QoderCN\\jev_finetune\\train_single.py
"""
from __future__ import annotations

import json
import os
import random
import sys
import time

os.environ.setdefault("HF_ENDPOINT", "https://hf-mirror.com")
os.environ.setdefault("USE_TF", "0")
os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
sys.stdout.reconfigure(encoding="utf-8")

WORK = os.path.dirname(os.path.abspath(__file__))
ITEMS = os.path.join(WORK, "train_items.pt")
OUT_DIR = os.path.join(WORK, "output", "model")

# ---- hyperparameters (mirror the official notebook unless noted; overridable via CLI) ----
EPOCHS = 4
MICRO_BATCH = 16        # 3090 24GB (T4 version used 8 per GPU x 2 GPUs)
GRAD_ACCUM = 4          # effective batch = 16*4 = 64 sequences, same as notebook
GROUP_SIZE = 4          # GRPO-style noise samples per decision
LR_ENCODER = 2.5e-5
LR_HEAD = 1.0e-4
SIGMA_START = 0.4
SIGMA_END = 0.1
W_SPH = 0.75            # spherical-score weight in the proper reward (soft-distribution emphasis)
CALIB_MAX = 400
CALIB_SEED = 20260922   # same fixed seed as the official notebook


def collate_train_batch(items, pad_id):
    import torch

    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), dtype=torch.float32)
    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, : len(it["target"])] = torch.tensor(it["target"], dtype=torch.float32)
    return {
        "input_ids": ids,
        "attention_mask": att,
        "marker_pos": mpos,
        "marker_mask": mmask,
        "target": target,
        "qtype": torch.tensor([it["qtype"] for it in items]),
        "label": torch.tensor([it["label"] for it in items]),
    }


def fit_one_temp(sel):
    """Temperature scaling by NLL on held-out (logits, soft-target) pairs."""
    import torch

    if len(sel) < 10:
        return 1.0
    kmax = max(len(z) for z, _ in sel)
    Z = torch.full((len(sel), kmax), -1e4)
    T = torch.zeros((len(sel), kmax))
    for i, (z, t) in enumerate(sel):
        Z[i, : len(z)] = torch.tensor(z)
        T[i, : len(t)] = torch.tensor(t, 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).item())


def temp_bucket(qtype_name: str, k: int) -> str:
    """Decision-Jef style per-cardinality buckets (same key format laya runtime reads)."""
    size = "2" if k <= 2 else "3-5" if k <= 5 else "6-10" if k <= 10 else "11+"
    return "%s:%s" % (qtype_name, size)


def main():
    global EPOCHS, GROUP_SIZE, SIGMA_END, W_SPH, LR_ENCODER, LR_HEAD, OUT_DIR
    import argparse

    ap = argparse.ArgumentParser(description="single-GPU RLCD fine-tune of laya on typed-decisions")
    ap.add_argument("--epochs", type=int, default=EPOCHS)
    ap.add_argument("--group-size", type=int, default=GROUP_SIZE)
    ap.add_argument("--w-sph", type=float, default=W_SPH, help="spherical reward weight (soft-distribution emphasis)")
    ap.add_argument("--sigma-end", type=float, default=SIGMA_END)
    ap.add_argument("--lr-scale", type=float, default=1.0, help="scale both LRs (<1 to reduce overfit on longer runs)")
    ap.add_argument("--tag", default="", help="variant name; empty = default output/model")
    ap.add_argument("--items", default=ITEMS, help="items .pt: list, or dict with 'train'/'calib' (pre-split)")
    a = ap.parse_args()
    EPOCHS, GROUP_SIZE, W_SPH, SIGMA_END = a.epochs, a.group_size, a.w_sph, a.sigma_end
    LR_ENCODER *= a.lr_scale
    LR_HEAD *= a.lr_scale
    if a.tag:
        OUT_DIR = os.path.join(WORK, "output", f"model_{a.tag}")
    print(f"[cfg] epochs={EPOCHS} group={GROUP_SIZE} w_sph={W_SPH} sigma_end={SIGMA_END} "
          f"lr_enc={LR_ENCODER:.2e} lr_head={LR_HEAD:.2e} -> {OUT_DIR}")

    import torch
    from safetensors.torch import load_file, save_file
    from transformers import AutoTokenizer
    from huggingface_hub import snapshot_download
    from laya.common import build_model, proper_reward, QTYPES, QTYPE_NAMES

    assert torch.cuda.is_available(), "CUDA torch required: use D:\\venvs\\voxcpm-cuda"
    device = torch.device("cuda", 0)
    p = torch.cuda.get_device_properties(0)
    print(f"[env] GPU: {p.name} ({p.total_memory/1e9:.1f} GB) | torch {torch.__version__}")

    print("[1/5] load base checkpoint (convaiinnovations/laya root) ...")
    model_dir = snapshot_download("convaiinnovations/laya")
    with open(os.path.join(model_dir, "rl_agent_config.json")) as f:
        cfg = json.load(f)
    cfg["gradient_checkpointing"] = True
    cfg["max_len"] = 1024
    cfg["head_max_len"] = 256

    tok = AutoTokenizer.from_pretrained(os.path.join(model_dir, "tokenizer"))
    model = build_model(cfg, encoder_dir=os.path.join(model_dir, "encoder"))
    weights = load_file(os.path.join(model_dir, "model.safetensors"))
    model.load_state_dict(weights, strict=True)
    model.encoder.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False})
    model.head_checkpointing = True
    model.to(device)
    model.train()
    n_params = sum(x.numel() for x in model.parameters())
    print(f"      params: {n_params/1e6:.1f}M (full fine-tune, bf16 autocast)")

    print("[2/5] items + calibration holdout ...")
    loaded = torch.load(a.items, weights_only=False)
    if isinstance(loaded, dict):  # pre-split by preprocess_aug.py (calib slice is clean)
        train_items = loaded["train"]
        calib_items = loaded["calib"]
    else:
        all_items = loaded
        order = list(range(len(all_items)))
        random.Random(CALIB_SEED).shuffle(order)
        n_calib = min(CALIB_MAX, len(all_items) // 10)
        calib_items = [all_items[i] for i in sorted(order[:n_calib])]
        train_items = [all_items[i] for i in sorted(order[n_calib:])]
    print(f"      train={len(train_items)} | calib held out={len(calib_items)} (never enters a training batch)")

    enc_params = [p_ for nm, p_ in model.named_parameters() if "encoder." in nm]
    head_params = [p_ for nm, p_ in model.named_parameters() if "encoder." not in nm]
    assert enc_params and head_params, "param split failed: encoder/head groups both must be non-empty"
    optimizer = torch.optim.AdamW(
        [{"params": enc_params, "lr": LR_ENCODER}, {"params": head_params, "lr": LR_HEAD}],
        weight_decay=0.01)
    total_updates = (len(train_items) // (MICRO_BATCH * GRAD_ACCUM)) * EPOCHS
    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=max(1, total_updates), eta_min=1e-6)
    print(f"      ~{total_updates} updates | micro={MICRO_BATCH} accum={GRAD_ACCUM} epochs={EPOCHS}")

    pad_id = tok.pad_token_id
    print("[3/5] training ...")
    t_all = time.time()
    for epoch in range(EPOCHS):
        random.seed(42 + epoch)
        random.shuffle(train_items)
        epoch_loss, n_batches = 0.0, 0
        optimizer.zero_grad(set_to_none=True)
        accum_step = 0
        progress = epoch / max(1, EPOCHS - 1)
        sigma = SIGMA_START + (SIGMA_END - SIGMA_START) * progress
        t0 = time.time()
        for b_idx in range(0, len(train_items), MICRO_BATCH):
            chunk = train_items[b_idx:b_idx + MICRO_BATCH]
            if len(chunk) < MICRO_BATCH:
                continue
            batch = collate_train_batch(chunk, pad_id)
            batch = {k: (v.to(device) if torch.is_tensor(v) else v) for k, v in batch.items()}

            with torch.autocast("cuda", dtype=torch.bfloat16):
                logits, act = model(batch["input_ids"], batch["attention_mask"],
                                     batch["marker_pos"], batch["marker_mask"], batch["qtype"])
            logits = logits.float()
            mask, target, qtype = batch["marker_mask"], batch["target"], batch["qtype"]
            k = mask.sum(-1, keepdim=True).float()

            # RLCD: zero-mean gaussian exploration on detached logits
            eps = torch.randn((GROUP_SIZE,) + 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), qtype, mask, w_sph=W_SPH, w_rps=1.0)
                adv = r - r.mean(0, keepdim=True)
                adv = adv / (adv.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 + 1.0 * loss_ce) / GRAD_ACCUM + 0.0 * act.sum()
            loss.backward()
            accum_step += 1
            if accum_step % GRAD_ACCUM == 0:
                torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
                optimizer.step()
                scheduler.step()
                optimizer.zero_grad(set_to_none=True)
            epoch_loss += loss.item() * GRAD_ACCUM
            n_batches += 1
            if n_batches % 50 == 0:
                el = time.time() - t0
                eta = el / n_batches * (len(train_items) // MICRO_BATCH - n_batches)
                mem = torch.cuda.max_memory_allocated() / 1e9
                print(f"  ep{epoch+1} step {n_batches} | loss {loss.item()*GRAD_ACCUM:.4f} | reward {r.mean().item():.3f} | "
                      f"lr {scheduler.get_last_lr()[0]:.2e} | {mem:.1f}GB | eta {eta/60:.1f}min")
        print(f"=== epoch {epoch+1}/{EPOCHS} in {time.time()-t0:.0f}s | avg loss {epoch_loss/max(1,n_batches):.4f} ===")

        os.makedirs(OUT_DIR, exist_ok=True)
        ckpt_sd = {kk: v.half().contiguous().cpu() for kk, v in model.state_dict().items()}
        save_file(ckpt_sd, os.path.join(OUT_DIR, "checkpoint_latest.safetensors"))
        with open(os.path.join(OUT_DIR, "checkpoint_meta.json"), "w") as f:
            json.dump({"epoch": epoch + 1, "total_epochs": EPOCHS,
                       "avg_loss": epoch_loss / max(1, n_batches)}, f, indent=2)

    print(f"[4/5] total train time {(time.time()-t_all)/60:.1f} min")
    print("[5/5] fitting calibration temperatures on the HELD-OUT slice ...")
    del optimizer, scheduler
    torch.cuda.empty_cache()
    model.eval()
    preds = []
    with torch.no_grad():
        for c_idx in range(0, len(calib_items), 16):
            c_chunk = calib_items[c_idx:c_idx + 16]
            cb = collate_train_batch(c_chunk, pad_id)
            cb = {kk: (v.to(device) if torch.is_tensor(v) else v) for kk, v in cb.items()}
            with torch.autocast("cuda", dtype=torch.bfloat16):
                l_sub, _ = model(cb["input_ids"], cb["attention_mask"], cb["marker_pos"],
                                  cb["marker_mask"], cb["qtype"])
            l_np = l_sub.float().cpu().numpy()
            for r_i, it in enumerate(c_chunk):
                kk = len(it["markers"])
                preds.append((it["qtype"], l_np[r_i, :kk], it["target"]))

    qtnames = {0: "choice", 1: "score", 2: "noul"}
    temps = [1.2, 1.2, 1.2]
    for qt in range(3):
        sel = [(z, t) for q_type, z, t in preds if q_type == qt]
        if sel:
            temps[qt] = fit_one_temp(sel)
    buckets = {}
    for bkey in sorted({temp_bucket(qtnames[qt], len(z)) for qt, z, _ in preds}):
        tname, size = bkey.split(":")
        lo, hi = {"2": (0, 2), "3-5": (3, 5), "6-10": (6, 10), "11+": (11, 10**9)}[size]
        sel = [(z, t) for q_type, z, t in preds
               if qtnames[q_type] == tname and lo <= len(z) <= hi]
        if len(sel) >= 30:  # below that the bucket fit is noise; runtime falls back to per-type
            buckets[bkey] = fit_one_temp(sel)
    print("  per-type temps (choice, score, noul):", [round(t, 3) for t in temps])
    print("  per-bucket temps:", {kk: round(v, 3) for kk, v in buckets.items()})

    sd = {kk: v.half().contiguous().cpu() for kk, v in model.state_dict().items()}
    save_file(sd, os.path.join(OUT_DIR, "model.safetensors"))
    model.encoder.config.save_pretrained(os.path.join(OUT_DIR, "encoder"))
    tok.save_pretrained(os.path.join(OUT_DIR, "tokenizer"))
    cfg["fine_tuned"] = True
    cfg["model_name"] = "laya-td-3090"
    cfg["temperature"] = temps
    cfg["temperature_by_options"] = buckets  # fitted on held-out data, unlike upstream #186
    cfg["training"] = {"updates": n_batches // GRAD_ACCUM * EPOCHS, "epochs_completed": EPOCHS,
                       "hours": round((time.time() - t_all) / 3600, 2), "world_size": 1,
                       "fine_tuned_from_checkpoint": True, "calib_held_out": len(calib_items)}
    with open(os.path.join(OUT_DIR, "rl_agent_config.json"), "w") as f:
        json.dump(cfg, f, indent=2)
    print(f"[DONE] model saved -> {OUT_DIR}")


if __name__ == "__main__":
    main()