"""Validate LocalEngramGather against a torch reference, single-process and multi-rank.

torchrun --nproc_per_node=8 scripts/h20_engram_check.py
"""

import torch
import torch.distributed as dist

from deepinfer.h20_ops.engram_local import LocalEngramGather

HIDDEN = 256
N_LOCAL = 4096  # rows per rank
ENTRIES = 24  # 3 n-gram tables x 8 heads
TOKENS = 128
LAYERS = 2


def quantize_table(rows: torch.Tensor, gen: torch.Generator) -> tuple[torch.Tensor, torch.Tensor]:
    """fp8 table + per-32 ue8m0 (stored as fp32) scales, like the checkpoint."""
    groups = rows.float().view(rows.shape[0], HIDDEN // 32, 32)
    amax = groups.abs().amax(-1).clamp_min(1e-30)
    exp = torch.ceil(torch.log2(amax / 448.0)).clamp_min(-126)
    scale = torch.pow(2.0, exp)
    q = (groups / scale.unsqueeze(-1)).to(torch.float8_e4m3fn).view(rows.shape[0], HIDDEN)
    return q, scale


def main() -> None:
    dist.init_process_group("nccl")
    rank, world = dist.get_rank(), dist.get_world_size()
    torch.cuda.set_device(rank)
    dev = torch.device("cuda", rank)
    gen = torch.Generator(device="cpu").manual_seed(0)

    torch.manual_seed(rank)
    tables, scales = [], []
    for _ in range(LAYERS):
        rows = torch.randn(N_LOCAL, HIDDEN, generator=gen)
        q, s = quantize_table(rows, gen)
        tables.append(q.to(dev))
        scales.append(s.to(dev))

    gather = LocalEngramGather(group=None, hidden=HIDDEN)
    # Use a real process group so all_to_all runs across ranks.
    gather.group = dist.group.WORLD
    gather.set_tables(tables, scales)
    assert gather.num_local_entries == [N_LOCAL] * LAYERS
    print(f"[rank {rank}] {gather.describe()}")

    total = world * N_LOCAL
    idx = torch.randint(0, total, (LAYERS, TOKENS, ENTRIES), generator=gen, dtype=torch.int32).to(dev)
    idx[0, 0, 0] = -1  # invalid slot must come back zeroed

    hooks = gather.fetch(idx)
    got = hooks[0]()

    # Reference: rebuild the global table from every rank.
    ref_rows = [torch.empty(world * N_LOCAL, HIDDEN, dtype=torch.float8_e4m3fn, device=dev) for _ in range(LAYERS)]
    ref_sf = [torch.empty(world * N_LOCAL, HIDDEN // 32, device=dev) for _ in range(LAYERS)]
    for layer in range(LAYERS):
        dist.all_gather_into_tensor(ref_rows[layer], tables[layer])
        dist.all_gather_into_tensor(ref_sf[layer], scales[layer])

    def reference(layer: int) -> torch.Tensor:
        flat = idx[layer].reshape(-1).long()
        valid = flat >= 0
        safe = flat.clamp_min(0)
        w = ref_rows[layer].index_select(0, safe).float()
        s = ref_sf[layer].index_select(0, safe)
        vals = w.reshape(-1, HIDDEN // 32, 32) * s.unsqueeze(-1)
        vals = vals.reshape(-1, HIDDEN).to(torch.bfloat16)
        vals[~valid] = 0
        return vals

    a, sf = got if isinstance(got, tuple) else (got, None)
    ref = reference(0)
    # `a` holds fp8 codes; the per-token 1x128 scales turn them back into values.
    got_vals = (a.float() * sf.float().repeat_interleave(HIDDEN * ENTRIES // sf.shape[1], dim=1)).view(TOKENS, ENTRIES, HIDDEN)
    # Compare in dequantized space: per-128 requantization is the only difference expected.
    ref_deq = ref.float().view(TOKENS, ENTRIES, HIDDEN)
    err = (got_vals - ref_deq).abs().max().item()
    rel = err / max(ref_deq.abs().max().item(), 1e-9)
    print(f"[rank {rank}] out={tuple(a.shape)} sf={None if sf is None else tuple(sf.shape)} max_abs={err:.4f} rel={rel:.4f}")
    print(f"[rank {rank}] invalid slot all-zero: {bool((got_vals[0, 0] == 0).all())}")

    # Non-owner traffic actually exercised?
    owners = torch.div(idx[0].reshape(-1).clamp_min(0), N_LOCAL, rounding_mode="floor")
    remote = int((owners != rank).sum())
    print(f"[rank {rank}] remote entries {remote}/{owners.numel()}")

    ok = rel < 0.05
    print(f"[rank {rank}] {'PASS' if ok else 'FAIL'}")
    dist.barrier()
    dist.destroy_process_group()
    if not ok:
        raise SystemExit(1)


if __name__ == "__main__":
    main()