"""
Optimizer benchmark on a ~500M parameter transformer, run once per dtype
(float32, float16, bfloat16).

Compares speed (ms/step) and peak VRAM across:
  - AdamW (torch, unfused — traditional Python loop)
  - AdamW8bit (bitsandbytes)
  - Adafactor
  - Automagic v1
  - Automagic v2 (fused-backward)
  - Automagic v3 (fused-backward and traditional/unfused)
  - Prodigy
"""
import contextlib
import gc
import io
import os
import sys
import time

import torch
import torch.nn as nn
import torch.nn.functional as F

# Allow running this file directly: `python test_optimizers.py` without setting PYTHONPATH.
# Toolkit imports happen inside main() so they pick this up.
_REPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", ".."))
if _REPO_ROOT not in sys.path:
    sys.path.insert(0, _REPO_ROOT)


# ---- model ---------------------------------------------------------------

class TransformerBlock(nn.Module):
    def __init__(self, d_model: int, n_heads: int, d_ff: int):
        super().__init__()
        self.n_heads = n_heads
        self.d_head = d_model // n_heads
        self.ln1 = nn.LayerNorm(d_model)
        self.q = nn.Linear(d_model, d_model, bias=False)
        self.k = nn.Linear(d_model, d_model, bias=False)
        self.v = nn.Linear(d_model, d_model, bias=False)
        self.o = nn.Linear(d_model, d_model, bias=False)
        self.ln2 = nn.LayerNorm(d_model)
        self.ffn_up = nn.Linear(d_model, d_ff, bias=False)
        self.ffn_down = nn.Linear(d_ff, d_model, bias=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        B, S, D = x.shape
        h = self.ln1(x)
        q = self.q(h).view(B, S, self.n_heads, self.d_head).transpose(1, 2)
        k = self.k(h).view(B, S, self.n_heads, self.d_head).transpose(1, 2)
        v = self.v(h).view(B, S, self.n_heads, self.d_head).transpose(1, 2)
        a = F.scaled_dot_product_attention(q, k, v, is_causal=True)
        a = a.transpose(1, 2).contiguous().view(B, S, D)
        x = x + self.o(a)
        h = self.ln2(x)
        x = x + self.ffn_down(F.gelu(self.ffn_up(h)))
        return x


class Transformer(nn.Module):
    def __init__(self, d_model=1024, n_heads=16, n_layers=40, d_ff=4096):
        super().__init__()
        self.blocks = nn.ModuleList([
            TransformerBlock(d_model, n_heads, d_ff) for _ in range(n_layers)
        ])
        self.norm = nn.LayerNorm(d_model)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        for b in self.blocks:
            x = b(x)
        return self.norm(x)


# ---- benchmark -----------------------------------------------------------

DEVICE = "cuda"
DTYPES = [torch.float32, torch.float16, torch.bfloat16]
D_MODEL = 1024
N_HEADS = 16
N_LAYERS = 40
D_FF = 4096
BATCH = 1
SEQ = 128
WARMUP = 3
ITERS = 10


def build_model(dtype):
    torch.manual_seed(0)
    return Transformer(D_MODEL, N_HEADS, N_LAYERS, D_FF).to(DEVICE, dtype=dtype)


def benchmark(results: list, label: str, opt_factory, dtype):
    gc.collect()
    torch.cuda.empty_cache()
    torch.cuda.reset_peak_memory_stats()

    model = build_model(dtype)
    # Some optimizers print on construction; mute that so the final table is clean.
    with contextlib.redirect_stdout(io.StringIO()):
        opt = opt_factory(model.parameters())
    x = torch.randn(BATCH, SEQ, D_MODEL, device=DEVICE, dtype=dtype)

    print(f"  running {label}...", flush=True)
    try:
        for _ in range(WARMUP):
            opt.zero_grad(set_to_none=True)
            model(x).sum().backward()
            opt.step()
        torch.cuda.synchronize()

        t0 = time.perf_counter()
        for _ in range(ITERS):
            opt.zero_grad(set_to_none=True)
            model(x).sum().backward()
            opt.step()
        torch.cuda.synchronize()
        dt = (time.perf_counter() - t0) / ITERS * 1000
        peak = torch.cuda.max_memory_allocated() / 1024**3
        results.append({"label": label, "ms": dt, "peak": peak, "ok": True})
    except torch.cuda.OutOfMemoryError:
        results.append({"label": label, "ms": float("inf"), "peak": float("inf"), "ok": False, "note": "OOM"})
    except Exception as e:
        # An optimizer may not support a given dtype; record it and keep going.
        print(f"    {label} failed: {type(e).__name__}: {e}", flush=True)
        results.append({"label": label, "ms": float("inf"), "peak": float("inf"), "ok": False, "note": "ERR"})
    finally:
        # Fused-backward optimizers register post-accumulate-grad hooks on the
        # params; without removing them the hook closures keep this model alive
        # past `del`, inflating the *next* fused optimizer's peak. Unregister so
        # each run reports its true isolated peak regardless of run order.
        for h in getattr(opt, "_hook_handles", []):
            h.remove()
        del opt, model
        gc.collect()
        torch.cuda.empty_cache()


def print_table(results: list, dtype_name: str):
    results = sorted(results, key=lambda r: r["peak"])

    headers = ["#", "Optimizer", "Peak VRAM", "Time/step"]
    rows = []
    for i, r in enumerate(results, 1):
        if not r["ok"]:
            rows.append([str(i), r["label"], r.get("note", "OOM"), "-"])
            continue
        rows.append([str(i), r["label"], f"{r['peak']:.2f} GB", f"{r['ms']:.1f} ms"])

    widths = [max(len(str(row[c])) for row in [headers] + rows) for c in range(len(headers))]

    def fmt(row, sep=" │ "):
        return sep.join(s.ljust(widths[c]) if c == 1 else s.rjust(widths[c]) for c, s in enumerate(row))

    line_top = "─" * (sum(widths) + 3 * (len(widths) - 1))
    print()
    print(f"  dtype: {dtype_name}")
    print(line_top)
    print(fmt(headers))
    print(line_top)
    for row in rows:
        print(fmt(row))
    print(line_top)


def main():
    n_params = sum(p.numel() for p in build_model(torch.float32).parameters())
    print(f"Model:  {N_LAYERS} blocks × d_model={D_MODEL} × d_ff={D_FF}")
    print(f"        {n_params/1e6:.1f}M params")
    print(f"Step:   batch={BATCH}, seq={SEQ}")
    print(f"Timing: {WARMUP} warmup + {ITERS} timed iters")
    print(f"Rounds: {', '.join(str(d).replace('torch.', '') for d in DTYPES)}")

    from toolkit.optimizers.automagic import Automagic
    from toolkit.optimizers.automagic2 import Automagic2
    from toolkit.optimizers.automagic3 import Automagic3
    from toolkit.optimizers.adafactor import Adafactor
    from prodigyopt import Prodigy
    import bitsandbytes as bnb

    optimizers = [
        ("AdamW", lambda p: torch.optim.AdamW(p, lr=1e-4, eps=1e-6, foreach=False, fused=False)),
        ("AdamW8bit", lambda p: bnb.optim.AdamW8bit(p, lr=1e-4, eps=1e-6)),
        ("Adafactor", lambda p: Adafactor(p, lr=1e-4, scale_parameter=False, relative_step=False, warmup_init=False)),
        ("Automagic v1", lambda p: Automagic(p, lr=1e-4)),
        ("Automagic v2", lambda p: Automagic2(p, lr=1e-4)),
        ("Automagic v3 fused", lambda p: Automagic3(p, lr=1e-4, fused=True)),
        ("Automagic v3 unfused", lambda p: Automagic3(p, lr=1e-4, fused=False)),
        ("Prodigy", lambda p: Prodigy(p, lr=1.0, eps=1e-6)),
    ]

    for dtype in DTYPES:
        dtype_name = str(dtype).replace("torch.", "")
        print(f"\n=== Round: {dtype_name} ===")
        results: list = []
        for label, factory in optimizers:
            benchmark(results, label, factory, dtype)
        print_table(results, dtype_name)


if __name__ == "__main__":
    main()