"""Numerically test the SM90 decode-time indexer score against a reference formula.

`_decode_p_sum_sm90` reimplements DeepGEMM's `fp8_fp4_paged_mqa_logits` in PyTorch, because
Hopper's paged kernel hard-requires 64-token pages while the index cache uses a different page
size. Nothing else exercises it: prefill goes through the real kernel, so a defect here would be
decode-only -- which is exactly where the remaining drift lives.

The reference below deliberately shares no code with the implementation: it reads the cache with
explicit per-token slicing in float64 rather than strided views, so an offset or stride mistake
in one does not cancel against the other.

    python scripts/h20_index_decode_check.py
"""

import torch

from deepinfer.model_executor.layers.tilelang_ops.cache_common import FP8_FLOAT32_SF_INDEX_KV_FORMAT
from deepinfer.model_executor.models.v4_modules.index_attn import MQA

DEV = "cuda"
HEAD_DIM = 128


def build_cache(num_blocks: int, page_tokens: int, scales: torch.Tensor, keys: torch.Tensor) -> torch.Tensor:
    """Cache in the store kernel's layout: block-wide values, then block-wide scales."""
    fmt = FP8_FLOAT32_SF_INDEX_KV_FORMAT
    raw = torch.zeros((num_blocks, page_tokens, 1, fmt.bytes_per_token), dtype=torch.uint8, device=DEV)
    flat = raw.view(-1)
    for block in range(num_blocks):
        base = block * fmt.bytes_per_token * page_tokens
        values = keys[block].to(torch.float8_e4m3fn).view(torch.uint8).reshape(-1)
        flat[base : base + page_tokens * HEAD_DIM] = values
        sf = scales[block].float().reshape(-1).view(torch.uint8)
        flat[base + page_tokens * HEAD_DIM : base + fmt.bytes_per_token * page_tokens] = sf
    return raw


def reference_p_sum(
    q: torch.Tensor,
    kv_cache: torch.Tensor,
    weights: torch.Tensor,
    klen: torch.Tensor,
    block_table: torch.Tensor,
    page_tokens: int,
) -> torch.Tensor:
    """float64 reference: gather by explicit offsets, dequantize, `sum_h relu(q.k) w`."""
    num_rows = q.shape[0]
    max_len = int(klen.max())
    out = torch.full((num_rows, max_len), float("-inf"), dtype=torch.float64, device=DEV)
    fmt_bytes = FP8_FLOAT32_SF_INDEX_KV_FORMAT.bytes_per_token
    for row in range(num_rows):
        ctx = int(klen[row])
        keys = []
        for pos in range(ctx):
            page = pos // page_tokens
            offset = pos % page_tokens
            block = int(block_table[row, page])
            base = block * fmt_bytes * page_tokens
            raw = kv_cache.view(-1)[base + offset * HEAD_DIM : base + (offset + 1) * HEAD_DIM]
            scale = kv_cache.view(-1)[base + page_tokens * HEAD_DIM + offset * 4 : base + page_tokens * HEAD_DIM + offset * 4 + 4]
            keys.append(raw.view(torch.float8_e4m3fn).double() * scale.view(torch.float32).double())
        key = torch.stack(keys)  # [ctx, head_dim]
        scores = q[row].double() @ key.t()
        out[row, :ctx] = (scores.relu() * weights[row].double().unsqueeze(-1)).sum(dim=0)
    return out


def run_case(page_tokens: int, num_blocks: int, num_rows: int, num_heads: int, lens: list[int], seed: int) -> None:
    torch.manual_seed(seed)
    total_tokens = num_blocks * page_tokens
    keys = torch.randn(num_blocks, page_tokens, HEAD_DIM, dtype=torch.float32, device=DEV) * 0.5
    # One fp32 scale per token, as `elements_per_sf == quantized_elements == 128` implies.
    scales = (keys.abs().amax(dim=-1) / 448.0).clamp_min(1e-8)
    cache = build_cache(num_blocks, page_tokens, scales, keys)

    q = torch.randn(num_rows, num_heads, HEAD_DIM, dtype=torch.float32, device=DEV) * 0.5
    q_sf = (q.abs().amax(dim=-1, keepdim=True) / 448.0).clamp_min(1e-8)  # one scale per (row, head)
    q_fp8 = (q / q_sf).to(torch.float8_e4m3fn)
    q_deq = q_fp8.float() * q_sf

    weights = torch.rand(num_rows, num_heads, dtype=torch.float32, device=DEV)
    klen = torch.tensor(lens, dtype=torch.int32, device=DEV)
    max_pages = (max(lens) + page_tokens - 1) // page_tokens
    block_table = torch.zeros(num_rows, max_pages, dtype=torch.int32, device=DEV)
    for row in range(num_rows):
        pages = (lens[row] + page_tokens - 1) // page_tokens
        block_table[row, :pages] = torch.randperm(num_blocks, device=DEV)[:pages]

    class _Stub:
        max_kv_len = total_tokens

    got = MQA._decode_p_sum_sm90(_Stub(), q_fp8, q_sf, cache, weights, klen, block_table)
    ref = reference_p_sum(q_deq, cache, weights, klen, block_table, page_tokens)

    got = got[:, : ref.shape[1]].double()
    # Positions beyond each row's context are `-inf` by design on both sides; compare only the
    # live window, which is what top-k actually sees.
    live = torch.arange(ref.shape[1], device=DEV).unsqueeze(0) < klen.double().unsqueeze(1)
    rel = (got[live] - ref[live]).norm() / ref[live].norm()
    cos = torch.nn.functional.cosine_similarity(got[live], ref[live], dim=0)
    worst = (got[live] - ref[live]).abs().max().item()

    ok = "OK" if rel < 1e-5 else "MISMATCH"
    print(
        f"page={page_tokens:<4} blocks={num_blocks:<3} rows={num_rows} heads={num_heads:<3} "
        f"lens={lens}  rel={rel:9.2e} cos={cos:8.5f} max|d|={worst:9.2e}  {ok}"
    )
    if rel >= 1e-5:
        # Localize: which row/position disagrees first.
        diff = (got - ref).abs().masked_fill(~live, -1)
        row, col = divmod(int(diff.argmax()), diff.shape[1])
        print(f"{'':4}worst at row {row} pos {col}: got={got[row, col].item():.6f} ref={ref[row, col].item():.6f}")
        raise SystemExit("decode indexer disagrees with the reference")


def main() -> None:
    print(f"device: {torch.cuda.get_device_name(0)}\n")
    for page_tokens in (32, 64):
        run_case(page_tokens, 4, 2, 8, [1, page_tokens * 2], seed=0)
        run_case(page_tokens, 4, 2, 8, [page_tokens - 1, page_tokens], seed=1)
        run_case(page_tokens, 4, 2, 8, [page_tokens * 3, page_tokens * 2 + 1], seed=2)
        run_case(page_tokens, 8, 3, 32, [7, 33, page_tokens * 5 - 1], seed=3)
    print("\nINDEX_DECODE_CHECK_DONE")


if __name__ == "__main__":
    main()