"""Fair per-step cost comparison of the five optimizers.

The wall-clock totals from the training runs are not comparable: they were
produced on a shared GPU, so a run that happened to start while the device was
busy looks slower than one that had it to itself, regardless of the update
rule.  This benchmark instead times every optimizer back-to-back on the same
model, same batch shape and same device, reporting the cost of the optimizer
*step* separately from the forward/backward pass that all of them share.

    python -m src.part2.benchmark
"""

from __future__ import annotations

import argparse
import time

import torch
import torch.nn.functional as F

from ..part1.model import Transformer
from ..utils import RUNS, save_json
from .optimizers import CATEGORY, DEFAULT_HPARAMS, OPTIMIZERS
from .optimizers.sophia import SophiaG
from .train import build_optimizer, model_config


def time_optimizer(name, args, device, warmup=5):
    torch.manual_seed(0)
    cfg = model_config(args.vocab_size, args.n_ctx)
    model = Transformer(cfg).to(device)
    opt = build_optimizer(name, model, dict(DEFAULT_HPARAMS[name]))
    is_sophia = isinstance(opt, SophiaG)
    tok_per_step = args.batch_size * args.n_ctx

    x = torch.randint(0, args.vocab_size, (args.batch_size, args.n_ctx),
                      device=device)
    y = torch.randint(0, args.vocab_size, (args.batch_size, args.n_ctx),
                      device=device)

    fwd_bwd, step_t = 0.0, 0.0
    for i in range(warmup + args.steps):
        if device.type == "cuda":
            torch.cuda.synchronize()
        t0 = time.perf_counter()
        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()
        if device.type == "cuda":
            torch.cuda.synchronize()
        t1 = time.perf_counter()

        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        if is_sophia:
            opt.step(bs_scale=tok_per_step)
        else:
            opt.step()
        if device.type == "cuda":
            torch.cuda.synchronize()
        t2 = time.perf_counter()

        if i >= warmup:                       # discard warmup iterations
            fwd_bwd += t1 - t0
            step_t += t2 - t1

    n_params = sum(p.numel() for p in model.parameters())
    state_mb = opt.memory_bytes() / 1e6
    del model, opt
    torch.cuda.empty_cache()

    fwd_bwd_ms = 1000 * fwd_bwd / args.steps
    step_ms = 1000 * step_t / args.steps
    # Sophia additionally pays for one extra forward+backward every k steps to
    # refresh its curvature estimate.  That cost is real but falls outside
    # opt.step(), so timing the step alone would understate it.
    extra = fwd_bwd_ms / args.sophia_k if is_sophia else 0.0
    return {
        "optimizer": name,
        "category": CATEGORY[name],
        "fwd_bwd_ms": fwd_bwd_ms,
        "step_ms": step_ms,
        "hessian_amortised_ms": extra,
        "total_ms": fwd_bwd_ms + step_ms + extra,
        "state_mb": state_mb,
        "state_multiple": state_mb * 1e6 / (n_params * 4),
    }


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--steps", type=int, default=30)
    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("--sophia_k", type=int, default=10,
                    help="curvature refresh interval, to amortise its cost")
    args = ap.parse_args()

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"[benchmark] {args.steps} timed steps each, batch "
          f"{args.batch_size}x{args.n_ctx}, on {device}", flush=True)
    print(f"\n{'optimizer':10s} {'fwd+bwd':>9s} {'opt.step':>9s} "
          f"{'hess':>7s} {'total':>9s} {'overhead':>9s} {'state':>8s}")
    out = []
    for name in OPTIMIZERS:
        r = time_optimizer(name, args, device)
        out.append(r)
        over = 100 * (r["step_ms"] + r["hessian_amortised_ms"]) / r["total_ms"]
        print(f"{name:10s} {r['fwd_bwd_ms']:8.1f}ms {r['step_ms']:8.1f}ms "
              f"{r['hessian_amortised_ms']:6.1f}ms {r['total_ms']:8.1f}ms "
              f"{over:8.1f}% {r['state_multiple']:7.1f}x", flush=True)

    base = next(r for r in out if r["optimizer"] == "adamw")
    print(f"\nrelative to AdamW (total ms/step):")
    for r in out:
        print(f"  {r['optimizer']:10s} {r['total_ms']/base['total_ms']:.3f}x")
    save_json(out, RUNS / "part2" / "benchmark.json")
    print(f"\nwrote {RUNS/'part2'/'benchmark.json'}")


if __name__ == "__main__":
    main()