"""Check `h20_ops.quant` against PyTorch's native FP4/FP8 semantics.

The SM90 path re-quantizes packed E2M1 experts into FP8 while loading, and that code makes three
assumptions the checkpoint does not spell out: which nibble holds the even element, the E2M1 code
table, and the UE8M0 bias. If any of them is wrong every expert weight is scrambled, which is
exactly the symptom the layer comparison shows (shared expert matches, routed experts do not).

    python scripts/h20_quant_check.py
"""

import torch

from deepinfer.h20_ops.quant import _E2M1_VALUES, _unpack_fp4, requant_expert_fp4, ue8m0_to_fp32

# PyTorch's FP4 cast kernel does not exist for SM90 (device-side assert on H20), and the decode
# under test is pure integer/table work, so everything here runs on the CPU.
DEV = "cpu"


def check_e2m1_table() -> None:
    """Every 4-bit code must decode to the value PyTorch assigns it."""
    codes = torch.arange(16, dtype=torch.uint8, device=DEV)
    # PyTorch's own decoder: pack two codes per byte and view as the FP4 dtype.
    packed = torch.stack([codes, torch.zeros_like(codes)], dim=-1).reshape(-1).contiguous()
    ref = packed.view(torch.float4_e2m1fn_x2).to(torch.float32)[0::2]
    table = torch.tensor(_E2M1_VALUES, dtype=torch.float32, device=DEV)
    ok = torch.equal(ref, table)
    print(f"E2M1 table: {'OK' if ok else 'MISMATCH'}")
    if not ok:
        for code in range(16):
            if ref[code] != table[code]:
                print(f"  code {code:04b}: torch={ref[code].item():+.1f} ours={table[code].item():+.1f}")
        raise SystemExit("E2M1 table disagrees with PyTorch")


def check_nibble_order() -> None:
    """`_unpack_fp4` must interleave in the same order as the packed dtype."""
    # Distinct, exactly representable values so an order swap is unambiguous.
    values = torch.tensor([1.0, 2.0, 4.0, 6.0, 0.5, 1.5, 3.0, -1.0], dtype=torch.float32, device=DEV)
    packed = values.to(torch.float4_e2m1fn_x2).view(torch.uint8)  # (4,) bytes
    ref = packed.view(torch.float4_e2m1fn_x2).to(torch.float32)
    ours = _unpack_fp4(packed.view(torch.int8)).reshape(-1)
    ok = torch.equal(ref, ours)
    print(f"nibble order: {'OK' if ok else 'MISMATCH'}")
    if not ok:
        print(f"  torch: {ref.tolist()}")
        print(f"  ours : {ours.tolist()}")
        swapped = ours.view(-1, 2).flip(-1).reshape(-1)
        print(f"  ours (nibbles swapped): {swapped.tolist()}")
        print(f"  -> swapping would {'fix it' if torch.equal(ref, swapped) else 'NOT fix it'}")
        raise SystemExit("nibble order disagrees with PyTorch")


def check_ue8m0() -> None:
    """`ue8m0_to_fp32` must match the checkpoint's scale dtype."""
    bits = torch.tensor([127, 128, 126, 130, 124], dtype=torch.uint8, device=DEV)
    scales = bits.view(torch.float8_e8m0fnu)
    ours = ue8m0_to_fp32(scales)
    ref = torch.pow(2.0, bits.to(torch.float32) - 127)
    ok = torch.equal(ours, ref)
    print(f"UE8M0 bias: {'OK' if ok else 'MISMATCH'}  ours={ours.tolist()} ref={ref.tolist()}")
    if not ok:
        raise SystemExit("UE8M0 decode disagrees")


def check_requant_roundtrip() -> None:
    """End to end: dequantize FP4 on the CPU, requantize, and measure the error."""
    torch.manual_seed(0)
    n_out, k = 256, 512
    x = torch.randn(n_out, k, device=DEV, dtype=torch.float32) * 0.05

    # Quantize to FP4 with the checkpoint's per-32 UE8M0 grid, so the round trip starts from a
    # realistic representation rather than an exact float. FP8 casts also stay on the CPU here.
    groups = x.view(n_out, k // 32, 32)
    amax = groups.abs().amax(dim=-1).clamp_min(torch.finfo(torch.float32).tiny)
    exp = torch.ceil(torch.log2(amax / 6.0)).clamp(-127, 127)
    scale = torch.pow(2.0, exp)
    q4 = (groups / scale.unsqueeze(-1)).to(torch.float4_e2m1fn_x2)
    packed = q4.view(torch.int8)
    sf_bits = (exp.to(torch.int32) + 127).to(torch.uint8)
    sf = sf_bits.view(torch.float8_e8m0fnu)

    deq_ref = q4.to(torch.float32) * scale.repeat_interleave(32, dim=1)

    weight, sfb = requant_expert_fp4(packed, sf, out_block=128)
    assert weight.dtype == torch.float8_e4m3fn, weight.dtype
    assert sfb.dtype == torch.float32, sfb.dtype
    assert sfb.shape == (k // 128, n_out // 128), f"scale layout {tuple(sfb.shape)}"

    deq_ours = weight.to(torch.float32) * sfb.t().repeat_interleave(128, 0).repeat_interleave(128, 1)[:n_out, :k]
    err = (deq_ours - deq_ref).norm() / deq_ref.norm()
    print(f"requant roundtrip: rel_err={err.item():.4f} (fp8 e4m3 floor ~0.01, fp4 source itself is ~0.1)")
    assert err < 0.05, "requantization lost more precision than FP8 should"


def main() -> None:
    check_e2m1_table()
    check_nibble_order()
    check_ue8m0()
    check_requant_roundtrip()
    print("QUANT_CHECK_DONE")


if __name__ == "__main__":
    main()