"""Correctness checks for the Dory PyTorch port (dory/modeling.py).

  norms     weight row/column norms (is the checkpoint unit-normalized, nGPT style?)
  nll       assistant-token NLL on held Beam SFT examples at 1-4 loops
  generate  greedy chat generation at 2 loops
  ckpt      gradient checkpointing gives identical loss and gradients
  train     a few full-parameter AdamW steps on 2 examples: loss must fall, no NaNs
"""

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

import torch

ROOT = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(ROOT))
from dory.data import encode_example  # noqa: E402
from dory.modeling import DoryForCausalLM  # noqa: E402
from transformers import PreTrainedTokenizerFast  # noqa: E402


def sync(dev):
    if dev == "mps":
        torch.mps.synchronize()
    elif dev == "cuda":
        torch.cuda.synchronize()


def check_norms(model):
    print("== weight norms (mean ± std); ~1.0 means unit-normalized along that dim")
    layer = model.backbone.layers
    named = {
        "embeddings [V,H]": model.backbone.embeddings.weight,
        "lm_head [V,H]": model.lm_head.weight,
        "L0 q_proj [HD,H]": layer[0].mixer.q_proj.weight,
        "L0 o_proj [H,HD]": layer[0].mixer.o_proj.weight,
        "L1 gate_proj [F,H]": layer[1].mixer.gate_proj.weight,
        "L1 down_proj [H,F]": layer[1].mixer.down_proj.weight,
        "L20 q_proj [HD,H]": layer[20].mixer.q_proj.weight,
    }
    for name, w in named.items():
        w = w.float()
        r, c = w.norm(dim=1), w.norm(dim=0)
        print(f"  {name:22s} rows {r.mean():.4f}±{r.std():.4f}   cols {c.mean():.4f}±{c.std():.4f}")


def load_examples(tok, n, max_len):
    out = []
    for line in open(ROOT / "data" / "final" / "sft_verified.jsonl"):
        ex = json.loads(line)
        enc = encode_example(tok, ex["messages"], max_len=max_len)
        if enc is not None:
            out.append((ex, *enc))
        if len(out) >= n:
            break
    return out


@torch.no_grad()
def check_nll(model, examples, dev):
    print("== assistant-token NLL on Beam traces (lower is better; garbage model ~>5)")
    model.eval()
    for loops in (1, 2, 3, 4):
        tot, cnt, t0 = 0.0, 0, time.time()
        for _, ids, labels in examples:
            loss, n = model(ids[None].to(dev), labels=labels[None].to(dev), n_loops=loops)
            tot += float(loss) * n
            cnt += n
        sync(dev)
        print(f"  loops={loops}: nll={tot / cnt:.4f} ppl={torch.tensor(tot / cnt).exp():.2f} "
              f"over {cnt} tokens ({time.time() - t0:.1f}s)")


@torch.no_grad()
def check_generate(model, tok, dev, prompt, max_new):
    print(f"== greedy generation, 2 loops: {prompt!r}")
    model.eval()
    text = tok.apply_chat_template([{"role": "user", "content": prompt}], tokenize=False, add_generation_prompt=True)
    ids = tok(text, add_special_tokens=False, return_tensors="pt")["input_ids"].to(dev)
    eos = {tok.convert_tokens_to_ids("<|im_end|>")}
    t0 = time.time()
    out = model.generate(ids, max_new_tokens=max_new, n_loops=2, eos_ids=eos)
    sync(dev)
    print(f"  ({out.shape[1]} tokens, {time.time() - t0:.1f}s)\n" + tok.decode(out[0]))


def check_ckpt(model, examples, dev):
    print("== gradient checkpointing equivalence")
    _, ids, labels = examples[0]
    ids, labels = ids[None, :256].to(dev), labels[None, :256].to(dev)
    labels[:, :8] = -100
    labels[:, 8:] = ids[:, 8:]
    layers = model.backbone.layers
    probe = [layers[2].mixer.q_proj.weight, layers[30].attn_alpha, layers[31].mixer.suv_up,
             layers[60].mixer.o_proj.weight, model.lm_head.logit_scale]
    results = []
    for ckpt in (False, True):
        model.train()
        model.backbone.gradient_checkpointing = ckpt
        model.zero_grad(set_to_none=True)
        loss, _ = model(ids, labels=labels, n_loops=2)
        loss.backward()
        results.append((float(loss), [p.grad.float().clone() for p in probe]))
    model.backbone.gradient_checkpointing = False
    model.zero_grad(set_to_none=True)
    (l0, g0), (l1, g1) = results
    diffs = [float((a - b).abs().max() / (a.abs().max() + 1e-12)) for a, b in zip(g0, g1)]
    print(f"  loss {l0:.6f} vs {l1:.6f}; max relative grad diff {max(diffs):.2e}; "
          f"grads finite: {all(torch.isfinite(g).all() for g in g0)}")


@torch.no_grad()
def check_pack(model, examples, dev):
    print("== packing equivalence: 2 docs packed with doc_ids/position_ids vs separately")
    model.eval()
    docs = [(ids[:400], labels[:400]) for _, ids, labels in examples[:2]]
    sep_sum, sep_n = 0.0, 0
    for ids, labels in docs:
        loss, n = model(ids[None].to(dev), labels=labels[None].to(dev), n_loops=2)
        sep_sum += float(loss) * n
        sep_n += n
    ids = torch.cat([d[0] for d in docs])
    labels = torch.cat([d[1] for d in docs]).clone()
    labels[len(docs[0][0])] = -100  # second doc's first token has no in-doc predecessor
    doc_ids = torch.cat([torch.full((len(d[0]),), i) for i, d in enumerate(docs)])
    pos = torch.cat([torch.arange(len(d[0])) for d in docs])
    loss, n = model(ids[None].to(dev), labels=labels[None].to(dev), n_loops=2,
                    doc_ids=doc_ids[None].to(dev), position_ids=pos[None].to(dev))
    print(f"  separate: {sep_sum / sep_n:.6f} over {sep_n} tokens; packed: {float(loss):.6f} over {n} tokens")


def check_train(model, examples, dev, steps, lr):
    print(f"== full-parameter AdamW smoke test: {steps} steps, lr={lr}, 2 loops, checkpointing on")
    model.train()
    model.gradient_checkpointing_enable()
    opt = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=0.0, foreach=False)
    batch = examples[:2]
    for step in range(steps):
        t0 = time.time()
        opt.zero_grad(set_to_none=True)
        tot = 0.0
        for _, ids, labels in batch:
            loss, _ = model(ids[None].to(dev), labels=labels[None].to(dev), n_loops=2)
            (loss / len(batch)).backward()
            tot += float(loss.detach()) / len(batch)
        gnorm = torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        opt.step()
        model.normalize_matrices_()
        sync(dev)
        print(f"  step {step}: loss={tot:.4f} grad_norm={float(gnorm):.3f} ({time.time() - t0:.1f}s)", flush=True)
    check_norms(model)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--model", default=str(ROOT / "models" / "dory-4b"))
    ap.add_argument("--device", default="mps" if torch.backends.mps.is_available() else "cpu")
    ap.add_argument("--checks", default="norms,nll,generate,ckpt")
    ap.add_argument("--n-examples", type=int, default=6)
    ap.add_argument("--max-len", type=int, default=1536)
    ap.add_argument("--max-new", type=int, default=200)
    ap.add_argument("--prompt", default="What is 17 * 23? Reply with just the number.")
    ap.add_argument("--train-steps", type=int, default=8)
    ap.add_argument("--lr", type=float, default=1e-5)
    args = ap.parse_args()
    checks = set(args.checks.split(","))

    torch.manual_seed(0)
    t0 = time.time()
    tok = PreTrainedTokenizerFast.from_pretrained(args.model)
    model = DoryForCausalLM.from_pretrained(args.model, device=args.device)
    print(f"loaded on {args.device} in {time.time() - t0:.1f}s; "
          f"{sum(p.numel() for p in model.parameters()) / 1e9:.3f}B params")
    examples = load_examples(tok, args.n_examples, args.max_len)
    print(f"{len(examples)} SFT examples, lengths {[len(e[1]) for e in examples]}")

    if "norms" in checks:
        check_norms(model)
    if "nll" in checks:
        check_nll(model, examples, args.device)
    if "generate" in checks:
        check_generate(model, tok, args.device, args.prompt, args.max_new)
    if "pack" in checks:
        check_pack(model, examples, args.device)
    if "ckpt" in checks:
        check_ckpt(model, examples, args.device)
    if "train" in checks:
        check_train(model, examples, args.device, args.train_steps, args.lr)


if __name__ == "__main__":
    main()