"""Text LM training from a world-model (iVideoGPT) init vs. controls.

Every arm uses the same LlamaConfig as iVideoGPT's transformer, re-sized to the
GPT-2 vocab. Only the body (everything but embed_tokens / lm_head) differs:

  --init scratch   HF default random init
  --init ivgpt     iVideoGPT body weights; embeddings + head freshly initialized
  --init shuffled  iVideoGPT body with each tensor's entries randomly permuted
                   (same per-tensor weight statistics, no learned structure)
  --init rescaled  iVideoGPT body with every matrix scaled to std 0.02 (the scratch
                   init scale); keeps learned structure, removes the ~2.5x larger scale

  --init sp        shrink-and-perturb (Ash & Adams 2020): lambda * iVideoGPT + fresh init
                   (--sp_lambda); RMSNorm gains become lambda * g + (1 - lambda)
  --init nobody    no pretrained body at all: only the --extra_layers (baseline for them)

  --freeze body    train only embed_tokens, lm_head and RMSNorm gains (+ LoRA / extra layers)
  --lora_rank R    LoRA on every body linear (frozen base); merged into the weights at save
  --extra_layers K append K freshly initialized trainable layers on top of the body

Launch: torchrun --nproc_per_node=N wp/train.py --init ivgpt --lr 1e-3 ...
"""

import argparse
import json
import math
import os
import time
from contextlib import nullcontext
from pathlib import Path

import numpy as np
import torch
import torch.distributed as dist
from huggingface_hub import hf_hub_download
from safetensors.torch import load_file
from transformers import AutoTokenizer, LlamaConfig, LlamaForCausalLM

GPT2_EOT = 50256
VOCAB = 50304  # GPT-2 vocab padded to a multiple of 128
IO_PREFIXES = ("model.embed_tokens.", "lm_head.")


def parse_args():
    p = argparse.ArgumentParser()
    p.add_argument("--init", choices=["scratch", "ivgpt", "shuffled", "rescaled", "sp", "nobody"], required=True)
    p.add_argument("--freeze", choices=["none", "body"], default="none")
    p.add_argument("--repo", default="thuml/ivideogpt-oxe-64-act-free-medium")
    p.add_argument("--data_dir", default="/data/fineweb-edu-gpt2")
    p.add_argument("--out_dir", required=True)
    p.add_argument("--ckpt_dir", default=None, help="rolling checkpoint location (default: out_dir)")
    p.add_argument("--final_dir", default=None, help="save final bf16 HF model here")
    p.add_argument("--tokens", type=float, default=2e9, help="training token budget")
    p.add_argument("--batch_tokens", type=int, default=524288, help="global tokens per step")
    p.add_argument("--seq_len", type=int, default=1024)
    p.add_argument("--micro_bs", type=int, default=32)
    p.add_argument("--lr", type=float, default=1e-3)
    p.add_argument("--min_lr_frac", type=float, default=0.1)
    p.add_argument("--warmup_steps", type=int, default=200, help="-1: 5%% of steps, capped at 200")
    p.add_argument("--wd", type=float, default=0.1)
    p.add_argument("--grad_clip", type=float, default=1.0)
    p.add_argument("--eval_every", type=int, default=100)
    p.add_argument("--val_tokens", type=int, default=4_000_000)
    p.add_argument("--ckpt_every", type=int, default=500)
    p.add_argument("--seed", type=int, default=0)
    p.add_argument("--no_compile", action="store_true")
    p.add_argument("--sp_lambda", type=float, default=0.5)
    p.add_argument("--sp_only", choices=["all", "attn", "mlp", "lower", "upper"], default="all",
                   help="apply S&P only to these weights; the rest keep the scratch init")
    p.add_argument("--sp_shuffle", action="store_true", help="shrink-and-perturb the shuffled body (statistics-matched control)")
    p.add_argument("--lora_rank", type=int, default=0)
    p.add_argument("--extra_layers", type=int, default=0)
    return p.parse_args()


def build_model(args):
    cfg = LlamaConfig.from_pretrained(args.repo, subfolder="transformer")
    n_body = 0 if args.init == "nobody" else cfg.num_hidden_layers
    cfg.num_hidden_layers = n_body + args.extra_layers
    cfg.vocab_size = VOCAB
    cfg.bos_token_id = cfg.eos_token_id = GPT2_EOT
    cfg.attention_dropout = 0.0
    cfg.max_position_embeddings = args.seq_len
    cfg.tie_word_embeddings = False
    torch.manual_seed(args.seed)
    model = LlamaForCausalLM(cfg)
    model.n_body = n_body
    if args.init in ("scratch", "nobody"):
        return model

    sd = load_file(hf_hub_download(args.repo, "transformer/model.safetensors"))
    # action-conditioned checkpoints nest the Llama under "llm." and add an action input layer
    sd = {k.removeprefix("llm."): v for k, v in sd.items() if not k.startswith("action_linear.")}
    body = {k: v for k, v in sd.items() if not k.startswith(IO_PREFIXES) and "rotary_emb" not in k}
    if args.init == "shuffled" or (args.init == "sp" and args.sp_shuffle):
        g = torch.Generator().manual_seed(args.seed + 1)
        body = {k: v.flatten()[torch.randperm(v.numel(), generator=g)].view_as(v) for k, v in body.items()}
    if args.init == "rescaled":
        body = {k: v * (cfg.initializer_range / v.std()) if v.dim() == 2 else v for k, v in body.items()}
    if args.init == "sp":
        fresh = model.state_dict()  # the random init drawn above is the perturbation
        lam = args.sp_lambda
        def selected(k):
            layer = int(k.split(".")[2]) if k.startswith("model.layers.") else -1
            return {"all": True, "attn": "self_attn" in k or "input_layernorm" in k,
                    "mlp": ".mlp." in k or "post_attention_layernorm" in k,
                    "lower": 0 <= layer < cfg.num_hidden_layers // 2,
                    "upper": layer >= cfg.num_hidden_layers // 2}[args.sp_only]
        body = {k: (lam * v + (fresh[k] if v.dim() == 2 else (1 - lam))) if selected(k) else fresh[k]
                for k, v in body.items()}
    missing, unexpected = model.load_state_dict(body, strict=False)
    assert not unexpected, f"unexpected keys: {unexpected}"
    is_extra = lambda k: k.startswith("model.layers.") and int(k.split(".")[2]) >= n_body
    assert all(k.startswith(IO_PREFIXES) or is_extra(k) for k in missing), f"missing body keys: {missing}"
    return model


class LoRALinear(torch.nn.Module):
    """y = W x + (B A x) * alpha/r with W frozen; B starts at zero so the init function is unchanged."""

    def __init__(self, base, r, alpha):
        super().__init__()
        self.base = base
        self.lora_A = torch.nn.Parameter(torch.empty(r, base.in_features))
        torch.nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5))
        self.lora_B = torch.nn.Parameter(torch.zeros(base.out_features, r))
        self.scale = alpha / r

    def forward(self, x):
        return self.base(x) + (x @ self.lora_A.t()) @ self.lora_B.t() * self.scale

    def merged(self):
        lin = self.base
        lin.weight.data += (self.lora_B @ self.lora_A).to(lin.weight.dtype) * self.scale
        return lin


LORA_TARGETS = {"self_attn": ("q_proj", "k_proj", "v_proj", "o_proj"), "mlp": ("gate_proj", "up_proj", "down_proj")}


def add_lora(model, r):
    for layer in model.model.layers[: model.n_body]:
        for parent, names in LORA_TARGETS.items():
            mod = getattr(layer, parent)
            for n in names:
                setattr(mod, n, LoRALinear(getattr(mod, n), r, alpha=r))


def merge_lora(model):
    for layer in model.model.layers:
        for parent, names in LORA_TARGETS.items():
            mod = getattr(layer, parent)
            for n in names:
                if isinstance(getattr(mod, n), LoRALinear):
                    setattr(mod, n, getattr(mod, n).merged())


def apply_freeze(model, freeze):
    if freeze == "none":
        return
    for name, p in model.named_parameters():
        extra = name.startswith("model.layers.") and int(name.split(".")[2]) >= model.n_body
        p.requires_grad = name.startswith(IO_PREFIXES) or "norm" in name or "lora_" in name or extra


class Batches:
    """Random seq_len+1 windows from a uint16 token memmap; same order for every arm."""

    def __init__(self, path, seq_len, micro_bs, seed, rank):
        self.data = np.memmap(path, dtype=np.uint16, mode="r")
        self.seq_len, self.micro_bs = seq_len, micro_bs
        self.rng = np.random.default_rng([seed, rank])

    def next(self, device):
        idx = self.rng.integers(0, len(self.data) - self.seq_len - 1, self.micro_bs)
        x = np.stack([self.data[i : i + self.seq_len + 1] for i in idx]).astype(np.int64)
        x = torch.from_numpy(x).pin_memory().to(device, non_blocking=True)
        return x[:, :-1], x[:, 1:]


@torch.no_grad()
def evaluate(model, val, args, rank, world, device):
    """Mean loss over a fixed set of non-overlapping windows from the start of val.bin."""
    model.eval()
    n_seq = args.val_tokens // args.seq_len
    starts = np.arange(n_seq)[rank::world] * args.seq_len
    tot, cnt = torch.zeros((), device=device), torch.zeros((), device=device)
    for i in range(0, len(starts), args.micro_bs):
        x = np.stack([val[s : s + args.seq_len + 1] for s in starts[i : i + args.micro_bs]]).astype(np.int64)
        x = torch.from_numpy(x).to(device)
        with torch.autocast("cuda", dtype=torch.bfloat16):
            logits = model(input_ids=x[:, :-1]).logits
        tot += torch.nn.functional.cross_entropy(
            logits.float().flatten(0, 1), x[:, 1:].flatten(), reduction="sum"
        )
        cnt += x[:, 1:].numel()
    if world > 1:
        dist.all_reduce(tot)
        dist.all_reduce(cnt)
    model.train()
    return (tot / cnt).item()


def lr_at(step, total, args):
    if step < args.warmup_steps:
        return args.lr * (step + 1) / args.warmup_steps
    t = (step - args.warmup_steps) / max(1, total - args.warmup_steps)
    return args.lr * (args.min_lr_frac + (1 - args.min_lr_frac) * 0.5 * (1 + math.cos(math.pi * t)))


def main():
    args = parse_args()
    ddp = "RANK" in os.environ
    if ddp:
        dist.init_process_group("nccl")
        rank, world = dist.get_rank(), dist.get_world_size()
        local = int(os.environ["LOCAL_RANK"])
    else:
        rank, world, local = 0, 1, 0
    device = torch.device("cuda", local)
    torch.cuda.set_device(device)
    torch.backends.cuda.matmul.allow_tf32 = True
    torch.backends.cudnn.allow_tf32 = True
    master = rank == 0
    out = Path(args.out_dir)
    if master:
        out.mkdir(parents=True, exist_ok=True)

    tok_per_micro = args.micro_bs * args.seq_len * world
    assert args.batch_tokens % tok_per_micro == 0, "batch_tokens must divide micro_bs*seq_len*world"
    accum = args.batch_tokens // tok_per_micro
    total_steps = int(args.tokens // args.batch_tokens)
    if args.warmup_steps < 0:
        args.warmup_steps = min(200, max(10, int(0.05 * total_steps)))

    model = build_model(args)
    if args.lora_rank:
        add_lora(model, args.lora_rank)
    model = model.to(device)
    apply_freeze(model, args.freeze)
    params = [p for p in model.parameters() if p.requires_grad]
    n_total = sum(p.numel() for p in model.parameters())
    n_train = sum(p.numel() for p in params)
    decay = [p for p in params if p.dim() >= 2]
    no_decay = [p for p in params if p.dim() < 2]
    opt = torch.optim.AdamW(
        [{"params": decay, "weight_decay": args.wd}, {"params": no_decay, "weight_decay": 0.0}],
        lr=args.lr, betas=(0.9, 0.95), eps=1e-8, fused=True,
    )

    step = 0
    ckpt_path = Path(args.ckpt_dir or out) / "ckpt.pt"
    if master:
        ckpt_path.parent.mkdir(parents=True, exist_ok=True)
    if ckpt_path.exists():
        ck = torch.load(ckpt_path, map_location=device)
        model.load_state_dict(ck["model"])
        opt.load_state_dict(ck["opt"])
        step = ck["step"]
        if master:
            print(f"resumed from step {step}")

    fwd = model if args.no_compile else torch.compile(model)
    if ddp:
        fwd = torch.nn.parallel.DistributedDataParallel(fwd, device_ids=[local])

    train = Batches(Path(args.data_dir) / "train.bin", args.seq_len, args.micro_bs, args.seed, rank)
    for _ in range(step * accum):  # replay the sampler so resumed runs see the same data
        train.rng.integers(0, len(train.data) - args.seq_len - 1, args.micro_bs)
    val = np.memmap(Path(args.data_dir) / "val.bin", dtype=np.uint16, mode="r")

    def log(rec):
        if master:
            print(json.dumps(rec), flush=True)
            with open(out / "metrics.jsonl", "a") as f:
                f.write(json.dumps(rec) + "\n")

    if master and step == 0:
        (out / "args.json").write_text(json.dumps(vars(args) | {"world": world, "accum": accum}, indent=2))
        log({"step": 0, "tokens": 0, "val_loss": evaluate(model, val, args, rank, world, device),
             "n_params": n_total, "n_trainable": n_train})
    elif step == 0:
        evaluate(model, val, args, rank, world, device)

    t0 = time.time()
    while step < total_steps:
        lr = lr_at(step, total_steps, args)
        for g in opt.param_groups:
            g["lr"] = lr
        loss_acc = 0.0
        for micro in range(accum):
            x, y = train.next(device)
            sync = nullcontext() if (not ddp or micro == accum - 1) else fwd.no_sync()
            with sync, torch.autocast("cuda", dtype=torch.bfloat16):
                logits = fwd(input_ids=x).logits
                loss = torch.nn.functional.cross_entropy(logits.float().flatten(0, 1), y.flatten())
                (loss / accum).backward()
            loss_acc += loss.item() / accum
        gnorm = torch.nn.utils.clip_grad_norm_(params, args.grad_clip).item()
        opt.step()
        opt.zero_grad(set_to_none=True)
        step += 1

        if step % 10 == 0 and master:
            dt = time.time() - t0
            t0 = time.time()
            log({"step": step, "tokens": step * args.batch_tokens, "train_loss": round(loss_acc, 4),
                 "lr": lr, "gnorm": round(gnorm, 3), "tok_per_s": int(10 * args.batch_tokens / dt)})
        if step % args.eval_every == 0 or step == total_steps:
            vl = evaluate(model, val, args, rank, world, device)
            log({"step": step, "tokens": step * args.batch_tokens, "val_loss": vl})
        if master and (step % args.ckpt_every == 0 or step == total_steps):
            torch.save({"model": model.state_dict(), "opt": opt.state_dict(), "step": step}, ckpt_path.with_suffix(".tmp"))
            os.replace(ckpt_path.with_suffix(".tmp"), ckpt_path)

    if master:
        if args.final_dir:
            merge_lora(model)
            model.to(torch.bfloat16).save_pretrained(args.final_dir)
            AutoTokenizer.from_pretrained("gpt2").save_pretrained(args.final_dir)
        (out / "DONE").touch()
    if ddp:
        dist.destroy_process_group()


if __name__ == "__main__":
    main()