"""Compare a Triton kernel against the TileLang one it replaces, through the shared interface.

Every interface module (``gather.py``, ``norm.py``, ...) calls ``ops.<name>_tl`` and then uses the
result. That means the strongest available check is not "does Triton match a Python reference"
but "does the same public call give the same answer when the kernel behind it changes". This
script does exactly that: it swaps ``ops.<name>_tl`` between the two backends and compares the
output of the public function.

    python scripts/h20_ops_ab.py            # every case
    python scripts/h20_ops_ab.py gather norm
"""

from __future__ import annotations

import argparse
import sys
from collections.abc import Callable
from typing import Any

import torch

from deepinfer.model_executor.layers.tilelang_ops import ops

# Each backend is a package whose submodules export `*_tl`; the packages' own `__init__` does not
# re-export, so collect the factories the same way the dispatcher does.
CUDA = ops._collect("cuda")  # noqa: SLF001 - the collector is the one place that knows the layout
TRITON = ops._collect("triton")  # noqa: SLF001


def swap(name: str, backend: str) -> Callable[[], None]:
    """Point ``ops.<name>`` at one backend, returning a function that restores it."""
    previous = getattr(ops, name)
    source = CUDA if backend == "cuda" else TRITON
    if name not in source:
        raise AttributeError(f"backend {backend!r} does not export {name!r}")
    setattr(ops, name, source[name])

    def restore() -> None:
        setattr(ops, name, previous)

    return restore


def run_both(name: str, fn: Callable[[], Any]) -> tuple[Any, Any]:
    """Run ``fn`` once against each backend, returning (cuda result, triton result)."""
    restore = swap(name, "cuda")
    try:
        expected = fn()
    finally:
        restore()

    restore = swap(name, "triton")
    try:
        actual = fn()
    finally:
        restore()
    return expected, actual


# Tolerance follows the dtype's own precision. Anything narrower than float32 is compared at its
# representable resolution, because a different reduction order in a different compiler moves the
# last bit and nothing more: for `norm` that is 10 of 655360 elements differing by one bf16 ulp,
# while the fp8 bytes and every scale factor come out identical.
_DTYPE_TOL: dict[torch.dtype, tuple[float, float]] = {
    torch.bfloat16: (1e-2, 1e-2),
    torch.float16: (1e-3, 1e-3),
    torch.float32: (1e-5, 1e-6),
}


def compare(name: str, expected: Any, actual: Any, *, exact: bool, rtol: float | None = None, atol: float | None = None) -> bool:
    """Report whether one case agrees; returns True on success."""
    if isinstance(expected, torch.Tensor):
        if expected.dtype != actual.dtype:
            print(f"  FAIL {name}: dtype {expected.dtype} vs {actual.dtype}")
            return False
        if expected.shape != actual.shape:
            print(f"  FAIL {name}: shape {tuple(expected.shape)} vs {tuple(actual.shape)}")
            return False
        if exact or not expected.is_floating_point():
            same = torch.equal(expected, actual)
            if not same:
                bad = expected != actual
                print(f"  FAIL {name}: {int(bad.sum())}/{expected.numel()} elements differ")
                idx = bad.nonzero()[0].tolist()
                print(f"        first at {idx}: {expected[tuple(idx)].item()} vs {actual[tuple(idx)].item()}")
            return bool(same)
        if expected.dtype in (torch.float8_e4m3fn, torch.float8_e5m2):
            # Compare at the format's own resolution: a one-ulp difference is a rounding tie that
            # moved, not a wrong value, so byte distance is the honest measure. Require the scales
            # to be exact elsewhere and the bytes to agree almost everywhere.
            delta = (expected.view(torch.uint8).int() - actual.view(torch.uint8).int()).abs()
            frac = float((delta > 0).float().mean())
            if int(delta.max()) <= 2 and frac <= 1e-3:
                return True
            print(f"  FAIL {name}: {frac * 100:.4f}% bytes differ, max byte distance {int(delta.max())}")
            return False

        if rtol is None or atol is None:
            default_rtol, default_atol = _DTYPE_TOL.get(expected.dtype, (1e-5, 1e-6))
            rtol = default_rtol if rtol is None else rtol
            atol = default_atol if atol is None else atol
        close = torch.allclose(expected.float(), actual.float(), rtol=rtol, atol=atol)
        if not close:
            diff = (expected.float() - actual.float()).abs()
            scale = expected.float().abs().clamp_min(1e-12)
            print(f"  FAIL {name}: max abs {diff.max().item():.3e}, max rel {(diff / scale).max().item():.3e}")
        return bool(close)

    if isinstance(expected, (tuple, list)):
        # A case that returns several tensors is only right if every element is; recurse so each
        # one gets its own report instead of collapsing into one boolean.
        if len(expected) != len(actual):
            print(f"  FAIL {name}: {len(expected)} outputs vs {len(actual)}")
            return False
        ok = True
        for i, (e, a) in enumerate(zip(expected, actual, strict=True)):
            ok &= compare(f"{name}[{i}]", e, a, exact=exact, rtol=rtol, atol=atol)
        return ok

    same = expected == actual
    if not same:
        print(f"  FAIL {name}: {expected!r} vs {actual!r}")
    return bool(same)


def cases() -> dict[str, Callable[[], list[tuple[str, bool, Callable[[], Any]]]]]:
    """Registry of per-kernel case builders, filled in as kernels are ported."""
    from scripts.ab_cases import CASES

    return CASES


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("kernels", nargs="*", help="subset of kernel names; default is all ported")
    parser.add_argument("--verbose", action="store_true")
    args = parser.parse_args()

    registry = cases()
    selected = args.kernels or sorted(registry)
    unknown = [name for name in selected if name not in registry]
    if unknown:
        raise SystemExit(f"no A/B cases for: {', '.join(unknown)}")

    failures = 0
    for kernel in selected:
        print(f"[{kernel}]")
        for case_name, exact, fn in registry[kernel]():
            try:
                expected, actual = run_both(kernel, fn)
            except Exception as exc:  # noqa: BLE001 - a raise on one backend is itself a result
                lines = str(exc).splitlines()
                print(f"  FAIL {case_name}: {type(exc).__name__}: {lines[0][:160] if lines else ''}")
                failures += 1
                continue
            if not compare(case_name, expected, actual, exact=exact):
                failures += 1
            elif args.verbose:
                print(f"  ok   {case_name}")

    print()
    if failures:
        print(f"AB_FAILED {failures}")
        sys.exit(1)
    print("AB_OK")


if __name__ == "__main__":
    main()