"""Closed-loop KV cache round trip on the layouts SM90 actually uses.

The unit tests next to each op compare it against a reference implementation, and both sides
derive from the same `KVCacheFormat`, so they can agree while both being wrong about what the
engine stores. This closes the loop instead: write a known BF16 tensor through the paged cache,
read it back, and compare against the original. What survives is what decode sees.

Both layouts SM90 runs are covered:

* `V41_FORMAT` -- main/swa KV, e4m3 with per-32 ue8m0 scales, cast on store;
* `FP8_FLOAT32_SF_INDEX_KV_FORMAT` -- index KV, 128 fp8 + one fp32 scale = 132 B/token.

    python scripts/h20_kv_roundtrip.py
"""

import torch

from deepinfer import PLATFORM
from deepinfer.model_executor.layers.tilelang_ops.cache_common import (
    FP8_FLOAT32_SF_INDEX_KV_FORMAT,
    V41_FORMAT,
    KVCacheFormat,
)
from deepinfer.model_executor.layers.tilelang_ops.load_paged_kv_cache import load_paged_kv_cache
from deepinfer.model_executor.layers.tilelang_ops.store_kv_cache import store_kv_cache

DEV = "cuda"


def alloc_cache(layout: KVCacheFormat, num_blocks: int, block_size: int) -> torch.Tensor:
    """Exactly what the engine allocates: a uint8 blob with `bytes_per_token` per token."""
    return torch.zeros((num_blocks, block_size, 1, layout.bytes_per_token), dtype=torch.uint8, device=DEV)


def roundtrip(layout: KVCacheFormat, block_size: int, num_tokens: int, seed: int = 0) -> None:
    torch.manual_seed(seed)
    num_blocks = (num_tokens + block_size - 1) // block_size

    # One token per slot, sequential; that is the shape prefill produces.
    kv = torch.randn(num_tokens, layout.num_elements, dtype=torch.bfloat16, device=DEV)
    slots = torch.arange(num_tokens, dtype=torch.int32, device=DEV)
    cache = alloc_cache(layout, num_blocks, block_size)

    store_kv_cache(kv, cache, slots, layout=layout)

    # One segment per source block, covering the tokens actually written.
    segments = []
    dst = 0
    for block in range(num_blocks):
        rows = min(block_size, num_tokens - block * block_size)
        segments.append([block, 0, rows, dst])
        dst += rows
    segments = torch.tensor(segments, dtype=torch.int32, device=DEV)

    out = load_paged_kv_cache(cache, segments, num_tokens, layout=layout)
    if out.bf16_data is not None:
        # `do_cast` layouts are dequantized by the kernel itself.
        back = out.bf16_data.float()
    else:
        # Index layouts return quantized values plus scales, in their natural dtype -- the
        # consumer is a kernel that wants raw bytes, not BF16. Dequantize here or the comparison
        # is raw fp8 against BF16 (which reads as a ~127x magnitude error, direction intact).
        assert out.sf_data is not None, "index layout must return scales"
        group = layout.elements_per_sf
        sf = out.sf_data.float()
        if sf.ndim == 1:
            sf = sf.unsqueeze(-1)
        scale = sf.repeat_interleave(group, dim=-1)[:, : layout.num_elements]
        back = out.quantized_data.float() * scale

    ref = kv.float()
    rel = (back - ref).norm() / ref.norm()
    cos = torch.nn.functional.cosine_similarity(back.flatten(), ref.flatten(), dim=0)

    # e4m3 keeps 3 mantissa bits, so ~3% is the floor for any value that goes through it.
    budget = 0.06
    ok = "OK" if rel < budget else "FAIL"
    print(
        f"{layout.note:34} block={block_size:<4} tokens={num_tokens:<6} "
        f"bytes/token={layout.bytes_per_token:<4} rel={rel:8.4f} cos={cos:8.5f}  {ok}"
    )
    if rel >= budget:
        raise SystemExit(f"round trip lost more than fp8 should: {layout.note}")

    # Scale handling is the part most likely to be silently wrong: check that the per-token
    # magnitudes survive, not just the direction.
    per_token = (back - ref).norm(dim=1) / ref.norm(dim=1).clamp_min(1e-6)
    print(f"{'':34} worst token rel={per_token.max():.4f} median={per_token.median():.4f}")


def main() -> None:
    print(f"platform: sm90={PLATFORM.is_sm90} sf_dtype={PLATFORM.fp8_sf_type} pack={PLATFORM.sf_pack_factor}\n")
    for layout in (V41_FORMAT, FP8_FLOAT32_SF_INDEX_KV_FORMAT):
        for block_size in (32, 64):
            for num_tokens in (1, 233, 1029):
                roundtrip(layout, block_size, num_tokens)
        print()
    print("KV_ROUNDTRIP_DONE")


if __name__ == "__main__":
    main()