"""Train the heads classifier on an export bundle. Run on the pod as ``python -m training.train``.

Classifier-plan §3.2-3.5, §4.3-4.4. Steps of ``--posts-per-step`` posts (32) in length buckets, split into
token-budget micro-batches (16k padded tokens) with per-step loss normalisation; masked targets; LoRA lr 2e-4,
heads lr 1e-3, AdamW wd 0.01, clip 1.0, cosine with 3% warmup; val every ``--eval-every`` steps (200, or four times
per epoch if that is fewer); ``best/`` = lowest mean val log-loss over all fields. The arm of the §4.4 screen is
``--checkpoint base|it`` (the base checkpoint dir is ``--model``) times ``--pool last|mean``. Per-option class
weights come from the bundle's taxonomy (nsb-v2: pnl breakeven 3.0) unless ``--class-weights`` overrides them; the
resolved weights are recorded in config.json.

Writes OUT/config.json, OUT/taxonomy.json, OUT/train_log.jsonl, OUT/last/ (resume point every eval: adapter,
heads.safetensors, optimizer, scheduler, RNG, data cursor, cumulative hours), OUT/best/, OUT/predictions/val.jsonl
and OUT/metrics_val.json for the best checkpoint. ``--max-hours`` is cumulative across resumes (state.json).
Exit code 3: stopped by --max-hours, --stop-after or SIGINT (``timeout``) at a saved resume point; rerun with
--resume. Exit code 4: non-finite loss or gradient (nothing saved; last/ keeps the previous good point).
"""
from __future__ import annotations

import argparse
import hashlib
import json
import math
import platform
import random
import time
from pathlib import Path

from .calibrate import nll, outputs
from .data import inference_batches, load_split, micro_batches, step_counts, steps, subset, write_jsonl
from .formats import FORMAT, encode
from .taxonomy import Taxonomy

# Settings a resume must not change.
CORE = ("data", "model", "checkpoint", "pool", "subset", "epochs", "max_steps", "posts_per_step", "token_budget",
        "max_len", "seed", "lr_lora", "lr_heads", "warmup", "weight_decay", "lora_r", "lora_alpha", "lora_dropout",
        "head_dropout", "dtype", "eval_every", "val_limit", "class_weights")


def parse(argv=None):
    p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    p.add_argument("--data", required=True, help="Bundle directory (train/val.jsonl, taxonomy.json, manifest.json)")
    p.add_argument("--model", required=True, help="Local base checkpoint directory")
    p.add_argument("--out", required=True)
    p.add_argument("--arm", default="", help="Arm name for the record (B-last, I-last, B-mean, I-mean)")
    p.add_argument("--checkpoint", default="base", choices=["base", "it"], help="Which Gemma 4 checkpoint --model is")
    p.add_argument("--pool", default="last", choices=["last", "mean"])
    p.add_argument("--subset", type=int, default=0, help="Train on this many rows in seeded order (0 = all)")
    p.add_argument("--epochs", type=float, default=2.0)
    p.add_argument("--max-steps", type=int, default=0, help="Cap the schedule at this many optimiser steps")
    p.add_argument("--posts-per-step", type=int, default=32)
    p.add_argument("--token-budget", type=int, default=16384, help="Padded tokens per training micro-batch")
    p.add_argument("--eval-token-budget", type=int, default=65536)
    p.add_argument("--max-len", type=int, default=2048)
    p.add_argument("--lr-lora", type=float, default=2e-4)
    p.add_argument("--lr-heads", type=float, default=1e-3)
    p.add_argument("--warmup", type=float, default=0.03)
    p.add_argument("--weight-decay", type=float, default=0.01)
    p.add_argument("--lora-r", type=int, default=32)
    p.add_argument("--lora-alpha", type=int, default=32)
    p.add_argument("--lora-dropout", type=float, default=0.05)
    p.add_argument("--head-dropout", type=float, default=0.1)
    p.add_argument("--class-weights", default=None,
                   help='Per-option loss weights as JSON, e.g. \'{"pnl": {"breakeven": 3.0}}\'; "{}" = none '
                        "(default: the bundle taxonomy's class_weights)")
    p.add_argument("--eval-every", type=int, default=200, help="Optimiser steps between val evaluations and saves")
    p.add_argument("--val-limit", type=int, default=0, help="Evaluate on the first N val rows in seeded order")
    p.add_argument("--log-every", type=int, default=10)
    p.add_argument("--max-hours", type=float, default=0, help="Cumulative training hours across resumes")
    p.add_argument("--stop-after", type=int, default=0, help="Stop this process after N more steps (resume test)")
    p.add_argument("--force-longest", type=int, default=0, help="First run one forward/backward on the N longest rows")
    p.add_argument("--seed", type=int, default=20261017)
    p.add_argument("--resume", action="store_true", help="Continue from last/ if it exists (settings must match)")
    p.add_argument("--device", default="cuda")
    p.add_argument("--dtype", default="bfloat16", choices=["bfloat16", "float32"])
    p.add_argument("--no-gradient-checkpointing", action="store_true")
    return p.parse_args(argv)


def versions() -> dict:
    out = dict(python=platform.python_version())
    for pkg in ("torch", "transformers", "peft", "accelerate", "safetensors", "tokenizers", "huggingface_hub"):
        try:
            out[pkg] = getattr(__import__(pkg), "__version__", "unknown")
        except Exception:
            out[pkg] = None
    try:
        import torch

        out["cuda"] = torch.version.cuda
        if torch.cuda.is_available():
            out["gpu"] = torch.cuda.get_device_name(0)
    except Exception:
        pass
    commit = Path(__file__).with_name(".source_commit")
    out["source_commit"] = commit.read_text().strip() if commit.exists() else None
    return out


def file_sha256(path: Path) -> str | None:
    if not path.exists():
        return None
    return hashlib.sha256(path.read_bytes()).hexdigest()


class Log:
    def __init__(self, path: Path):
        self.fh = open(path, "a", encoding="utf-8")

    def __call__(self, **row):
        row = dict(time=round(time.time(), 1), **row)
        line = json.dumps(row, ensure_ascii=False)
        self.fh.write(line + "\n")
        self.fh.flush()
        print(line, flush=True)


def prediction_rows(rows, logits_by_index, infos, split=None) -> list[dict]:
    """Prediction rows hold ids and numbers only (never text)."""
    return [{"id": r["id"], "split": split or r.get("split"), "n_tokens": infos[i]["n_tokens"],
             "truncated": infos[i]["cut"], "logits": logits_by_index[i]} for i, r in enumerate(rows)]


def val_report(preds: list[dict], rows: list[dict], tax: Taxonomy) -> dict:
    """Mean val log-loss over fields (selection), per-field log-loss and gated accuracy (uncalibrated)."""
    pairs = list(zip(preds, rows))
    per_field = {f: nll(pairs, tax, f, 1.0) for f in tax.names}
    scored = [v for v in per_field.values() if v is not None]
    acc = {}
    labels = [outputs(p["logits"], tax)["labels"] for p in preds]
    for f in tax.names:
        gold = [(lab[f], r["targets"].get(f)) for lab, r in zip(labels, rows)
                if f not in r.get("mask", ()) and r["targets"].get(f) is not None]
        acc[f] = sum(1 for a, g in gold if a == g) / len(gold) if gold else None
    eval_pairs = [(p, r) for p, r in pairs if r.get("eval")]
    screen = [nll(eval_pairs, tax, f, 1.0) for f in ("relevance", "direction", "commercial")]
    screen = [v for v in screen if v is not None]
    return {"selection_loss": sum(scored) / len(scored) if scored else float("inf"), "log_loss": per_field,
            "accuracy": acc, "rows": len(rows), "eval_rows": len(eval_pairs),
            "screen_primary": sum(screen) / len(screen) if screen else None}


def resolve_class_weights(tax: Taxonomy, override: str | None) -> Taxonomy:
    """The taxonomy with ``--class-weights`` applied (JSON object; None keeps the bundle's weights)."""
    if override is None:
        return tax
    value = json.loads(override)
    if not isinstance(value, dict):
        raise SystemExit("--class-weights must be a JSON object {field: {option: weight}}")
    return tax.replace(class_weights=value)


def schedule_eval_every(eval_every: int, steps_per_epoch: int) -> int:
    """200 steps, or four times per epoch if that is fewer steps (classifier-plan §4.3)."""
    return max(1, min(eval_every, math.ceil(steps_per_epoch / 4)))


def main(argv=None) -> int:
    args = parse(argv)
    import torch

    from .model import HUB_IDS, bos_id, build, load_tokenizer, pad, predict, rng_state, save_checkpoint, \
        set_rng_state, step_loss, tokenizer_fn

    out = Path(args.out)
    out.mkdir(parents=True, exist_ok=True)
    cfg_path = out / "config.json"
    tax = resolve_class_weights(Taxonomy.load(args.data), args.class_weights)
    cfg = {k: v for k, v in vars(args).items() if k not in ("resume", "device", "stop_after", "force_longest")}
    cfg["gradient_checkpointing"] = not args.no_gradient_checkpointing
    cfg["class_weights"] = tax.class_weights or None  # resolved; None (not {}) keeps older unweighted runs resumable
    if cfg_path.exists():
        if not args.resume:
            raise SystemExit(f"refusing to start a new run over {out}; use --resume or a new --out")
        old = json.loads(cfg_path.read_text())
        changed = sorted(k for k in CORE if old.get(k) != cfg.get(k))
        if changed:
            raise SystemExit(f"--resume with different settings: {changed}")
        if old.get("max_hours") != cfg.get("max_hours"):
            raise SystemExit("--resume with a different --max-hours; the cap is cumulative")
    manifest_path = Path(args.data) / "manifest.json"
    manifest = json.loads(manifest_path.read_text()) if manifest_path.exists() else {}
    cfg.update(versions=versions(), format=FORMAT, base_model=HUB_IDS[args.checkpoint],
               data_version=manifest.get("version"), data_manifest_sha256=file_sha256(manifest_path),
               taxonomy_version=manifest.get("taxonomy_version"), render_version=manifest.get("render_version"))
    cfg_path.write_text(json.dumps(cfg, indent=1))
    log = Log(out / "train_log.jsonl")

    random.seed(args.seed)
    torch.manual_seed(args.seed)
    device = torch.device(args.device)
    (out / "taxonomy.json").write_text(json.dumps(tax.to_json(), ensure_ascii=False, indent=1))
    tok = load_tokenizer(args.model)
    tokenize, bos = tokenizer_fn(tok), bos_id(tok)

    train_rows = subset(load_split(args.data, "train"), args.subset, args.seed)
    val_rows = subset(load_split(args.data, "val"), args.val_limit, args.seed)
    t0 = time.time()
    enc_train = [encode(r["segments"], tokenize, args.max_len, bos) for r in train_rows]
    enc_val = [encode(r["segments"], tokenize, args.max_len, bos) for r in val_rows]
    ids = [e[0] for e in enc_train]
    lengths = [len(x) for x in ids]
    labels = [tax.encode(r["targets"], r.get("mask", ())) for r in train_rows]
    val_ids = [e[0] for e in enc_val]
    val_infos = [e[1] for e in enc_val]
    val_batches = inference_batches([len(x) for x in val_ids], args.eval_token_budget)
    log(event="data", train=len(train_rows), val=len(val_rows), tokens=sum(lengths), max_len=max(lengths),
        mean_len=round(sum(lengths) / len(lengths), 1), truncated_rows=sum(1 for e in enc_train if e[1]["cut"]),
        val_tokens=sum(len(x) for x in val_ids), seconds=round(time.time() - t0, 1))

    last_dir, best_dir = out / "last", out / "best"
    resume = args.resume and (last_dir / "state.json").exists()
    model = build(cfg, tax.sizes(), device, train=True)
    if resume:
        from peft import set_peft_model_state_dict
        from safetensors.torch import load_file

        from .model import load_heads

        set_peft_model_state_dict(model.backbone, load_file(str(last_dir / "adapter" / "adapter_model.safetensors")))
        load_heads(model, last_dir / "heads.safetensors")
    lora = [p for n, p in model.named_parameters() if p.requires_grad and "lora_" in n]
    heads = [p for n, p in model.named_parameters() if p.requires_grad and "lora_" not in n]
    opt = torch.optim.AdamW([dict(params=lora, lr=args.lr_lora), dict(params=heads, lr=args.lr_heads)],
                            weight_decay=args.weight_decay)
    steps_per_epoch = math.ceil(len(train_rows) / args.posts_per_step)
    total = max(1, int(steps_per_epoch * args.epochs))
    if args.max_steps:
        total = min(total, args.max_steps)
    eval_every = schedule_eval_every(args.eval_every, steps_per_epoch)
    from transformers import get_cosine_schedule_with_warmup

    sched = get_cosine_schedule_with_warmup(opt, int(args.warmup * total), total)
    state = dict(step=0, epoch=0, step_in_epoch=0, best=None, best_step=None, hours_used=0.0)
    if resume:
        state = json.loads((last_dir / "state.json").read_text())
        trainer = torch.load(last_dir / "trainer.pt", map_location="cpu", weights_only=False)
        opt.load_state_dict(trainer["opt"])
        sched.load_state_dict(trainer["sched"])
        set_rng_state(trainer["rng"])
        log(event="resume", **state)
    log(event="start", trainable=sum(p.numel() for p in lora + heads), steps=total, steps_per_epoch=steps_per_epoch,
        eval_every=eval_every, arm=args.arm, checkpoint=args.checkpoint, pool=args.pool)

    amp = dict(device_type=device.type, dtype=torch.bfloat16, enabled=device.type == "cuda")
    process_start = time.time()
    hours_before = state["hours_used"]
    weights = tax.weights()
    class_weights = [None if v is None else torch.tensor(v, dtype=torch.float32, device=device)
                     for v in tax.class_weight_vectors()]

    def hours_now():
        return hours_before + (time.time() - process_start) / 3600

    def save_last():
        state["hours_used"] = round(hours_now(), 4)
        save_checkpoint(model, last_dir, tax.names, dict(state),
                        dict(opt=opt.state_dict(), sched=sched.state_dict(), rng=rng_state()))

    def evaluate():
        t = time.time()
        rng = rng_state()  # evaluation must not move the training RNG (dropout) stream
        logits = predict(model, val_ids, tok.pad_token_id, val_batches, device, tax.names)
        set_rng_state(rng)
        preds = prediction_rows(val_rows, logits, val_infos, "val")
        rep = val_report(preds, val_rows, tax)
        seconds = time.time() - t
        log(event="eval", step=state["step"], selection_loss=round(rep["selection_loss"], 5),
            screen_primary=rep["screen_primary"],
            log_loss={f: (round(v, 4) if v is not None else None) for f, v in rep["log_loss"].items()},
            accuracy={f: (round(v, 4) if v is not None else None) for f, v in rep["accuracy"].items()},
            seconds=round(seconds, 1), posts_per_second=round(len(val_rows) / max(seconds, 1e-9), 1))
        if state["best"] is None or rep["selection_loss"] < state["best"]:
            state["best"], state["best_step"] = rep["selection_loss"], state["step"]
            save_checkpoint(model, best_dir, tax.names, dict(state))
            write_jsonl(out / "predictions" / "val.jsonl", preds)
            (out / "metrics_val.json").write_text(json.dumps(dict(step=state["step"], **rep), indent=1))

    if args.force_longest and not resume:
        # Worst-case micro-batches for peak memory: the N longest rows of the whole train split (not the subset).
        pool = sorted(load_split(args.data, "train"), key=lambda r: -sum(len(x["text"]) for x in r["segments"]))
        cand = [(encode(r["segments"], tokenize, args.max_len, bos)[0], tax.encode(r["targets"], r.get("mask", ())))
                for r in pool[:4 * args.force_longest]]
        cand = sorted(cand, key=lambda c: -len(c[0]))[:args.force_longest]
        f_ids, f_labels = [c[0] for c in cand], [c[1] for c in cand]
        f_len = [len(x) for x in f_ids]
        parts = micro_batches(list(range(len(cand))), f_len, args.token_budget)
        for mb in parts:
            x, m, last = pad([f_ids[i] for i in mb], tok.pad_token_id, device)
            y = torch.tensor([f_labels[i] for i in mb], device=device)
            with torch.autocast(**amp):
                loss, _ = step_loss(model(x, m, last), y, step_counts(f_labels), weights, class_weights)
            loss.backward()
        opt.zero_grad(set_to_none=True)
        mem = round(torch.cuda.max_memory_allocated() / 2 ** 30, 2) if device.type == "cuda" else None
        log(event="forced_longest", rows=len(cand), lengths=f_len, micro_batches=len(parts),
            padded_tokens=[max(f_len[i] for i in mb) * len(mb) for mb in parts], max_memory_gb=mem)
        if device.type == "cuda":
            torch.cuda.reset_peak_memory_stats()

    model.train()
    try:
        stop = run_steps(args, state, total, eval_every, model, opt, sched, lora, heads, ids, lengths, labels, tax,
                         weights, device, amp, tok, log, hours_now, evaluate, save_last, class_weights)
    except KeyboardInterrupt:  # timeout sends SIGINT: keep a resume point at the last completed step
        log(event="interrupted", step=state["step"])
        save_last()
        return 3
    if stop == 4:
        return 4
    save_last()
    if state["best"] is None:
        evaluate()
    log(event="done" if stop is None else "stopped", **state)
    return stop or 0


def run_steps(args, state, total, eval_every, model, opt, sched, lora, heads, ids, lengths, labels, tax, weights,
              device, amp, tok, log, hours_now, evaluate, save_last, class_weights=None):
    """The optimiser loop. Returns None (finished), 3 (cap or planned stop; the caller saves the resume point)
    or 4 (non-finite loss or gradient: nothing is saved, so last/ keeps the previous good resume point)."""
    import torch

    from .model import pad, step_loss

    def fresh_window():
        return dict(loss=0.0, tokens=0, padded=0, t=time.time(), field_sums=[0.0] * len(tax.names),
                    field_counts=[0] * len(tax.names), n=0)

    window = fresh_window()
    stop, done_here = None, 0
    while state["step"] < total and stop is None:
        epoch_steps = steps(lengths, args.posts_per_step, args.seed * 1000 + state["epoch"])
        while state["step_in_epoch"] < len(epoch_steps) and state["step"] < total:
            if args.max_hours and hours_now() >= args.max_hours:
                stop = 3
                break
            step = epoch_steps[state["step_in_epoch"]]
            step_labels = [labels[i] for i in step]
            counts = step_counts(step_labels)
            step_total = 0.0
            for mb in micro_batches(step, lengths, args.token_budget):
                x, m, last = pad([ids[i] for i in mb], tok.pad_token_id, device)
                y = torch.tensor([labels[i] for i in mb], device=device)
                with torch.autocast(**amp):
                    loss, per_field = step_loss(model(x, m, last), y, counts, weights, class_weights)
                loss.backward()
                step_total += float(loss.detach())
                window["field_sums"] = [a + b for a, b in zip(window["field_sums"], per_field)]
                window["tokens"] += int(m.sum())
                window["padded"] += int(m.numel())
            window["field_counts"] = [a + sum(1 for r in step_labels if r[f] >= 0)
                                      for f, a in enumerate(window["field_counts"])]
            norm = float(torch.nn.utils.clip_grad_norm_(lora + heads, 1.0))
            if not (math.isfinite(step_total) and math.isfinite(norm)):
                log(event="non_finite", step=state["step"] + 1, loss=step_total, grad_norm=norm)
                opt.zero_grad(set_to_none=True)
                return 4
            opt.step()
            sched.step()
            opt.zero_grad(set_to_none=True)
            state["step"] += 1
            state["step_in_epoch"] += 1
            done_here += 1
            window["loss"] += step_total
            window["n"] += 1
            if state["step"] % args.log_every == 0 or state["step"] == total:
                dt = max(time.time() - window["t"], 1e-9)
                row = dict(event="train", step=state["step"], epoch=state["epoch"],
                           loss=round(window["loss"] / window["n"], 4),
                           fields={f: round(s / c, 4) if c else None for f, s, c in
                                   zip(tax.names, window["field_sums"], window["field_counts"])},
                           lr=sched.get_last_lr()[0], tokens_per_second=round(window["tokens"] / dt),
                           padded_tokens_per_second=round(window["padded"] / dt), hours=round(hours_now(), 3))
                if device.type == "cuda":
                    row["max_memory_gb"] = round(torch.cuda.max_memory_allocated() / 2 ** 30, 1)
                log(**row)
                window = fresh_window()
            if state["step_in_epoch"] >= len(epoch_steps):
                state["epoch"] += 1
                state["step_in_epoch"] = 0
            if state["step"] % eval_every == 0 or state["step"] == total:
                evaluate()
                save_last()
            if args.stop_after and done_here >= args.stop_after:
                stop = 3
                break
            if state["step_in_epoch"] == 0:
                break  # new epoch: new step order
    return stop


if __name__ == "__main__":
    raise SystemExit(main())