"""Part 2: train the dense decoder-only model with each from-scratch optimizer.

The training regime is one pass over the corpus ("1x the dataset"), with
validation loss and test BLEU measured every 0.1x, mirroring Figure 2 of
*Fantastic Pretraining Optimizers and Where to Find Them*.

    python -m src.part2.train --optimizer muon
    python -m src.part2.train --all
"""

from __future__ import annotations

import argparse
import math
import time

import numpy as np
import torch
import torch.nn.functional as F

from ..part1.config import TransformerConfig
from ..part1.model import Transformer
from ..utils import RUNS, human, save_json, set_seed
from .data import EOS, Blocks, bleu_prompts, get_stream, get_tokenizer, make_loader
from .optimizers import CATEGORY, DEFAULT_HPARAMS, OPTIMIZERS
from .optimizers.sophia import SophiaG


def model_config(vocab_size, n_ctx):
    """The dense baseline of Part 1, reused verbatim."""
    return TransformerConfig(d_model=512, n_ctx=n_ctx, n_layers=6, n_heads=8,
                             vocab_size=vocab_size, d_ff=2048, use_moe=False,
                             name="part2_dense")


def best_swept_lr(name):
    """Peak LR selected by ``src.part2.sweep``, if that has been run."""
    f = RUNS / "part2" / "sweep.json"
    if not f.exists():
        return None
    from ..utils import load_json
    runs = load_json(f).get(name, [])
    ok = [r for r in runs if not r.get("diverged") and r["val_loss"] == r["val_loss"]]
    return min(ok, key=lambda r: r["val_loss"])["lr"] if ok else None


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


@torch.no_grad()
def eval_loss(model, loader, device, max_batches=50):
    model.eval()
    tot, n = 0.0, 0
    for i, (x, y) in enumerate(loader):
        if i >= max_batches:
            break
        x, y = x.to(device), y.to(device)
        with torch.autocast("cuda", dtype=torch.bfloat16, enabled=device.type == "cuda"):
            logits, _, _ = model(x)
        tot += F.cross_entropy(logits.view(-1, logits.size(-1)).float(),
                               y.reshape(-1)).item()
        n += 1
    model.train()
    loss = tot / max(n, 1)
    return loss, float(np.exp(min(loss, 20)))


@torch.no_grad()
def eval_bleu(model, tok, prompts, device, max_new=48, n_ctx=512, batch=48):
    """Greedy continuation BLEU. All prompts share a length, so no padding."""
    import sacrebleu
    model.eval()
    hyps, refs = [], []
    for s in range(0, len(prompts), batch):
        chunk = prompts[s:s + batch]
        seq = torch.tensor([p for p, _ in chunk], device=device)
        gen = [[] for _ in chunk]
        done = torch.zeros(len(chunk), dtype=torch.bool, device=device)
        for _ in range(max_new):
            with torch.autocast("cuda", dtype=torch.bfloat16, enabled=device.type == "cuda"):
                logits, _, _ = model(seq[:, -n_ctx:])
            nxt = logits[:, -1].float().argmax(-1)
            for j, t in enumerate(nxt.tolist()):
                if not done[j]:
                    if t == EOS:
                        done[j] = True
                    else:
                        gen[j].append(t)
            if bool(done.all()):
                break
            seq = torch.cat([seq, nxt.unsqueeze(1)], dim=1)
        hyps.extend(tok.decode(g) for g in gen)
        refs.extend(r for _, r in chunk)
    model.train()
    return sacrebleu.corpus_bleu(hyps, [refs]).score, list(zip(hyps, refs))[:6]


def build_optimizer(name, model, hparams):
    """Two param groups: matrices get weight decay, gains and embeddings do not."""
    decay = [p for _, p in model.named_parameters() if p.dim() >= 2]
    nodecay = [p for _, p in model.named_parameters() if p.dim() < 2]
    cls = OPTIMIZERS[name]
    if name == "muon":
        # Embeddings and the tied head are excluded from the orthogonalised
        # path and fall back to AdamW, per the reference implementation.
        hidden = [p for n, p in model.named_parameters()
                  if p.dim() == 2 and "embed" not in n and "lm_head" not in n]
        other = [p for n, p in model.named_parameters()
                 if not (p.dim() == 2 and "embed" not in n and "lm_head" not in n)]
        return cls([{"params": hidden, "use_muon": True},
                    {"params": other, "use_muon": False}], **hparams)
    return cls([{"params": decay},
                {"params": nodecay, "weight_decay": 0.0}], **hparams)


def train_one(name, args, streams, device):
    set_seed(args.seed)
    tok = streams["tokenizer"]
    cfg = model_config(args.vocab_size, args.n_ctx)
    model = Transformer(cfg).to(device)
    hparams = dict(DEFAULT_HPARAMS[name])
    if args.use_sweep:
        swept = best_swept_lr(name)
        if swept is not None:
            hparams["lr"] = swept
    if args.lr is not None:
        hparams["lr"] = args.lr
    opt = build_optimizer(name, model, hparams)

    tok_per_step = args.batch_size * args.n_ctx
    total_steps = args.token_budget // tok_per_step
    eval_every = max(total_steps // args.n_evals, 1)
    n_params = sum(p.numel() for p in model.parameters())
    chinchilla = args.token_budget / (20 * n_params)

    print(f"\n=== {name}  [{CATEGORY[name]}] ===", flush=True)
    print(f"    hparams {hparams}", flush=True)
    print(f"    {total_steps} steps x {tok_per_step} tok = "
          f"{human(args.token_budget)} tokens = {chinchilla:.3f}x Chinchilla "
          f"for {human(n_params)} params", flush=True)

    train_loader = make_loader(Blocks(streams["train"], args.n_ctx),
                               args.batch_size, True, args.seed)
    val_loader = make_loader(Blocks(streams["val"], args.n_ctx),
                             args.batch_size, False, num_workers=1)
    prompts = streams["prompts"]

    hist = {"step": [], "tokens": [], "frac": [], "chinchilla": [],
            "train_loss": [], "val_loss": [], "val_ppl": [], "bleu": []}
    it = iter(train_loader)
    t0 = time.perf_counter()
    running, seen = 0.0, 0
    is_sophia = isinstance(opt, SophiaG)

    for step in range(total_steps):
        lr = lr_at(step, total_steps, hparams["lr"], args.warmup)
        mult = lr / hparams["lr"]
        for g in opt.param_groups:
            g["lr"] = lr
            # Muon routes embeddings, the tied head and the 1D gains through an
            # internal AdamW with its own rate.  That rate lives under a
            # different key, so without this it would miss the warmup and the
            # cosine decay entirely and train at full rate from step 0.
            if "adamw_lr" in g:
                g["adamw_lr"] = hparams.get("adamw_lr", 3e-4) * mult
        try:
            x, y = next(it)
        except StopIteration:
            it = iter(train_loader)
            x, y = next(it)
        x, y = x.to(device, non_blocking=True), y.to(device, non_blocking=True)

        with torch.autocast("cuda", dtype=torch.bfloat16, enabled=device.type == "cuda"):
            logits, _, _ = model(x)
        loss = F.cross_entropy(logits.view(-1, logits.size(-1)).float(),
                               y.reshape(-1))
        opt.zero_grad(set_to_none=True)
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        if is_sophia:
            opt.step(bs_scale=tok_per_step)
        else:
            opt.step()

        # Sophia refreshes its curvature estimate every k steps with one extra
        # backward pass on labels sampled from the model's own distribution
        # (the Gauss-Newton-Bartlett estimator).
        if is_sophia and (step + 1) % args.sophia_k == 0:
            with torch.no_grad():
                samp = torch.multinomial(
                    logits.detach().float().view(-1, logits.size(-1)).softmax(-1),
                    1).squeeze(-1)
            with torch.autocast("cuda", dtype=torch.bfloat16, enabled=device.type == "cuda"):
                logits2, _, _ = model(x)
            l2 = F.cross_entropy(logits2.view(-1, logits2.size(-1)).float(), samp)
            opt.zero_grad(set_to_none=True)
            l2.backward()
            opt.update_hessian()
            opt.zero_grad(set_to_none=True)

        running += loss.item()
        seen += 1

        if (step + 1) % eval_every == 0 or step == total_steps - 1:
            vl, vp = eval_loss(model, val_loader, device, args.eval_batches)
            bleu, _ = eval_bleu(model, tok, prompts, device, n_ctx=args.n_ctx)
            toks = (step + 1) * tok_per_step
            hist["step"].append(step + 1)
            hist["tokens"].append(toks)
            hist["frac"].append(toks / args.token_budget)
            hist["chinchilla"].append(toks / (20 * n_params))
            hist["train_loss"].append(running / max(seen, 1))
            hist["val_loss"].append(vl)
            hist["val_ppl"].append(vp)
            hist["bleu"].append(bleu)
            print(f"    {toks/args.token_budget:4.2f}x  step {step+1:5d}  "
                  f"train {running/max(seen,1):.4f}  val {vl:.4f}  "
                  f"ppl {vp:7.2f}  BLEU {bleu:5.2f}  "
                  f"{(time.perf_counter()-t0)/60:.1f}m", flush=True)
            running, seen = 0.0, 0

    minutes = (time.perf_counter() - t0) / 60
    test_loader = make_loader(Blocks(streams["test"], args.n_ctx),
                              args.batch_size, False, num_workers=1)
    test_loss, test_ppl = eval_loss(model, test_loader, device, args.test_batches)
    bleu, samples = eval_bleu(model, tok, prompts, device, n_ctx=args.n_ctx)
    state_mb = opt.memory_bytes() / 1e6

    print(f"    FINAL  test loss {test_loss:.4f}  ppl {test_ppl:.2f}  "
          f"BLEU {bleu:.2f}  state {state_mb:.0f}MB  {minutes:.1f}m", flush=True)

    out = RUNS / "part2" / name
    out.mkdir(parents=True, exist_ok=True)
    torch.save({"model": model.state_dict(), "config": cfg.to_dict()},
               out / "model.pt")
    save_json({
        "optimizer": name, "category": CATEGORY[name], "hparams": hparams,
        "n_params": n_params, "tokens": args.token_budget,
        "chinchilla_ratio": chinchilla,
        "history": hist, "test_loss": test_loss, "test_ppl": test_ppl,
        "test_bleu": bleu, "minutes": minutes,
        "optimizer_state_mb": state_mb,
        "state_multiple": state_mb * 1e6 / (n_params * 4),
        "samples": samples,
    }, out / "metrics.json")
    del model, opt
    torch.cuda.empty_cache()


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--optimizer", default=None)
    ap.add_argument("--all", action="store_true")
    ap.add_argument("--token_budget", type=int, default=None,
                    help="default: 1x the training corpus")
    ap.add_argument("--batch_size", type=int, default=24)
    ap.add_argument("--n_ctx", type=int, default=512)
    ap.add_argument("--vocab_size", type=int, default=16000)
    ap.add_argument("--lr", type=float, default=None)
    ap.add_argument("--use_sweep", action="store_true",
                    help="take each optimizer's peak LR from sweep.json")
    ap.add_argument("--warmup", type=int, default=200)
    ap.add_argument("--n_evals", type=int, default=10)
    ap.add_argument("--eval_batches", type=int, default=60)
    ap.add_argument("--test_batches", type=int, default=200)
    ap.add_argument("--bleu_prompts", type=int, default=200)
    ap.add_argument("--sophia_k", type=int, default=10)
    ap.add_argument("--seed", type=int, default=1337)
    args = ap.parse_args()

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    torch.backends.cuda.matmul.allow_tf32 = True

    tok = get_tokenizer(args.vocab_size)
    streams = {"tokenizer": tok}
    for s in ("train", "val", "test"):
        streams[s] = get_stream(tok, s)
        print(f"[part2] {s}: {human(len(streams[s]))} tokens", flush=True)
    streams["prompts"] = bleu_prompts(tok, n=args.bleu_prompts)

    if args.token_budget is None:
        args.token_budget = int(len(streams["train"]))
        print(f"[part2] token budget = 1.0x corpus = "
              f"{human(args.token_budget)} tokens", flush=True)

    names = list(OPTIMIZERS) if args.all else [args.optimizer]
    for n in names:
        train_one(n, args, streams, device)


if __name__ == "__main__":
    main()