"""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 os
import sys
from collections.abc import Callable
from typing import Any

import torch

from deepinfer.model_executor.layers.tilelang_ops import ops


def _load_backend(name: str) -> dict[str, Any] | None:
    """Collect one backend, or None when it cannot be imported on this host.

    The default backend is TileLang, which needs `tile_kernels` and a recent `tilelang`. On a host
    that only carries Triton -- the whole point of the port, and the state of a vendor image --
    the comparison has nothing to compare against, so the run degrades to "build and save the
    selected backend's output" instead of failing to start.
    """
    try:
        return ops._collect(name)  # noqa: SLF001 - the collector is the one place that knows the layout
    except ImportError as exc:
        print(f"[ops] backend {name!r} is not importable here: {str(exc).splitlines()[0]}", file=sys.stderr, flush=True)
        return None


CUDA = _load_backend("cuda")
TRITON = _load_backend("triton")
BASELINE = CUDA is not None


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 source is None or 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).

    Without a baseline backend the first element is None and only the selected backend's output
    comes back, which the caller turns into a dump or a comparison against another host's dump.
    """
    if BASELINE:
        restore = swap(name, "cuda")
        try:
            expected = fn()
        finally:
            restore()
    else:
        expected = None

    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,
    agreement: float | None = None,
) -> bool:
    """Report whether one case agrees; returns True on success."""
    if (
        agreement is not None
        and isinstance(expected, torch.Tensor)
        and expected.dtype
        in (
            torch.int32,
            torch.int64,
        )
    ):
        # Random-draw outputs: require a high agreement rate rather than equality, and say what the
        # rate was so a regression that halves it cannot hide behind the threshold.
        same = int((expected == actual).sum())
        total = expected.numel()
        rate = same / max(total, 1)
        if rate >= agreement:
            if rate < 1.0:
                print(f"  note {name}: {total - same}/{total} differ (random draw on a boundary)")
            return True
        print(f"  FAIL {name}: agreement {rate:.4f} < {agreement} ({total - same}/{total})")
        return False

    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, agreement=agreement)
        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 _dump_file(root: str, kernel: str, case: str) -> str:
    return os.path.join(root, f"{kernel}__{case}.pt")


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")
    parser.add_argument(
        "--dump",
        metavar="DIR",
        help="save the selected backend's output per case instead of comparing to the default backend. "
        "Use it on a host where the default backend is not installed, then pass the directory to "
        "--compare on another host to check that two Triton builds agree.",
    )
    parser.add_argument("--compare", metavar="DIR", help="compare against dumps written by --dump on another host")
    args = parser.parse_args()
    if args.dump and args.compare:
        raise SystemExit("--dump and --compare are mutually exclusive")

    if args.compare:
        os.makedirs(args.compare, exist_ok=True)
    elif args.dump:
        os.makedirs(args.dump, exist_ok=True)
    elif not BASELINE:
        raise SystemExit(
            "no baseline backend here, so there is nothing to compare against: pass --dump DIR "
            "(save the outputs) or --compare DIR (check against another host's dumps)"
        )

    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)}")

    # Every `*_tl` symbol the selected backend exports must be callable. A factory that falls off
    # the end of its body returns None, and nothing notices until the engine calls it -- a stray
    # edit once dropped a `return launch` and the first symptom was a TypeError eight ranks deep
    # in an end-to-end run. Checking here turns that into one line before any case runs.
    broken = [name for name in dir(ops) if name.endswith("_tl") and not callable(getattr(ops, name))]
    if broken:
        raise SystemExit(f"backend {ops._BACKEND!r} exports non-callable kernel factories: {', '.join(broken)}")

    failures = 0
    skipped = 0
    saved = 0
    for kernel in selected:
        print(f"[{kernel}]")
        try:
            built = registry[kernel]()
        except ImportError as exc:
            # A case builder that needs `tile_kernels` cannot run on a Triton-only host; say so
            # rather than counting it as a failure, so the report stays honest about coverage.
            print(f"  skip (needs a package this host lacks): {str(exc).splitlines()[0]}")
            skipped += 1
            continue
        for case in built:
            # A case is `(name, exact, fn)`, optionally with a fourth element: the fraction of
            # entries allowed to differ. Only for outputs whose answer is a random draw, where the
            # last bit of a cumulative mass decides which side of a boundary the token falls on;
            # everything else is compared exactly or within a dtype tolerance.
            case_name, exact, fn = case[:3]
            agreement = case[3] if len(case) > 3 else None
            try:
                expected, actual = run_both(kernel, fn)
                if args.dump:
                    torch.save(actual, _dump_file(args.dump, kernel, case_name))
                    saved += 1
                    continue
                if args.compare:
                    path = _dump_file(args.compare, kernel, case_name)
                    if not os.path.exists(path):
                        print(f"  FAIL {case_name}: no dump at {path}")
                        failures += 1
                        continue
                    expected = torch.load(path, weights_only=False)
            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, agreement=agreement):
                failures += 1
            elif args.verbose:
                print(f"  ok   {case_name}")

    print()
    if args.dump:
        print(f"AB_DUMPED {saved} case(s) to {args.dump}" + (f", {skipped} kernel(s) skipped" if skipped else ""))
        return
    if failures:
        print(f"AB_FAILED {failures}" + (f" ({skipped} kernel(s) skipped)" if skipped else ""))
        sys.exit(1)
    print("AB_OK" + (f" ({skipped} kernel(s) skipped)" if skipped else ""))


if __name__ == "__main__":
    main()