#!/usr/bin/env python
"""Stage A: backbone pretraining (plain causal decoder; workspace disabled) on a uint16 token stream.

Usage: python -m training.pretrain --config configs/experiments/pretrain_103m.yaml [--dry-run] [--device mps]
The experiment YAML has keys: model (path to model yaml or inline dict), data (train_bin, val_bin), train
(batch_size, seq_len, steps, lr, warmup, weight_decay, grad_clip, eval_every, eval_batches, ckpt_every), seed.
"""
from __future__ import annotations

import argparse, json, math, sys, time
from pathlib import Path

import torch

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from prisme.config import PrismeConfig, load_config  # noqa: E402
from prisme.model import PrismeModel  # noqa: E402
from prisme.losses import lm_loss  # noqa: E402
from training.common import (Run, TokenBin, cosine_lr, load_checkpoint, load_yaml, make_adamw, peak_memory_bytes,  # noqa: E402
                             pick_device, save_checkpoint, seed_all, sha256_file, sync)


def build_model_cfg(exp: dict) -> PrismeConfig:
    m = exp["model"]
    cfg = load_config(m) if isinstance(m, str) else PrismeConfig.from_dict(m)
    for k, v in exp.get("model_overrides", {}).items():
        setattr(cfg, k, v)
    cfg.validate()
    return cfg


@torch.no_grad()
def evaluate(model, bin_: TokenBin, batch_size, device, max_batches):
    model.eval()
    tot, n = 0.0, 0
    for x, y in bin_.sequential_batches(batch_size, device, max_batches):
        logits = model(x, None)["logits"]
        tot += lm_loss(logits, y).item() * y.numel(); n += y.numel()
    model.train()
    return tot / max(n, 1)


def main(argv=None) -> int:
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("--config", required=True)
    ap.add_argument("--dry-run", action="store_true", help="build model/data, run 2 steps on CPU-sized batch, exit")
    ap.add_argument("--device", default="auto")
    ap.add_argument("--seed", type=int, default=None)
    ap.add_argument("--resume", default=None, help="checkpoint dir to resume weights from")
    ap.add_argument("--run-root", default="experiments/runs")
    args = ap.parse_args(argv)
    exp = load_yaml(args.config)
    seed = args.seed if args.seed is not None else exp.get("seed", 1)
    seed_all(seed)
    device = pick_device(args.device)
    cfg = build_model_cfg(exp)
    cfg.use_workspace = False  # Stage A: plain backbone
    tr = exp["train"]
    if args.dry_run:
        tr = {**tr, "steps": 2, "batch_size": 2, "seq_len": min(64, tr["seq_len"]), "eval_batches": 1, "eval_every": 1, "ckpt_every": 10**9}
        device = torch.device("cpu")
    model = PrismeModel(cfg).to(device)
    n_params = sum(p.numel() for p in model.parameters())
    train_bin = TokenBin(exp["data"]["train_bin"], tr["seq_len"])
    val_bin = TokenBin(exp["data"]["val_bin"], tr["seq_len"])
    run = Run(exp.get("name", "pretrain"), {"experiment": exp, "model": cfg.to_dict(), "resolved_train": tr}, root=args.run_root, seed=seed)
    run.write_manifest(device=str(device), backend=device.type, precision="fp32", n_params=n_params,
                       dataset_manifest_sha256=sha256_file(Path(exp["data"]["train_bin"]).parent / "manifest.json") if (Path(exp["data"]["train_bin"]).parent / "manifest.json").exists() else None,
                       checkpoint_parent=args.resume)
    if args.resume:
        meta = load_checkpoint(args.resume, model, strict=False)
        run.log(event="resumed", missing=meta["_missing"], unexpected=meta["_unexpected"])
    opt = make_adamw(model.parameters(), tr["lr"], tr.get("weight_decay", 0.1))
    steps = tr["steps"]
    tokens_per_step = tr["batch_size"] * tr["seq_len"]
    batches = train_bin.batches(tr["batch_size"], seed, device)
    model.train()
    t0 = time.time(); seen = 0; best_val = math.inf
    for step in range(steps):
        lr = cosine_lr(step, tr.get("warmup", 200), steps, tr["lr"])
        for g in opt.param_groups:
            g["lr"] = lr
        x, y = next(batches)
        loss = lm_loss(model(x, None)["logits"], y)
        opt.zero_grad(set_to_none=True)
        loss.backward()
        gn = torch.nn.utils.clip_grad_norm_(model.parameters(), tr.get("grad_clip", 1.0))
        opt.step()
        seen += tokens_per_step
        if step % tr.get("log_every", 10) == 0 or step == steps - 1:
            sync(device)
            el = time.time() - t0
            run.log(step=step, loss=loss.item(), lr=lr, grad_norm=float(gn), tokens=seen, tok_per_s=seen / max(el, 1e-6))
            print(f"step {step} loss {loss.item():.4f} lr {lr:.2e} gn {float(gn):.2f} tok/s {seen/max(el,1e-6):.0f}", flush=True)
        if (step + 1) % tr["eval_every"] == 0 or step == steps - 1:
            vl = evaluate(model, val_bin, tr["batch_size"], device, tr.get("eval_batches", 20))
            run.log(step=step, val_loss=vl, val_ppl=math.exp(vl), tokens=seen)
            print(f"  val loss {vl:.4f} ppl {math.exp(vl):.2f}", flush=True)
            if vl < best_val:
                best_val = vl
        if (step + 1) % tr["ckpt_every"] == 0 or step == steps - 1:
            sha = save_checkpoint(run.dir / "checkpoint", model, cfg.to_dict(), {"step": step + 1, "tokens": seen, "run_id": run.run_id})
            run.log(step=step, checkpoint_sha256=sha)
    sync(device)
    run.finish(train_tokens=seen, estimated_train_flops=6 * n_params * seen, peak_memory_bytes=peak_memory_bytes(device),
               final_val_loss=best_val, metrics_file=str(run.dir / "metrics.jsonl"))
    print("run", run.run_id, "done; val", best_val)
    return 0


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