"""Compare the compat layer's two backends on identical inputs, one per process.

The per-kernel A/B scripts compare `tile_kernels` against `tile_kernels.torch` directly. This
compares `deepinfer.tk_compat` against itself in both modes, which is a different question: it
exercises the adapters -- the `num_per_channels` -> `block_size` translation, the out-buffer
copies, the in-place rotary, the reordered engram-gate arguments, the composed big_fuse -- on the
signatures deepinfer actually calls.

Two processes because the backend is chosen once, at import.

    DEEPINFER_KERNELS=cuda python compat_ab.py --dump /tmp/cuda.pt
    DEEPINFER_KERNELS=ref  python compat_ab.py --compare /tmp/cuda.pt

Every case seeds from scratch, so both runs build byte-identical inputs.
"""

import argparse
import sys
import traceback
from collections.abc import Callable
from typing import TypeVar

import torch

from deepinfer.tk_compat import config, engram, mhc, moe, quant, transform

DEV = "cuda"
F = TypeVar("F", bound=Callable[[], dict])
CASES: dict[str, Callable[[], dict]] = {}


def case(name: str) -> Callable[[F], F]:
    """Register a function returning a dict of tensors for one operator."""

    def wrap(fn: F) -> F:
        CASES[name] = fn
        return fn

    return wrap


def _inp(*shape: int, dtype: torch.dtype = torch.float32, scale: float = 1.0) -> torch.Tensor:
    """Deterministic input: the seed is set per case, so the order of calls does not matter."""
    return (torch.rand(shape, device=DEV, dtype=torch.float32) - 0.5) * 2 * scale


# ------------------------------------------------------------------------------ quant


@case("per_token_cast")
def _c_per_token_cast() -> dict:
    torch.manual_seed(0)
    x = _inp(128, 2048, dtype=torch.bfloat16).to(torch.bfloat16)
    q, sf = quant.per_token_cast(x, fmt="e4m3", num_per_channels=128, round_sf=True, use_tma_aligned_col_major_sf=True)
    return {"q": q.view(torch.uint8), "sf": sf}


@case("per_token_cast_out")
def _c_per_token_cast_out() -> dict:
    torch.manual_seed(1)
    x = _inp(64, 512, dtype=torch.bfloat16).to(torch.bfloat16)
    kwargs = dict(fmt="e4m3", num_per_channels=128, round_sf=True, use_tma_aligned_col_major_sf=False)
    # Take the layout the interface itself allocates: the compiled kernel asserts specific strides
    # on the sf buffer, and `zeros_like` does not reproduce them (it fails with "strides[1]
    # expected 64, got 1"). Cloning the allocating result's storage gives a valid out buffer.
    ref_q, ref_sf = quant.per_token_cast(x, **kwargs)
    out_q = ref_q.clone()
    out_sf = ref_sf.clone()
    out_q.zero_()
    out_sf.zero_()
    quant.per_token_cast(x, out=(out_q, out_sf), **kwargs)
    return {"q": out_q.view(torch.uint8), "sf": out_sf}


@case("cast_back")
def _c_cast_back() -> dict:
    torch.manual_seed(2)
    x = _inp(64, 512, dtype=torch.bfloat16).to(torch.bfloat16)
    q = quant.per_token_cast(x, fmt="e4m3", num_per_channels=128, round_sf=True, use_tma_aligned_col_major_sf=True)
    return {"deq": quant.cast_back(tuple(q), "bf16", (1, 128))}


@case("swiglu_cast")
def _c_swiglu() -> dict:
    torch.manual_seed(3)
    x = _inp(64, 1024, dtype=torch.bfloat16).to(torch.bfloat16)
    q, sf = quant.swiglu_forward_and_per_token_cast(x, fmt="e4m3", num_per_channels=128, round_sf=True, use_tma_aligned_col_major_sf=True)
    return {"q": q.view(torch.uint8), "sf": sf}


# -------------------------------------------------------------------------------- moe


@case("moe_gate")
def _c_moe_gate() -> dict:
    torch.manual_seed(4)
    logits = _inp(128, 384, scale=2.0)
    bias = _inp(384, scale=0.1)
    idx, w = moe.moe_topk_gate(
        logits,
        num_topk=6,
        use_shared_as_routed=False,
        num_shared_experts=0,
        routed_scaling_factor=1.5,
        ep_rank=0,
        scoring_func="sqrtsoftplus",
        bias=bias,
    )
    return {"idx": idx, "w": w}


@case("moe_gate_mask")
def _c_moe_gate_mask() -> dict:
    torch.manual_seed(5)
    logits = _inp(64, 384, scale=2.0)
    bias = _inp(384, scale=0.1)
    mask = torch.ones((64,), device=DEV, dtype=torch.bool)
    mask[::5] = False
    idx, w = moe.moe_topk_gate(
        logits,
        num_topk=6,
        use_shared_as_routed=False,
        num_shared_experts=0,
        routed_scaling_factor=1.5,
        ep_rank=0,
        mask=mask,
        bias=bias,
    )
    return {"idx": idx, "w": w}


# --------------------------------------------------------------------------- transform


@case("apply_rotary")
def _c_rotary() -> dict:
    torch.manual_seed(6)
    tokens, heads, rot_dim = 32, 8, 128
    q = _inp(tokens, heads, rot_dim, dtype=torch.bfloat16).to(torch.bfloat16)
    k = _inp(tokens, heads, rot_dim, dtype=torch.bfloat16).to(torch.bfloat16)
    positions = torch.arange(tokens, device=DEV, dtype=torch.int64)
    angles = _inp(tokens, rot_dim // 2) * 3.0
    cache = torch.cat([angles.cos(), angles.sin()], dim=-1).contiguous()
    transform.apply_rotary(q, cache, k, positions, interleaved=False)
    return {"q": q, "k": k}


# ----------------------------------------------------------------------------- engram


@case("engram_hash")
def _c_engram_hash() -> dict:
    torch.manual_seed(7)
    layers, ngram, tables, vocab = 2, 4, 3, 1000
    ids = torch.randint(0, vocab, (128, ngram), device=DEV, dtype=torch.int32)
    mult = torch.randint(1, 2**31, (layers, ngram), device=DEV, dtype=torch.int64)
    vocab_sizes = torch.full((layers, ngram - 1, tables), vocab, device=DEV, dtype=torch.int32)
    offsets = torch.zeros((layers, (ngram - 1) * tables), device=DEV, dtype=torch.int32)
    return {"hash": engram.engram_hash(ids, mult, vocab_sizes, offsets, None)}


@case("engram_gate")
def _c_engram_gate() -> dict:
    torch.manual_seed(8)
    tokens, hc_mult, hidden = 64, 4, 512
    hs = _inp(tokens, hc_mult, hidden, dtype=torch.bfloat16, scale=0.5).to(torch.bfloat16)
    kv = _inp(tokens, hc_mult + 1, hidden, dtype=torch.bfloat16, scale=0.5).to(torch.bfloat16)
    wh = (_inp(hc_mult, hidden) * 0.5 + 0.5).to(torch.bfloat16)
    we = (_inp(hc_mult, hidden) * 0.5 + 0.5).to(torch.bfloat16)
    fused = engram.fused_weight(wh, we)
    out, *_ = engram.engram_gate_fwd(hs, kv, fused, 1e-6, 10.0, save_for_backward=False, image_token_mask=None, inplace=False)
    return {"out": out}


# -------------------------------------------------------------------------------- mhc


@case("mhc_expand")
def _c_expand() -> dict:
    torch.manual_seed(9)
    x = _inp(64, 128, dtype=torch.bfloat16).to(torch.bfloat16)
    o = torch.empty((64, 4, 128), device=DEV, dtype=torch.bfloat16)
    mhc.expand_to_mhc_fwd(x, o)
    return {"o": o}


@case("mhc_normw_merge")
def _c_normw() -> dict:
    torch.manual_seed(10)
    fn = _inp(12, 256)
    normw = _inp(256) * 0.5 + 0.5
    out = torch.empty_like(fn)
    mhc.mhc_fn_normw_merge_fwd(fn, normw, out)
    return {"out": out}


@case("mhc_head_mix")
def _c_head_mix() -> dict:
    torch.manual_seed(11)
    mix = _inp(64, 12)
    scale = _inp(1) * 2
    base = _inp(12)
    out = torch.empty_like(mix)
    mhc.mhc_head_compute_mix_fwd(mix, scale, base, out, 1e-6)
    return {"out": out}


@case("mhc_post")
def _c_post() -> dict:
    torch.manual_seed(12)
    seqs, tokens, mult, hidden = 2, 16, 4, 128
    x = _inp(seqs, tokens, hidden, dtype=torch.bfloat16).to(torch.bfloat16)
    residual = _inp(seqs, tokens, mult, hidden, dtype=torch.bfloat16).to(torch.bfloat16)
    post_mix = _inp(seqs, tokens, mult, 1) * 0.5 + 0.5
    comb = _inp(seqs, tokens, mult, mult) * 0.5 + 0.5
    return {"out": mhc.mhc_post_fwd(x, residual, post_mix, comb)}


@case("mhc_apply_mix")
def _c_apply_mix() -> dict:
    torch.manual_seed(13)
    x = _inp(64, 4, 128, dtype=torch.bfloat16).to(torch.bfloat16)
    mix = _inp(64, 4) * 0.5 + 0.5
    o = torch.empty((64, 128), device=DEV, dtype=torch.bfloat16)
    mhc.mhc_pre_apply_mix_fwd(x, mix, o)
    return {"o": o}


@case("mhc_split_mixes")
def _c_split_mixes() -> dict:
    torch.manual_seed(14)
    tokens, mult = 64, 4
    mult3 = mult * 2 + mult * mult
    mixes = _inp(tokens, mult3)
    scale = _inp(3) * 2
    base = _inp(mult3)
    pre = torch.empty((tokens, mult), device=DEV, dtype=torch.float32)
    post = torch.empty((tokens, mult), device=DEV, dtype=torch.float32)
    comb = torch.empty((tokens, mult * mult), device=DEV, dtype=torch.float32)
    mhc.mhc_pre_split_mixes_fwd(mixes, scale, base, pre, post, comb, 0.5, 1e-6)
    return {"pre": pre, "post": post, "comb": comb}


@case("mhc_reduce_rmsnorm")
def _c_reduce() -> dict:
    torch.manual_seed(15)
    tokens, mult3, groups, n_splits = 32, 12, 512, 16
    mul = _inp(n_splits, tokens, 1, mult3)
    sqr = _inp(n_splits, tokens, 1) * 0.5 + 1.0
    o_mul = torch.empty_like(mul[0])
    o_sqr = torch.empty_like(sqr[0])
    out = torch.empty((tokens, mult3), device=DEV, dtype=torch.float32)
    mhc.mhc_reduce_partials_and_rmsnorm_fwd(mul, sqr, o_mul, o_sqr, out, rms_group_size=groups, rms_eps=1e-6, n_splits=n_splits)
    return {"o_mul": o_mul, "o_sqr": o_sqr, "out": out}


@case("mhc_sinkhorn")
def _c_sinkhorn() -> dict:
    torch.manual_seed(16)
    x = _inp(64, 4, 4)
    out = torch.empty_like(x)
    mhc.mhc_sinkhorn_fwd(x, out, 10, 1e-6)
    return {"out": out}


@case("mhc_big_fuse")
def _c_big_fuse() -> dict:
    torch.manual_seed(17)
    tokens, mult, hidden, n_splits = 32, 4, 128, 16
    mult3 = mult * 2 + mult * mult
    residual = _inp(tokens, mult, hidden, dtype=torch.bfloat16).to(torch.bfloat16)
    mul = _inp(n_splits, tokens, mult3)
    sqr = _inp(n_splits, tokens) * 0.5 + float(hidden)
    scale = _inp(3) * 2
    base = _inp(mult3)
    post_mix = torch.empty((tokens, mult), device=DEV, dtype=torch.float32)
    comb = torch.empty((tokens, mult * mult), device=DEV, dtype=torch.float32)
    layer_in = torch.empty((tokens, hidden), device=DEV, dtype=torch.bfloat16)
    mhc.mhc_pre_big_fuse_fwd(
        mul,
        sqr,
        scale,
        base,
        residual,
        post_mix,
        comb,
        layer_in,
        1e-6,
        1e-6,
        1e-6,
        0.5,
        10,
        n_splits=n_splits,
    )
    return {"post": post_mix, "comb": comb, "layer_in": layer_in}


# -------------------------------------------------------------------------------- main


def _tolerance(kernel: str, key: str, ref: torch.Tensor) -> tuple[float, str]:
    """How close the two backends must be, and why.

    Integer outputs and the quantised bytes must match exactly: they are indices and encoded
    values, so "close" would be a different token. Floats get a tolerance scaled by the tensor's
    own magnitude, because several of these tensors sit near zero by construction.
    """
    if ref.dtype in (torch.int32, torch.int64, torch.uint8, torch.bool) or ref.dtype.is_complex:
        return 0.0, "exact"
    if ref.dtype == torch.bfloat16:
        return 2.0**-8 * ref.float().abs().max().clamp_min(1e-6).item() * 2, "2 bf16 ULP at magnitude"
    return ref.float().abs().max().clamp_min(1e-6).item() * 1e-5, "1e-5 of magnitude"


def run(compare_path: str | None, dump_path: str | None) -> int:
    print(f"backend = {config.backend()}")
    results = {}
    failed = []
    for name, fn in CASES.items():
        try:
            results[name] = fn()
        except Exception:
            print(f"  ERROR {name}")
            traceback.print_exc()
            failed.append(name)

    if dump_path:
        torch.save(results, dump_path)
        print(f"dumped {len(results)} case(s) to {dump_path}")
        return 0 if not failed else 1

    expected = torch.load(compare_path, weights_only=False)
    missing = sorted(set(expected) - set(results))
    if missing:
        print(f"cases present in the dump but not produced here: {missing}")
    for name in CASES:
        if name not in expected or name not in results:
            continue
        exp, got = expected[name], results[name]
        for key in exp:
            a, b = exp[key], got[key]
            if a.shape != b.shape:
                print(f"  FAIL {name}.{key}: shape {tuple(a.shape)} vs {tuple(b.shape)}")
                failed.append(f"{name}.{key}")
                continue
            if a.dtype != b.dtype:
                print(f"  FAIL {name}.{key}: dtype {a.dtype} vs {b.dtype}")
                failed.append(f"{name}.{key}")
                continue
            tol, why = _tolerance(name, key, a)
            diff = (a.float() - b.float()).abs().max().item()
            ok = diff <= tol
            print(f"  {'ok  ' if ok else 'FAIL'} {name}.{key}: max diff {diff:.3e} (tol {tol:.3e}, {why})")
            if not ok:
                failed.append(f"{name}.{key}")

    print()
    if failed:
        print(f"MISMATCH {len(failed)}: {sorted(set(failed))}")
        return 1
    print(f"ALL_MATCH ({len(results)} case(s))")
    return 0


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--dump", metavar="PATH")
    parser.add_argument("--compare", metavar="PATH")
    args = parser.parse_args()
    if bool(args.dump) == bool(args.compare):
        parser.error("pass exactly one of --dump or --compare")
    sys.exit(run(args.compare, args.dump))