#!/usr/bin/env python3
"""Benchmark flashrt-adaptive-norms against PyTorch eager references."""

from __future__ import annotations

import argparse
import ctypes
import ctypes.util
import importlib
import json
import math
import os
import sys
from dataclasses import asdict, dataclass
from pathlib import Path

import torch


ROOT = Path(__file__).resolve().parents[2]
PACKAGE = ROOT / "flashrt-adaptive-norms"
REGISTRATION_INCLUDE = (
    ROOT.parent
    / "kernels"
    / "kernel-builder"
    / "src"
    / "pyproject"
    / "templates"
    / "torch"
)

SHAPES = {
    "small": (64, 1024),
    "vla_2k": (2520, 3072),
    "vla_4k": (4096, 3072),
}
SHAPE_GROUPS = {
    "smoke": ["small"],
    "headline": ["vla_2k"],
    "all": list(SHAPES.keys()),
}


@dataclass
class Result:
    shape: str
    rows: int
    dim: int
    kernel: str
    flashrt_us: float
    torch_eager_us: float
    speedup_vs_eager: float
    p99_abs: float
    cosine: float
    status: str


class SourceOps:
    def __init__(self, namespace: str) -> None:
        self._ops = getattr(torch.ops, namespace)

    def ada_rms_norm_style_bf16(self, x, weight, style, eps, out, gate_out):
        self._ops.ada_rms_norm_style_bf16(x, weight, style, float(eps), out, gate_out)
        return out, gate_out

    def gate_residual_ada_norm_fp8_static_bf16(self, residual, x, gate, weight, style, scale, eps, out, gate_out):
        self._ops.gate_residual_ada_norm_fp8_static_bf16(
            residual, x, gate, weight, style, scale, float(eps), out, gate_out
        )
        return residual, out, gate_out


def _preload_cublaslt() -> None:
    for parent in Path(torch.__file__).resolve().parents:
        candidate = parent / "nvidia" / "cublas" / "lib" / "libcublasLt.so.12"
        if candidate.exists():
            ctypes.CDLL(str(candidate), mode=ctypes.RTLD_GLOBAL)
            return
    library = ctypes.util.find_library("cublasLt")
    if library:
        ctypes.CDLL(library, mode=ctypes.RTLD_GLOBAL)


def _current_arch_list() -> str:
    major, minor = torch.cuda.get_device_capability(0)
    return f"{major}.{minor}"


def load_source_ops() -> SourceOps:
    from torch.utils.cpp_extension import load

    if not REGISTRATION_INCLUDE.is_dir():
        raise RuntimeError(f"missing kernel-builder registration include: {REGISTRATION_INCLUDE}")
    _preload_cublaslt()
    os.environ.setdefault("TORCH_CUDA_ARCH_LIST", _current_arch_list())
    namespace = "flashrt_adaptive_norms_benchmark"
    load(
        name=namespace,
        sources=[
            str(PACKAGE / "torch-ext" / "torch_binding.cpp"),
            str(PACKAGE / "csrc" / "adaptive_norms.cu"),
        ],
        extra_include_paths=[str(PACKAGE / "csrc"), str(REGISTRATION_INCLUDE)],
        extra_cflags=["-O3", "-DCUDA_KERNEL"],
        extra_cuda_cflags=["-O3", "--expt-relaxed-constexpr", "-DCUDA_KERNEL"],
        verbose=False,
    )
    return SourceOps(namespace)


def load_installed_ops(artifact: str | None):
    if artifact:
        sys.path.insert(0, artifact)
    try:
        return importlib.import_module("flashrt_adaptive_norms")
    finally:
        if artifact:
            sys.path.remove(artifact)


def make_case(rows: int, dim: int):
    x = torch.randn((rows, dim), device="cuda", dtype=torch.bfloat16)
    residual = torch.randn_like(x)
    gate = torch.randn_like(x)
    weight = (1.0 + 0.1 * torch.randn((dim,), device="cuda", dtype=torch.bfloat16)).contiguous()
    style = (0.05 * torch.randn((rows, 3 * dim), device="cuda", dtype=torch.bfloat16)).contiguous()
    scale = torch.tensor([0.04], device="cuda", dtype=torch.float32)
    return x, residual, gate, weight, style, scale


def rms_norm(x, weight, eps):
    rms = torch.rsqrt(torch.mean(x.float() * x.float(), dim=-1, keepdim=True) + eps)
    return x.float() * rms * weight.float()


def ref_ada(x, weight, style, eps):
    dim = x.shape[1]
    normed = rms_norm(x, weight, eps)
    style_scale = style[:, :dim].float().expand(x.shape[0], -1)
    style_shift = style[:, dim : 2 * dim].float().expand(x.shape[0], -1)
    style_gate = style[:, 2 * dim :].expand(x.shape[0], -1)
    y = normed * (1.0 + style_scale) + style_shift
    return y.to(torch.bfloat16), style_gate.contiguous().to(torch.bfloat16)


def ref_fused(residual, x, gate, weight, style, scale, eps):
    updated = (residual.float() + x.float() * gate.float()).to(torch.bfloat16)
    dim = updated.shape[1]
    normed = rms_norm(updated, weight, eps)
    style_scale = style[:, :dim].float().expand(updated.shape[0], -1)
    style_shift = style[:, dim : 2 * dim].float().expand(updated.shape[0], -1)
    style_gate = style[:, 2 * dim :].expand(updated.shape[0], -1)
    y = (normed * (1.0 + style_scale) + style_shift) / scale.float().reshape(())
    return updated, y.to(torch.float8_e4m3fn), style_gate.contiguous().to(torch.bfloat16)


def time_us(fn, warmup, iters):
    for _ in range(warmup):
        fn()
    torch.cuda.synchronize()
    start = torch.cuda.Event(enable_timing=True)
    end = torch.cuda.Event(enable_timing=True)
    start.record()
    for _ in range(iters):
        fn()
    end.record()
    torch.cuda.synchronize()
    return start.elapsed_time(end) * 1000.0 / iters


def percentile(x, q):
    flat = x.flatten()
    k = max(1, min(flat.numel(), math.ceil(q * flat.numel())))
    return flat.kthvalue(k).values


def metrics(got, expected):
    diff = (got.float() - expected.float()).abs().flatten()
    return float(percentile(diff, 0.99).item()), float(
        torch.nn.functional.cosine_similarity(got.float().flatten(), expected.float().flatten(), dim=0).item()
    )


def run_one(ops, name, shape, args, broadcast: bool):
    rows, dim = shape
    x, residual, gate, weight, style, scale = make_case(rows, dim)
    if broadcast:
        style = style[:1].contiguous()
        name = f"{name}_broadcast"
    out = torch.empty_like(x)
    gate_out = torch.empty_like(x)
    fp8_out = torch.empty_like(x, dtype=torch.float8_e4m3fn)
    fused_gate_out = torch.empty_like(x)
    ada_got, _ = ops.ada_rms_norm_style_bf16(x, weight, style, args.eps, out, gate_out)
    ada_exp, _ = ref_ada(x, weight, style, args.eps)
    ada_p99, ada_cos = metrics(ada_got, ada_exp)
    ada_flashrt_us = time_us(lambda: ops.ada_rms_norm_style_bf16(x, weight, style, args.eps, out, gate_out), args.warmup, args.iters)
    ada_eager_us = time_us(lambda: ref_ada(x, weight, style, args.eps), args.warmup, args.iters)

    residual_work = residual.clone()
    fused_got = ops.gate_residual_ada_norm_fp8_static_bf16(
        residual_work, x, gate, weight, style, scale, args.eps, fp8_out, fused_gate_out
    )[1]
    _, fused_exp, _ = ref_fused(residual, x, gate, weight, style, scale, args.eps)
    fused_p99, fused_cos = metrics(fused_got.float(), fused_exp.float())
    fused_flashrt_us = time_us(
        lambda: ops.gate_residual_ada_norm_fp8_static_bf16(
            residual_work, x, gate, weight, style, scale, args.eps, fp8_out, fused_gate_out
        ),
        args.warmup,
        args.iters,
    )
    fused_eager_us = time_us(lambda: ref_fused(residual, x, gate, weight, style, scale, args.eps), args.warmup, args.iters)

    return [
        Result(name, rows, dim, "ada_rms_norm_style_bf16", ada_flashrt_us, ada_eager_us, ada_eager_us / ada_flashrt_us, ada_p99, ada_cos, "PASS"),
        Result(name, rows, dim, "gate_residual_ada_norm_fp8_static_bf16", fused_flashrt_us, fused_eager_us, fused_eager_us / fused_flashrt_us, fused_p99, fused_cos, "PASS"),
    ]


def write_markdown(path: Path, results: list[Result]) -> None:
    lines = [
        "# Source Benchmark Results",
        "",
        "Environment: NVIDIA GeForce RTX 5090 local source-extension build.",
        "Baseline: PyTorch eager tensor reference with matching BF16/FP8 math contract.",
        "",
        "| Shape | Rows,Dim | Kernel | FlashRT us | Eager us | vs eager | p99 abs | Cosine | Status |",
        "|---|---:|---|---:|---:|---:|---:|---:|---|",
    ]
    for r in results:
        lines.append(
            f"| {r.shape} | {r.rows},{r.dim} | {r.kernel} | {r.flashrt_us:.3f} | "
            f"{r.torch_eager_us:.3f} | {r.speedup_vs_eager:.2f}x | {r.p99_abs:.6f} | "
            f"{r.cosine:.8f} | {r.status} |"
        )
    path.write_text("\n".join(lines) + "\n")


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--backend", choices=["source", "installed"], default="source")
    parser.add_argument("--artifact", default=None)
    parser.add_argument("--shapes", choices=sorted(SHAPE_GROUPS), default="smoke")
    parser.add_argument("--warmup", type=int, default=5)
    parser.add_argument("--iters", type=int, default=20)
    parser.add_argument("--eps", type=float, default=1e-6)
    parser.add_argument(
        "--style-mode", choices=["per-row", "broadcast", "both"],
        default="both"
    )
    parser.add_argument("--output", default=None)
    parser.add_argument("--markdown", default=None)
    args = parser.parse_args()
    if not torch.cuda.is_available():
        raise SystemExit("CUDA is required")
    torch.manual_seed(53)
    ops = load_source_ops() if args.backend == "source" else load_installed_ops(args.artifact)
    results = []
    for name in SHAPE_GROUPS[args.shapes]:
        if args.style_mode in ("per-row", "both"):
            results.extend(run_one(ops, name, SHAPES[name], args, False))
        if args.style_mode in ("broadcast", "both"):
            results.extend(run_one(ops, name, SHAPES[name], args, True))
    for r in results:
        print(
            f"{r.status} {r.shape}/{r.kernel}: flashrt={r.flashrt_us:.3f}us "
            f"eager={r.torch_eager_us:.3f}us speedup={r.speedup_vs_eager:.2f}x p99={r.p99_abs:.6f}"
        )
    if args.output:
        Path(args.output).parent.mkdir(parents=True, exist_ok=True)
        Path(args.output).write_text(json.dumps([asdict(r) for r in results], indent=2) + "\n")
    if args.markdown:
        Path(args.markdown).parent.mkdir(parents=True, exist_ok=True)
        write_markdown(Path(args.markdown), results)


if __name__ == "__main__":
    main()