"""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"


# `inference/convert.py` in the checkpoint ships the authoritative table and unpack order; they
# are inlined here so this check does not depend on the model directory, and so a drift in either
# place shows up as a mismatch rather than as two copies of the same constant agreeing with itself.
REF_FP4_TABLE = (
    0.0,
    0.5,
    1.0,
    1.5,
    2.0,
    3.0,
    4.0,
    6.0,
    0.0,
    -0.5,
    -1.0,
    -1.5,
    -2.0,
    -3.0,
    -4.0,
    -6.0,
)


def check_e2m1_table() -> None:
    """Every 4-bit code must decode to the value the reference assigns it."""
    ours = tuple(_E2M1_VALUES)
    ok = ours == REF_FP4_TABLE
    print(f"E2M1 table: {'OK' if ok else 'MISMATCH'}")
    if not ok:
        for code, (a, b) in enumerate(zip(ours, REF_FP4_TABLE, strict=True)):
            if a != b:
                print(f"  code {code:04b}: reference={b:+.1f} ours={a:+.1f}")
        raise SystemExit("E2M1 table disagrees with inference/convert.py")


def check_nibble_order() -> None:
    """`_unpack_fp4` must interleave low nibble first, as the reference does."""
    # Two bytes of hand-picked codes; a swap would give a different sequence.
    raw = torch.tensor([0b0011_0001, 0b0111_0110], dtype=torch.int8, device=DEV)
    table = torch.tensor(REF_FP4_TABLE, dtype=torch.float32, device=DEV)

    # Reference: low nibble is the even element.
    x = raw.view(torch.uint8)
    low = (x & 0x0F).long()
    high = ((x >> 4) & 0x0F).long()
    ref = torch.stack([table[low], table[high]], dim=-1).flatten()

    ours = _unpack_fp4(raw).reshape(-1)
    ok = torch.equal(ref, ours)
    print(f"nibble order: {'OK' if ok else 'MISMATCH'}  values={ours.tolist()}")
    if not ok:
        swapped = ours.view(-1, 2).flip(-1).reshape(-1)
        print(f"  reference: {ref.tolist()}")
        print(f"  ours     : {ours.tolist()}")
        print(f"  swapped  : {swapped.tolist()}  -> swapping would {'fix it' if torch.equal(ref, swapped) else 'NOT fix it'}")
        raise SystemExit("nibble order disagrees with inference/convert.py")


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 _quantize_fp4_codes(x: torch.Tensor, table: torch.Tensor) -> torch.Tensor:
    """Nearest E2M1 code for every element, by table lookup.

    Spelled out instead of casting to `torch.float4_e2m1fn_x2` because that cast kernel exists
    neither on SM90 nor on the CPU, and the decode under test is pure integer/table work anyway.
    """
    # (..., 16) distances -> argmin gives the code index in REF_FP4_TABLE order.
    return (x.unsqueeze(-1) - table).abs().argmin(dim=-1).to(torch.uint8)


def check_requant_roundtrip() -> None:
    """End to end: quantize to FP4 / per-32 UE8M0, then 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
    table = torch.tensor(REF_FP4_TABLE, dtype=torch.float32, device=DEV)

    # 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.
    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)
    codes = _quantize_fp4_codes(groups / scale.unsqueeze(-1), table)
    # Pack low nibble first, matching `_unpack_fp4`.
    packed = (codes[..., 0::2] | (codes[..., 1::2] << 4)).to(torch.int8).reshape(n_out, k // 2)
    sf_bits = (exp.to(torch.int32) + 127).to(torch.uint8)
    sf = sf_bits.view(torch.float8_e8m0fnu)

    deq_ref = table[codes.long()].reshape(n_out, k) * 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()