"""Train a DecisionModel.

    python -m jevlike.train --train data/train.jsonl --val data/val.jsonl --out checkpoints/jevlike-base

Writes the best checkpoint (by --select on val) to --out, a resumable ``last/`` state (model +
optimizer, overwritten at every eval), ``train_args.json`` and ``log.jsonl``.

Context length: ``--max-len`` (default 512, or the --init-from / --resume checkpoint's value) is
used to serialize the data and is stored in the saved config (``--save-max-len`` stores a
different one, e.g. train at 2048 but serve at 4096). Continue a 512 checkpoint at 4096 with
``--init-from checkpoints/jevlike-v2 --max-len 4096``. ``--max-tokens`` is raised to --max-len if
smaller so every micro-batch fits one full-length sequence.
"""
from __future__ import annotations

import os

os.environ.setdefault("PYTORCH_ALLOC_CONF", "expandable_segments:True")
os.environ.setdefault("TOKENIZERS_PARALLELISM", "true")

import argparse
import json
import math
import random
import shutil
import time
from collections import defaultdict

import numpy as np
import torch
from torch.utils.data import DataLoader
from transformers import AutoTokenizer

from jevlike.model import DecisionConfig, DecisionModel, compute_loss, ece
from jevlike.serialize import EncodedDataset, Serializer, TokenBudgetBatchSampler
from jevlike.types import QTYPES, read_jsonl


def to_device(b: dict[str, torch.Tensor], device) -> dict:
    out = {k: v.to(device, non_blocking=True) for k, v in b.items()}
    out["k_max"] = int(b["n_options"].max())
    return out


def autocast(device: torch.device):
    return torch.autocast(device.type, dtype=torch.bfloat16, enabled=device.type == "cuda")


# ---------------------------------------------------------------- evaluation (shared with calibrate.py)

@torch.no_grad()  # not inference_mode: calibrate.py fits on these tensors with autograd
def collect(model: DecisionModel, ds: EncodedDataset, device, max_tokens: int = 16384) -> dict:
    """Raw per-question outputs over a dataset, in dataset order."""
    model.eval()
    sampler = TokenBudgetBatchSampler(ds.lengths, max_tokens=max_tokens, max_batch=256, shuffle=False)
    logits: list = [None] * len(ds)
    esc = torch.zeros(len(ds))
    for idx in sampler:
        b = to_device(ds.collate(idx), device)
        with autocast(device):
            out = model(**b)
        lg, e = out["logits"].float().cpu(), out["escalate_logit"].float().cpu()
        for j, i in enumerate(idx):
            logits[i] = lg[j, :len(ds.encs[i].markers)]
        esc[idx] = e
    return {"logits": logits, "escalate_logit": esc, "labels": torch.tensor(ds.labels),
            "qtype": torch.tensor([e.qtype for e in ds.encs]),
            "n_options": torch.tensor([len(e.markers) for e in ds.encs]),
            "source": [r.source for r in ds.records]}


def summarize(c: dict, temps: list[float] | None = None, by_source: bool = False) -> dict:
    """NLL / accuracy / ECE overall and per type (+ RPS for score, escalate BCE)."""
    lp = [(l / (temps[i] if temps else 1.0)).log_softmax(-1) for i, l in enumerate(c["logits"])]
    labels = c["labels"]
    nll = torch.stack([-l[y] for l, y in zip(lp, labels.tolist())])
    conf = torch.stack([l.exp().max() for l in lp])
    correct = torch.tensor([int(l.argmax()) == y for l, y in zip(lp, labels.tolist())]).float()
    esc_bce = torch.nn.functional.binary_cross_entropy_with_logits(c["escalate_logit"], correct, reduction="none")
    res = {"n": len(labels), "nll": nll.mean().item(), "acc": correct.mean().item(), "ece": ece(conf, correct),
           "esc_bce": esc_bce.mean().item()}
    for t, name in enumerate(QTYPES):
        m = c["qtype"] == t
        if m.any():
            res[f"{name}/n"] = int(m.sum())
            res[f"{name}/acc"] = correct[m].mean().item()
            res[f"{name}/nll"] = nll[m].mean().item()
            res[f"{name}/ece"] = ece(conf[m], correct[m])
    sc = (c["qtype"] == QTYPES.index("score")).nonzero().flatten().tolist()
    if sc:
        rps = [((lp[i].exp().cumsum(0) - (torch.arange(len(lp[i])) >= labels[i]).float()) ** 2).sum().item()
               / (len(lp[i]) - 1) for i in sc]
        res["score/rps"] = float(np.mean(rps))
    types = [t for t in QTYPES if f"{t}/acc" in res]
    res["macro_acc"] = float(np.mean([res[f"{t}/acc"] for t in types]))
    if by_source:
        src = defaultdict(list)
        for s, ok in zip(c["source"], correct.tolist()):
            src[s].append(ok)
        res["source_acc"] = {s: round(float(np.mean(v)), 4) for s, v in sorted(src.items())}
    return res


# ---------------------------------------------------------------- training

def lr_lambda(warmup: int, total: int, schedule: str):
    def f(step: int) -> float:
        if step < warmup:
            return (step + 1) / warmup
        x = min(1.0, (step - warmup) / max(1, total - warmup))
        return 0.5 * (1 + math.cos(math.pi * x)) if schedule == "cosine" else 1 - x
    return f


def param_groups(model: DecisionModel, a) -> list[dict]:
    groups = {(bb, dec): [] for bb in (True, False) for dec in (True, False)}
    for n, p in model.named_parameters():
        if p.requires_grad:
            groups[(n.startswith("backbone."), p.ndim >= 2 and "embeddings" not in n)].append(p)
    return [{"params": ps, "lr": a.lr_backbone if bb else a.lr_head, "weight_decay": a.weight_decay if dec else 0.0}
            for (bb, dec), ps in groups.items() if ps]


def main(argv: list[str] | None = None) -> None:
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--train", required=True)
    ap.add_argument("--val", required=True)
    ap.add_argument("--out", required=True)
    ap.add_argument("--backbone", default="answerdotai/ModernBERT-base")
    ap.add_argument("--max-len", type=int, default=None,
                    help="training context length (default: the --init-from/--resume checkpoint's, else 512)")
    ap.add_argument("--head-max-len", type=int, default=None, help="question-part budget (default: checkpoint's, else 192)")
    ap.add_argument("--save-max-len", type=int, default=None,
                    help="max_len stored in the saved config (default: --max-len)")
    ap.add_argument("--device", default=None, help="cpu / cuda (default: cuda if available)")
    ap.add_argument("--epochs", type=float, default=2.0)
    ap.add_argument("--max-steps", type=int, default=0, help="optimizer steps; overrides --epochs if > 0")
    ap.add_argument("--max-tokens", type=int, default=2048, help="padded tokens per micro-batch")
    ap.add_argument("--max-batch", type=int, default=64, help="questions per micro-batch")
    ap.add_argument("--grad-accum", type=int, default=8)
    ap.add_argument("--lr-backbone", type=float, default=3e-5)
    ap.add_argument("--lr-head", type=float, default=3e-4)
    ap.add_argument("--weight-decay", type=float, default=0.01)
    ap.add_argument("--warmup", type=float, default=0.06, help="fraction of total steps")
    ap.add_argument("--schedule", choices=["linear", "cosine"], default="linear")
    ap.add_argument("--clip", type=float, default=1.0)
    ap.add_argument("--dropout", type=float, default=0.1)
    ap.add_argument("--rps-weight", type=float, default=1.0)
    ap.add_argument("--esc-weight", type=float, default=0.2)
    ap.add_argument("--grad-checkpointing", action="store_true")
    ap.add_argument("--train-embeddings", action="store_true", help="default: token embeddings frozen")
    ap.add_argument("--eval-every", type=int, default=1000)
    ap.add_argument("--log-every", type=int, default=20)
    ap.add_argument("--select", choices=["nll", "acc"], default="nll", help="best-checkpoint metric on val")
    ap.add_argument("--limit-train", type=int, default=0)
    ap.add_argument("--limit-val", type=int, default=0)
    ap.add_argument("--seed", type=int, default=0)
    ap.add_argument("--resume", action="store_true", help="continue from OUT/last if present")
    ap.add_argument("--init-from", default=None, help="start from a trained checkpoint dir (fresh optimizer)")
    ap.add_argument("--optim", choices=["adamw", "adamw8bit"], default="adamw")
    ap.add_argument("--freeze-layers", type=int, default=0, help="freeze the bottom N backbone encoder layers")
    a = ap.parse_args(argv)

    random.seed(a.seed), np.random.seed(a.seed), torch.manual_seed(a.seed)
    device = torch.device(a.device or ("cuda" if torch.cuda.is_available() else "cpu"))
    os.makedirs(a.out, exist_ok=True)
    last_dir = os.path.join(a.out, "last")
    resume = a.resume and os.path.exists(os.path.join(last_dir, "trainer.pt"))
    src_cfg = DecisionConfig.load(last_dir if resume else a.init_from) if resume or a.init_from else None
    a.max_len = a.max_len or (src_cfg.max_len if src_cfg else 512)
    a.head_max_len = a.head_max_len or (src_cfg.head_max_len if src_cfg else 192)
    limit = int((src_cfg.backbone_config if src_cfg else {}).get("max_position_embeddings") or 8192)
    for name in ("max_len", "save_max_len"):
        v = getattr(a, name)
        if v is not None and not a.head_max_len < v <= limit:
            ap.error(f"--{name.replace('_', '-')} {v} must be in ({a.head_max_len}, {limit}]")
    if a.max_tokens < a.max_len:
        print(f"--max-tokens {a.max_tokens} < --max-len {a.max_len}: raising --max-tokens to {a.max_len} "
              f"(a micro-batch must fit one full-length sequence)")
        a.max_tokens = a.max_len
    log_f = open(os.path.join(a.out, "log.jsonl"), "a")

    def log(d: dict, show: str | None = None) -> None:
        log_f.write(json.dumps(d) + "\n")
        log_f.flush()
        print(show or json.dumps(d), flush=True)

    tok = AutoTokenizer.from_pretrained(last_dir if resume else a.init_from or a.backbone)
    ser = Serializer(tok, a.max_len, a.head_max_len)
    t0 = time.time()
    rng = random.Random(a.seed)
    tr_recs, va_recs = read_jsonl(a.train), read_jsonl(a.val)
    if a.limit_train:
        tr_recs = rng.sample(tr_recs, min(a.limit_train, len(tr_recs)))
    if a.limit_val:
        va_recs = rng.sample(va_recs, min(a.limit_val, len(va_recs)))
    tr, va = EncodedDataset(tr_recs, ser), EncodedDataset(va_recs, ser)
    trunc = np.bincount([e.trunc for e in tr.encs], minlength=6).tolist()
    print(f"encoded {len(tr)} train / {len(va)} val in {time.time() - t0:.0f}s; train tokens mean "
          f"{tr.lengths.mean():.0f} max {tr.lengths.max()}; trunc codes [none,state,instr,texts,keys,overflow]={trunc}")

    sampler = TokenBudgetBatchSampler(tr.lengths, a.max_tokens, a.max_batch, shuffle=True, seed=a.seed)
    loader = DataLoader(tr, batch_sampler=sampler, collate_fn=tr.collate, pin_memory=device.type == "cuda")
    n_micro = len(sampler)
    steps_per_epoch = math.ceil(n_micro / a.grad_accum)
    total = a.max_steps or math.ceil(steps_per_epoch * a.epochs)

    if resume:
        model = DecisionModel.from_pretrained(last_dir, device)
    elif a.init_from:
        model = DecisionModel.from_pretrained(a.init_from, device)
        model.cfg.name = os.path.basename(os.path.normpath(a.out))
    else:
        cfg = DecisionConfig(backbone=a.backbone, dropout=a.dropout, max_len=a.max_len, head_max_len=a.head_max_len,
                             name=os.path.basename(os.path.normpath(a.out)))
        model = DecisionModel.from_backbone(cfg).to(device)
    if not resume:  # the context length the checkpoint is saved (and served) with
        model.cfg.max_len, model.cfg.head_max_len = a.save_max_len or a.max_len, a.head_max_len
    if not a.train_embeddings:
        model.backbone.get_input_embeddings().weight.requires_grad_(False)
    for layer in model.backbone.layers[:a.freeze_layers]:  # bottom encoder layers (memory for -large)
        layer.requires_grad_(False)
    if a.grad_checkpointing:
        model.backbone.gradient_checkpointing_enable({"use_reentrant": False})
    if a.optim == "adamw8bit":  # 8-bit optimizer states: needed to fit ModernBERT-large in 6 GB
        import bitsandbytes as bnb
        opt = bnb.optim.AdamW8bit(param_groups(model, a), betas=(0.9, 0.98), eps=1e-6)
    else:
        opt = torch.optim.AdamW(param_groups(model, a), betas=(0.9, 0.98), eps=1e-6, fused=device.type == "cuda")
    sched = torch.optim.lr_scheduler.LambdaLR(opt, lr_lambda(int(a.warmup * total), total, a.schedule))
    step, epoch, consumed, best = 0, 0, 0, None
    if resume:
        st = torch.load(os.path.join(last_dir, "trainer.pt"), map_location=device, weights_only=False)
        opt.load_state_dict(st["optimizer"]), sched.load_state_dict(st["scheduler"])
        step, epoch, consumed, best = st["step"], st["epoch"], st["consumed"], st["best"]
        torch.set_rng_state(st["rng"])
        print(f"resumed at step {step} (epoch {epoch}, batch {consumed})")
    else:
        with open(os.path.join(a.out, "train_args.json"), "w") as f:
            json.dump(vars(a), f, indent=1)
    n_params = sum(p.numel() for p in model.parameters()) / 1e6
    n_train = sum(p.numel() for p in model.parameters() if p.requires_grad) / 1e6
    print(f"model {n_params:.1f}M params ({n_train:.1f}M trainable); {n_micro} micro-batches/epoch, "
          f"{steps_per_epoch} steps/epoch, {total} total steps, device {device}")

    def evaluate() -> None:
        nonlocal best
        res = summarize(collect(model, va, device), by_source=True)
        model.train()
        key = res["nll"] if a.select == "nll" else -res["macro_acc"]
        improved = best is None or key < best
        log({"kind": "val", "step": step, **res},
            f"[val step {step}] nll {res['nll']:.4f} acc {res['acc']:.4f} ece {res['ece']:.4f} " +
            " ".join(f"{t}:{res[f'{t}/acc']:.3f}/{res[f'{t}/nll']:.3f}" for t in QTYPES if f"{t}/acc" in res) +
            (" *best*" if improved else ""))
        if improved:
            best = key
            model.save_pretrained(a.out, tok)
        tmp = last_dir + ".tmp"
        model.save_pretrained(tmp, tok)
        torch.save({"optimizer": opt.state_dict(), "scheduler": sched.state_dict(), "step": step, "epoch": epoch,
                    "consumed": consumed, "best": best, "rng": torch.get_rng_state()}, os.path.join(tmp, "trainer.pt"))
        shutil.rmtree(last_dir, ignore_errors=True)
        os.replace(tmp, last_dir)

    model.train()
    if device.type == "cuda":
        torch.cuda.reset_peak_memory_stats()
    win = defaultdict(float)
    t_win = time.time()
    while step < total:
        sampler.set_epoch(epoch, skip=consumed)
        it = iter(loader)
        while micro := [b for _, b in zip(range(a.grad_accum), it)]:
            n_q = sum(len(b["labels"]) for b in micro)
            for b in micro:
                b = to_device(b, device)
                with autocast(device):
                    out = model(**b)
                loss, stats = compute_loss(out, b, a.rps_weight, a.esc_weight)
                (loss / n_q).backward()
                for k, v in stats.items():
                    win[k] += v
                win["loss"] += loss.item()
                win["tokens"] += b["attention_mask"].sum().item()
            consumed += len(micro)
            gnorm = torch.nn.utils.clip_grad_norm_(model.parameters(), a.clip)
            opt.step()
            sched.step()
            opt.zero_grad(set_to_none=True)
            step += 1
            if step % a.log_every == 0 or step == total:
                dt = time.time() - t_win
                n = win["n"]
                mem = torch.cuda.max_memory_allocated() / 2**30 if device.type == "cuda" else 0.0
                d = {"kind": "train", "step": step, "epoch": round(epoch + consumed / n_micro, 3),
                     "loss": win["loss"] / n, "nll": win["nll"] / n, "acc": win["correct"] / n,
                     "rps": win["rps"] / max(1, win["n_score"]), "esc_bce": win["bce"] / n,
                     "lr": sched.get_last_lr()[0], "gnorm": float(gnorm), "ex_s": n / dt,
                     "tok_s": win["tokens"] / dt, "peak_gb": round(mem, 3)}
                log(d, f"step {step}/{total} ep {d['epoch']:.2f} loss {d['loss']:.4f} nll {d['nll']:.4f} "
                       f"acc {d['acc']:.3f} gnorm {d['gnorm']:.2f} lr {d['lr']:.2e} {d['ex_s']:.1f} ex/s "
                       f"{d['tok_s']:.0f} tok/s peak {mem:.2f}GB")
                win.clear()
                t_win = time.time()
            if step % a.eval_every == 0 or step == total:
                evaluate()
                t_win = time.time()
            if step >= total:
                break
        else:  # epoch exhausted
            epoch, consumed = epoch + 1, 0
    log_f.close()
    print(f"done: best val {a.select} {abs(best):.4f}; checkpoint in {a.out}")


if __name__ == "__main__":
    main()