"""Full-parameter SFT of Dory1.0-4B on Beam traces (torchrun, DDP, packed sequences).

  torchrun --nproc_per_node 8 scripts/train_dory.py --data a.jsonl b.jsonl --out /ckpt/run1

Master weights and AdamW state are fp32; matmuls run under bf16 autocast. After every
optimizer step the matrices are projected back to unit norm (nGPT). Loops per step
are either fixed (--loops 2) or sampled identically on all ranks (--loops 1,2,3,4
--loop-probs 0.2,0.4,0.2,0.2), so the loop count stays elastic after training.
"""

import argparse
import contextlib
import json
import math
import os
import random
import shutil
import sys
import time
from pathlib import Path

import torch
import torch.distributed as dist
from safetensors.torch import save_file
from torch.nn.parallel import DistributedDataParallel as DDP

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

BASE_FILES = ("config.json", "generation_config.json", "tokenizer.json", "tokenizer_config.json",
              "special_tokens_map.json", "chat_template.jinja")


def log0(*a):
    if int(os.environ.get("RANK", 0)) == 0:
        print(*a, flush=True)


def load_examples(paths, tok, max_len, medium_effort):
    out, skipped = [], 0
    for path in paths:
        for line in open(path):
            ex = json.loads(line)
            enc = encode_example(tok, ex["messages"], medium_effort=medium_effort, max_len=max_len)
            if enc is None:
                skipped += 1
                continue
            out.append((ex["id"], enc[0], enc[1]))
    return out, skipped


def pack(examples, seq_len, seed):
    """First-fit-decreasing into bins of <= seq_len tokens; returns list of bins (lists of examples)."""
    order = sorted(examples, key=lambda e: -len(e[1]))
    bins, space = [], []
    for ex in order:
        n = len(ex[1])
        for i, free in enumerate(space):
            if n <= free:
                bins[i].append(ex)
                space[i] -= n
                break
        else:
            bins.append([ex])
            space.append(seq_len - n)
    random.Random(seed).shuffle(bins)
    return bins


def collate(bin_, seq_len, pad_id):
    """Concatenate a bin's documents and pad to exactly seq_len (one static shape, so
    FlexAttention compiles once). Padding gets doc id -1: it attends only to itself."""
    ids, labels, docs, pos = [], [], [], []
    for d, (_, i, l) in enumerate(bin_):
        l = l.clone()
        l[0] = -100  # no in-document predecessor
        ids.append(i)
        labels.append(l)
        docs.append(torch.full((len(i),), d, dtype=torch.int32))
        pos.append(torch.arange(len(i)))
    pad = seq_len - sum(len(i) for i in ids)
    if pad > 0:
        ids.append(torch.full((pad,), pad_id, dtype=torch.long))
        labels.append(torch.full((pad,), -100, dtype=torch.long))
        docs.append(torch.full((pad,), -1, dtype=torch.int32))
        pos.append(torch.zeros(pad, dtype=torch.long))
    return {k: torch.cat(v)[None] for k, v in
            (("input_ids", ids), ("labels", labels), ("doc_ids", docs), ("position_ids", pos))}


def lr_at(step, total, peak, warmup, floor=0.1):
    if step < warmup:
        return peak * (step + 1) / warmup
    p = (step - warmup) / max(1, total - warmup)
    return peak * (floor + (1 - floor) * 0.5 * (1 + math.cos(math.pi * p)))


@torch.no_grad()
def evaluate(model, bins, dev, loops_list, attn, seq_len, pad_id):
    model.eval()
    res = {}
    for loops in loops_list:
        tot = torch.zeros(2, device=dev, dtype=torch.float64)
        for b in bins:
            batch = {k: v.to(dev) for k, v in collate(b, seq_len, pad_id).items()}
            with torch.autocast("cuda", dtype=torch.bfloat16):
                loss, n = model(**batch, n_loops=loops, attn_impl=attn)
            tot += torch.tensor([float(loss) * n, n], device=dev, dtype=torch.float64)
        dist.all_reduce(tot)
        res[loops] = float(tot[0] / tot[1])
    model.train()
    return res


def save(model, base, out, args, extra):
    out.mkdir(parents=True, exist_ok=True)
    state = {k: v.detach().to(torch.bfloat16).contiguous().cpu() for k, v in model.state_dict().items()}
    save_file(state, out / "model.safetensors", metadata={"format": "pt"})
    for f in BASE_FILES:
        if (base / f).exists():
            shutil.copy(base / f, out / f)
    (out / "training_args.json").write_text(json.dumps({**vars(args), **extra}, indent=1))


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--data", nargs="+", required=True)
    ap.add_argument("--model", required=True)
    ap.add_argument("--out", required=True)
    ap.add_argument("--epochs", type=float, default=2)
    ap.add_argument("--lr", type=float, default=1e-5)
    ap.add_argument("--warmup-frac", type=float, default=0.03)
    ap.add_argument("--seq-len", type=int, default=16384)
    ap.add_argument("--max-example-len", type=int, default=12288)
    ap.add_argument("--grad-accum", type=int, default=1)
    ap.add_argument("--loops", default="2")
    ap.add_argument("--loop-probs", default=None)
    ap.add_argument("--medium-effort", action="store_true", help="train under Dory's efficient-reasoning tag")
    ap.add_argument("--attn", default="flex", choices=["flex", "sdpa"])
    ap.add_argument("--eval-frac", type=float, default=0.02)
    ap.add_argument("--eval-loops", default="1,2,4")
    ap.add_argument("--max-steps", type=int, default=0, help="stop early (smoke tests)")
    ap.add_argument("--save-every", type=int, default=0)
    ap.add_argument("--seed", type=int, default=1234)
    ap.add_argument("--check-attn", action="store_true", help="compare flex vs sdpa loss before training")
    args = ap.parse_args()

    dist.init_process_group("nccl")
    rank, world = dist.get_rank(), dist.get_world_size()
    local = int(os.environ["LOCAL_RANK"])
    torch.cuda.set_device(local)
    dev = torch.device("cuda", local)
    torch.manual_seed(args.seed)
    torch.backends.cuda.matmul.allow_tf32 = True

    base, out = Path(args.model), Path(args.out)
    tok = PreTrainedTokenizerFast.from_pretrained(base)
    pad_id = tok.pad_token_id if tok.pad_token_id is not None else 0
    t0 = time.time()
    examples, skipped = load_examples(args.data, tok, args.max_example_len, args.medium_effort)
    rng = random.Random(args.seed)
    rng.shuffle(examples)
    n_eval = max(world, int(len(examples) * args.eval_frac))
    eval_ex, train_ex = examples[:n_eval], examples[n_eval:]
    eval_bins = pack(eval_ex, args.seq_len, args.seed)[rank::world]
    train_tokens = sum(len(e[1]) for e in train_ex)
    target_tokens = sum(int((e[2] != -100).sum()) for e in train_ex)
    log0(f"examples: {len(train_ex)} train / {len(eval_ex)} eval ({skipped} over length), "
         f"{train_tokens / 1e6:.1f}M tokens/epoch ({target_tokens / 1e6:.1f}M trained), "
         f"tokenized in {time.time() - t0:.0f}s")

    model = DoryForCausalLM.from_pretrained(base, dtype=torch.float32, device=dev)
    model.gradient_checkpointing_enable()
    model.train()
    ddp = DDP(model, device_ids=[local], gradient_as_bucket_view=True)
    opt = torch.optim.AdamW(model.parameters(), lr=args.lr, betas=(0.9, 0.95), eps=1e-8,
                            weight_decay=0.0, fused=True)

    loops = [int(x) for x in args.loops.split(",")]
    probs = [float(x) for x in args.loop_probs.split(",")] if args.loop_probs else [1 / len(loops)] * len(loops)
    eval_loops = [int(x) for x in args.eval_loops.split(",")]

    bins_epoch0 = pack(train_ex, args.seq_len, args.seed)
    steps_per_epoch = len(bins_epoch0) // (world * args.grad_accum)
    total_steps = int(steps_per_epoch * args.epochs)
    if args.max_steps:
        total_steps = min(total_steps, args.max_steps)
    warmup = max(1, int(total_steps * args.warmup_frac))
    log0(f"world={world} bins/epoch={len(bins_epoch0)} steps/epoch={steps_per_epoch} total_steps={total_steps}")
    if args.check_attn:
        for impl in ("sdpa", "flex"):
            log0(f"attn check {impl}:", evaluate(model, eval_bins[:2], dev, [2], impl, args.seq_len, pad_id))
    log0("eval before:", evaluate(model, eval_bins, dev, eval_loops, args.attn, args.seq_len, pad_id))

    log_f = None
    if rank == 0:
        out.mkdir(parents=True, exist_ok=True)
        log_f = open(out / "train_log.jsonl", "a")
    step, epoch, t_last, tok_since = 0, 0, time.time(), 0
    while step < total_steps:
        bins = pack(train_ex, args.seq_len, args.seed + epoch)
        mine = bins[rank::world][: steps_per_epoch * args.grad_accum]
        for i in range(0, len(mine), args.grad_accum):
            if step >= total_steps:
                break
            micro = [collate(b, args.seq_len, pad_id) for b in mine[i : i + args.grad_accum]]
            n_loops = random.Random(args.seed * 100003 + step).choices(loops, probs)[0]
            local_tgt = sum(int((m["labels"][:, 1:] != -100).sum()) for m in micro)
            tgt = torch.tensor(local_tgt, device=dev, dtype=torch.float64)
            dist.all_reduce(tgt)
            for g in opt.param_groups:
                g["lr"] = lr_at(step, total_steps, args.lr, warmup)
            loss_sum = 0.0
            for j, m in enumerate(micro):
                m = {k: v.to(dev, non_blocking=True) for k, v in m.items()}
                sync = j == len(micro) - 1
                ctx = ddp.no_sync() if not sync else contextlib.nullcontext()
                with ctx, torch.autocast("cuda", dtype=torch.bfloat16):
                    loss, n = ddp(**m, n_loops=n_loops, attn_impl=args.attn)
                    # token-weighted mean over the global batch; DDP averages over ranks
                    (loss * n * world / float(tgt)).backward()
                loss_sum += float(loss.detach()) * n
                tok_since += int((m["doc_ids"] >= 0).sum())
            gnorm = torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            opt.step()
            opt.zero_grad(set_to_none=True)
            model.normalize_matrices_()
            step += 1
            stats = torch.tensor([loss_sum, local_tgt], device=dev, dtype=torch.float64)
            dist.all_reduce(stats)
            if rank == 0:
                dt = time.time() - t_last
                rec = {"step": step, "epoch": round(step / max(1, steps_per_epoch), 3),
                       "loss": round(float(stats[0] / stats[1]), 4), "loops": n_loops,
                       "lr": opt.param_groups[0]["lr"], "grad_norm": round(float(gnorm), 3),
                       "tok_per_s": round(tok_since * world / dt), "time": time.strftime("%H:%M:%S")}
                log_f.write(json.dumps(rec) + "\n")
                log_f.flush()
                if step % 10 == 0 or step <= 5:
                    print(json.dumps(rec), flush=True)
            t_last, tok_since = time.time(), 0
            if args.save_every and step % args.save_every == 0 and rank == 0:
                save(model, base, out / f"step-{step}", args, {"step": step})
        epoch += 1
        ev = evaluate(model, eval_bins, dev, eval_loops, args.attn, args.seq_len, pad_id)
        log0(f"eval after epoch {epoch}:", ev)

    if rank == 0:
        save(model, base, out / "final", args, {"step": step, "final_eval": ev})
        log0(f"saved {out / 'final'}")
    dist.barrier()
    dist.destroy_process_group()


if __name__ == "__main__":
    main()