"""
Reproducible speed / VRAM / accuracy benchmarks for the toolkit quantization
backends.

Compares bf16 against the custom OstrisLinear backends (convrot8, convrot4 for
now; add more qtypes to QTYPES as they land).

Measures, per qtype:
  - layer inference latency across DiT-representative shapes (vs bf16),
    eager and torch.compile'd
  - layer training latency (forward + backward through the frozen layer),
    eager and torch.compile'd
  - VRAM on a transformer-ish block stack: resident weights, peak during a
    no-grad forward, peak during a train step; forward/train-ckpt peaks also
    under torch.compile
  - accuracy drift vs bf16: output relative error per layer shape and
    accumulated through the block stack
  - weight reconstruction error and one-time quantize (conversion) time

Usage:
    python scripts/test_quantizations.py --gpu 1
    python scripts/test_quantizations.py --gpu 1 --qtypes bf16 convrot8
"""

import argparse
import math
import os
import sys
import time

# set cuda bus ordering to be pcie
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

import torch  # noqa: E402

# (tokens, in_features, out_features) — FLUX/Wan-class projections
SPEED_SHAPES = [
    (4096, 3072, 3072),
    (4096, 3072, 12288),
    (4096, 12288, 3072),
    (1024, 3072, 12288),
]

# block stack used for the vram/drift tests (mimics a DiT block's linears)
VRAM_BLOCKS = 8
VRAM_BLOCK_SHAPES = [(3072, 12288), (12288, 3072), (3072, 3072), (3072, 3072)]
VRAM_TOKENS = 4096

QTYPES = [
    "bf16", "qfloat8", "float8", "convrot8", "convrot4",
    "convrotint7", "convrotint6", "convrotint5", "convrotint4", "convrotint3",
    "convrotint2", "convrotbitnet", "convrotcomfyw4a4",
]

STACK_KEY = f"{VRAM_BLOCKS}-block stack"


def convert(module: torch.nn.Linear, qtype: str) -> torch.nn.Linear:
    """Quantize a linear with the given qtype. Returns the (possibly replaced)
    module — quanto swaps the module object, the ostris backends convert in place."""
    if qtype == "bf16":
        return module
    from toolkit.util.ostris_quant import convert_linear_to_ostris, get_ostris_quantizer

    q = get_ostris_quantizer(qtype)
    if q is not None:
        assert convert_linear_to_ostris(module, q), f"conversion refused for {qtype}"
        return module

    # quanto / torchao qtypes go through the shared toolkit quantize flow; use a
    # holder so quanto's module replacement has a parent to swap into
    from optimum.quanto import freeze
    from toolkit.util.quantize import get_qtype, quantize

    holder = torch.nn.Sequential(module)
    quantize(holder, weights=get_qtype(qtype))
    freeze(holder)
    return holder[0]


def fp_weight(module: torch.nn.Linear) -> torch.Tensor:
    """Dequantized weight in float32, whatever the backend."""
    if hasattr(module, "dequantize_weight"):
        return module.dequantize_weight().float()
    w = module.weight
    if hasattr(w, "dequantize"):
        return w.dequantize().float()
    return w.detach().float()


def bench(fn, iters: int, device) -> float:
    for _ in range(max(3, iters // 5)):
        fn()
    torch.cuda.synchronize(device)
    t0 = time.perf_counter()
    for _ in range(iters):
        fn()
    torch.cuda.synchronize(device)
    return (time.perf_counter() - t0) / iters * 1000  # ms


def gb(nbytes: int) -> str:
    return f"{nbytes / 1e9:6.2f} GB"


def make_layer(k: int, n: int, device) -> torch.nn.Linear:
    lin = torch.nn.Linear(k, n, bias=True, dtype=torch.bfloat16, device=device)
    with torch.no_grad():
        lin.weight.mul_(0.02)
    # the train benches model lora-style training: base frozen, grads flow to
    # the input only. without this, bf16/quanto accumulate weight grads that
    # inflate every later vram measurement
    lin.requires_grad_(False)
    return lin


def make_stack(device) -> torch.nn.ModuleList:
    # default nn.Linear init (~1/sqrt(in) std) so block branches contribute at a
    # realistic O(1) scale to the residual stream — scaling weights down further
    # makes accumulated quantization drift look artificially tiny
    torch.manual_seed(0)
    blocks = torch.nn.ModuleList()
    for _ in range(VRAM_BLOCKS):
        blocks.append(torch.nn.ModuleList([
            torch.nn.Linear(k, n, bias=True, dtype=torch.bfloat16, device=device)
            for k, n in VRAM_BLOCK_SHAPES
        ]))
    # frozen base (see make_layer)
    blocks.requires_grad_(False)
    return blocks


def block_forward(b, h):
    # pre-norm residual block like a real transformer, so activations stay at a
    # sane scale and quantization drift accumulates realistically across depth
    r = torch.nn.functional.layer_norm(h, h.shape[-1:])
    r = b[0](r)          # 3072 -> 12288
    r = torch.nn.functional.gelu(r)
    h = h + b[1](r)      # 12288 -> 3072
    r = torch.nn.functional.layer_norm(h, h.shape[-1:])
    return h + b[3](b[2](r))  # 3072 -> 3072 -> 3072


def stack_forward(blocks, x, checkpoint=False):
    h = x
    for b in blocks:
        if checkpoint:
            h = torch.utils.checkpoint.checkpoint(
                block_forward, b, h, use_reentrant=False
            )
        else:
            h = block_forward(b, h)
    return h


def run_speed(qtype: str, device, iters: int, results: dict):
    for m, k, n in SPEED_SHAPES:
        torch.manual_seed(0)
        lin = make_layer(k, n, device)
        lin = convert(lin, qtype)
        x = torch.randn(m, k, device=device, dtype=torch.bfloat16)

        with torch.no_grad():
            t_inf = bench(lambda: lin(x), iters, device)

        def train_step():
            xi = x.detach().requires_grad_(True)
            lin(xi).sum().backward()

        t_train = bench(train_step, max(10, iters // 3), device)
        results[(qtype, "inf", (m, k, n))] = t_inf
        results[(qtype, "train", (m, k, n))] = t_train

        # compiled variants (compilation happens during bench warmup, so it
        # isn't charged to the timing; a backend that won't compile records
        # nothing and shows as '-')
        lin_c = torch.compile(lin, dynamic=False)
        try:
            with torch.no_grad():
                results[(qtype, "inf_comp", (m, k, n))] = bench(
                    lambda: lin_c(x), iters, device
                )
        except Exception as e:
            print(f"  [{qtype}] compiled inference failed for {m}x{k}->{n}: {e}")

        def train_step_c():
            xi = x.detach().requires_grad_(True)
            lin_c(xi).sum().backward()

        try:
            results[(qtype, "train_comp", (m, k, n))] = bench(
                train_step_c, max(10, iters // 3), device
            )
        except Exception as e:
            print(f"  [{qtype}] compiled train failed for {m}x{k}->{n}: {e}")
        torch.cuda.empty_cache()


def _stack_fwd_peak(blocks, x, device, base) -> int:
    # warm up first so lazy-init allocations (and compilation) are not counted
    # as steady-state peak
    with torch.no_grad():
        stack_forward(blocks, x)
    torch.cuda.synchronize(device)
    torch.cuda.reset_peak_memory_stats(device)
    with torch.no_grad():
        stack_forward(blocks, x)
    torch.cuda.synchronize(device)
    return torch.cuda.max_memory_allocated(device) - base


def _stack_train_peak(blocks, x, device, base, checkpoint) -> int:
    # frozen base; grads flow to the input like lora training
    def train_step():
        xi = x.detach().requires_grad_(True)
        stack_forward(blocks, xi, checkpoint).float().pow(2).mean().backward()

    train_step()
    torch.cuda.synchronize(device)
    torch.cuda.reset_peak_memory_stats(device)
    train_step()
    torch.cuda.synchronize(device)
    return torch.cuda.max_memory_allocated(device) - base


def run_vram(qtype: str, device, results: dict):
    torch.cuda.empty_cache()
    base = torch.cuda.memory_allocated(device)

    blocks = make_stack(device)
    for b in blocks:
        for i in range(len(b)):
            b[i] = convert(b[i], qtype)
    torch.cuda.empty_cache()
    results[(qtype, "vram_weights")] = torch.cuda.memory_allocated(device) - base

    x = torch.randn(VRAM_TOKENS, 3072, device=device, dtype=torch.bfloat16)

    results[(qtype, "vram_fwd_peak")] = _stack_fwd_peak(blocks, x, device, base)
    # real training checkpoints, but the plain train peak still gets reported
    results[(qtype, "vram_train_peak")] = _stack_train_peak(blocks, x, device, base, False)
    results[(qtype, "vram_train_ckpt_peak")] = _stack_train_peak(blocks, x, device, base, True)

    # same peaks with every linear compiled (mirrors the trainer's block compile)
    for b in blocks:
        for i in range(len(b)):
            b[i] = torch.compile(b[i], dynamic=False)
    try:
        results[(qtype, "vram_fwd_peak_comp")] = _stack_fwd_peak(blocks, x, device, base)
        results[(qtype, "vram_train_ckpt_peak_comp")] = _stack_train_peak(
            blocks, x, device, base, True
        )
    except Exception as e:
        print(f"  [{qtype}] compiled vram measurement failed: {e}")

    blocks = x = None  # release before the allocator accounting of the next run
    torch.cuda.empty_cache()


def run_drift(qtype: str, device, results: dict):
    """Output error vs the bf16 reference, per layer shape and through the stack."""
    for m, k, n in SPEED_SHAPES:
        torch.manual_seed(0)
        lin = make_layer(k, n, device)
        x = torch.randn(m, k, device=device, dtype=torch.bfloat16)
        with torch.no_grad():
            y_ref = lin(x).float()
            lin = convert(lin, qtype)
            y_q = lin(x).float()
        results[(qtype, "drift", (m, k, n))] = ((y_q - y_ref).norm() / y_ref.norm()).item()
        torch.cuda.empty_cache()

    blocks = make_stack(device)
    x = torch.randn(VRAM_TOKENS, 3072, device=device, dtype=torch.bfloat16)
    with torch.no_grad():
        y_ref = stack_forward(blocks, x).float()
        for b in blocks:
            for i in range(len(b)):
                b[i] = convert(b[i], qtype)
        y_q = stack_forward(blocks, x).float()
    results[(qtype, "drift", STACK_KEY)] = ((y_q - y_ref).norm() / y_ref.norm()).item()
    blocks = x = None
    torch.cuda.empty_cache()


def run_quality_and_quantize_time(qtype: str, device, results: dict):
    torch.manual_seed(0)
    lin = make_layer(3072, 3072, device)
    w0 = lin.weight.detach().float().clone()
    torch.cuda.synchronize(device)
    t0 = time.perf_counter()
    lin = convert(lin, qtype)
    torch.cuda.synchronize(device)
    results[(qtype, "quantize_ms")] = (time.perf_counter() - t0) * 1000
    if qtype == "bf16":
        results[(qtype, "weight_err")] = 0.0
    else:
        wq = fp_weight(lin)
        results[(qtype, "weight_err")] = ((wq - w0).norm() / w0.norm()).item()
    torch.cuda.empty_cache()


def print_speed_table(title: str, kind: str, qts, results):
    # speedups always reference EAGER bf16, compiled kinds included, so the
    # comp columns answer "what do I gain over plain bf16"
    ref_kind = kind.removesuffix("_comp")
    print(f"\n=== {title} (ms; speedup vs eager bf16) ===")
    print(f"{'M x K -> N':<22}" + "".join(f"{qt:>18}" for qt in qts))
    for shape in SPEED_SHAPES:
        m, k, n = shape
        row = f"{f'{m} x {k} -> {n}':<22}"
        ref = results.get(("bf16", ref_kind, shape))
        for qt in qts:
            t = results.get((qt, kind, shape))
            if t is None:
                row += f"{'-':>18}"
                continue
            is_self_ref = qt == "bf16" and kind == ref_kind
            sp = f" ({ref / t:4.2f}x)" if ref and not is_self_ref else " " * 8
            row += f"{t:8.3f}ms{sp}"
        print(row)


def main():
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("--gpu", type=int, default=0, help="cuda device id to run on")
    ap.add_argument("--qtypes", nargs="+", default=QTYPES, help=f"subset of {QTYPES}")
    ap.add_argument("--iters", type=int, default=50, help="timing iterations per case")
    args = ap.parse_args()

    device = torch.device(f"cuda:{args.gpu}")
    torch.cuda.set_device(device)
    props = torch.cuda.get_device_properties(device)
    print(f"device: cuda:{args.gpu} ({props.name}, sm_{props.major}{props.minor}, "
          f"{props.total_memory / 1e9:.0f} GB)")
    print(f"torch {torch.__version__}\n")

    # warm the toolkit import chain (module imports + custom-op registration) so it
    # isn't charged to the first qtype's quantize timing
    if any(qt != "bf16" for qt in args.qtypes):
        from toolkit.util.ostris_quant import get_ostris_quantizer
        for qt in args.qtypes:
            if qt != "bf16":
                get_ostris_quantizer(qt)

    # many distinct module instances share one forward code object; the default
    # per-code cache limit (8) would silently fall back to eager and corrupt the
    # compiled columns
    torch._dynamo.config.cache_size_limit = 4096

    results = {}
    for qt in args.qtypes:
        print(f"benchmarking {qt} ...")
        torch._dynamo.reset()  # drop the previous qtype's compiled artifacts
        run_quality_and_quantize_time(qt, device, results)
        run_drift(qt, device, results)
        run_speed(qt, device, args.iters, results)
        run_vram(qt, device, results)

    qts = args.qtypes
    print_speed_table("layer latency, inference", "inf", qts, results)
    print_speed_table("layer latency, inference (compiled)", "inf_comp", qts, results)
    print_speed_table("layer latency, train fwd+bwd", "train", qts, results)
    print_speed_table("layer latency, train fwd+bwd (compiled)", "train_comp", qts, results)

    print(f"\n=== vram on the block stack ({VRAM_BLOCKS} blocks, {VRAM_TOKENS} tokens) ===")
    print(f"{'':<28}" + "".join(f"{qt:>18}" for qt in qts))
    for key, label in (("vram_weights", "weights resident"),
                       ("vram_fwd_peak", "peak, no-grad fwd"),
                       ("vram_train_peak", "peak, train step"),
                       ("vram_train_ckpt_peak", "peak, train step (ckpt)"),
                       ("vram_fwd_peak_comp", "peak, no-grad fwd (comp)"),
                       ("vram_train_ckpt_peak_comp", "peak, train ckpt (comp)")):
        row = f"{label:<28}"
        for qt in qts:
            v = results.get((qt, key))
            row += f"{gb(v):>18}" if v is not None else f"{'-':>18}"
        print(row)

    print("\n=== accuracy drift vs bf16 (output rel err, no-grad) ===")
    print(f"{'':<28}" + "".join(f"{qt:>18}" for qt in qts))
    for shape in SPEED_SHAPES + [STACK_KEY]:
        label = f"{shape[0]} x {shape[1]} -> {shape[2]}" if isinstance(shape, tuple) else shape
        row = f"{label:<28}"
        for qt in qts:
            row += f"{results[(qt, 'drift', shape)]:>18.5f}"
        print(row)

    print("\n=== quantization ===")
    print(f"{'':<28}" + "".join(f"{qt:>18}" for qt in qts))
    row = f"{'weight rel err':<28}"
    for qt in qts:
        row += f"{results[(qt, 'weight_err')]:>18.5f}"
    print(row)
    row = f"{'quantize time (ms)':<28}"
    for qt in qts:
        row += f"{results[(qt, 'quantize_ms')]:>18.1f}"
    print(row)

    # ---- clean per-qtype breakdown: speed (geomean over shapes) + accuracy ----
    def geomean_speedup(qt, kind):
        # every speedup references EAGER bf16 (compiled kinds included), so the
        # comp columns answer "what do I gain over plain bf16"
        ref_kind = kind.removesuffix("_comp")
        logs = []
        for shape in SPEED_SHAPES:
            ref = results.get(("bf16", ref_kind, shape))
            t = results.get((qt, kind, shape))
            if ref and t:
                logs.append(math.log(ref / t))
        return math.exp(sum(logs) / len(logs)) if logs else None

    def fmt_speed(v):
        return f"{v:.2f}x" if v is not None else "-"

    print("\n=== summary (speed = geomean speedup vs bf16; drift lower is better) ===")
    print(f"{'':<18}{'inference':>12}{'inference comp':>16}{'train':>12}{'train comp':>12}"
          f"{'accuracy drift':>16}{'max vram':>12}{'max vram comp':>15}")
    for qt in qts:
        # real training checkpoints, so the ckpt peak is the meaningful train
        # number; the no-grad fwd peak still matters for sampling
        max_vram = max(results[(qt, "vram_fwd_peak")], results[(qt, "vram_train_ckpt_peak")])
        fwd_c = results.get((qt, "vram_fwd_peak_comp"))
        ckpt_c = results.get((qt, "vram_train_ckpt_peak_comp"))
        max_vram_comp = max(fwd_c, ckpt_c) if fwd_c is not None and ckpt_c is not None else None
        print(f"{qt:<18}"
              f"{fmt_speed(geomean_speedup(qt, 'inf')):>12}"
              f"{fmt_speed(geomean_speedup(qt, 'inf_comp')):>16}"
              f"{fmt_speed(geomean_speedup(qt, 'train')):>12}"
              f"{fmt_speed(geomean_speedup(qt, 'train_comp')):>12}"
              f"{results[(qt, 'drift', STACK_KEY)]:>16.5f}"
              f"{gb(max_vram).strip():>12}"
              f"{(gb(max_vram_comp).strip() if max_vram_comp is not None else '-'):>15}")


if __name__ == "__main__":
    main()