"""Compare our FP4->FP8 expert requantization against DeepSeek's own reference conversion.

`inference/convert.py::cast_e2m1fn_to_e4m3fn` is the authoritative implementation: it is what
produced the TP8 checkpoint that the reference model actually runs on. Two things are worth
checking there, and both matter because a wrong answer scrambles every routed expert while
leaving the (already-FP8) shared expert intact -- exactly the asymmetry the layer comparison
showed.

1. Equality of the *decoded* weights. Both paths should represent the same real values; any
   difference is accuracy we are throwing away.
2. The scale grid. The reference keeps a per-32 e8m0 scale (relative offsets folded into the FP8
   values), while the SM90 path re-quantizes on a per-128 fp32 grid. That is a legitimate design
   difference only if the magnitude error stays small.

    python scripts/h20_expert_weight_check.py --layer 0 --expert 370
"""

import argparse
import json
from pathlib import Path

import torch
from safetensors import safe_open

from deepinfer.h20_ops.quant import requant_expert_fp4, ue8m0_to_fp32

# Copied verbatim from inference/convert.py so the comparison uses the reference algorithm
# rather than a paraphrase of it.
FP4_TABLE = torch.tensor(
    [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],
    dtype=torch.float32,
)


def cast_e2m1fn_to_e4m3fn(x: torch.Tensor, scale: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
    """DeepSeek's lossless FP4 -> FP8 cast (inference/convert.py)."""
    assert x.dtype == torch.int8 and x.ndim == 2
    out_dim, in_dim = x.size()
    in_dim *= 2
    fp8_block_size = 32
    fp4_block_size = 32
    assert in_dim % fp8_block_size == 0 and out_dim % fp8_block_size == 0
    assert scale.size(0) == out_dim and scale.size(1) == in_dim // fp4_block_size

    x = x.view(torch.uint8)
    low = x & 0x0F
    high = (x >> 4) & 0x0F
    x = torch.stack([FP4_TABLE[low.long()], FP4_TABLE[high.long()]], dim=-1).flatten(2)

    MAX_OFFSET_BITS = 6
    b_out = out_dim // fp8_block_size
    b_in = in_dim // fp8_block_size
    x = x.view(b_out, fp8_block_size, b_in, fp8_block_size).transpose(1, 2)
    scale = scale.float().view(b_out, fp8_block_size, b_in, -1).transpose(1, 2).flatten(2)
    scale_max_offset_bits = scale.amax(dim=-1, keepdim=True) / (2**MAX_OFFSET_BITS)
    offset = scale / scale_max_offset_bits
    offset = offset.unflatten(-1, (fp8_block_size, -1)).repeat_interleave(fp4_block_size, dim=-1)
    x = (x * offset).transpose(1, 2).reshape(out_dim, in_dim)
    return x.to(torch.float8_e4m3fn), scale_max_offset_bits.squeeze(-1).to(torch.float8_e8m0fnu)


def read_tensor(hf_dir: Path, index: dict, name: str) -> torch.Tensor:
    shard = index["weight_map"][name]
    with safe_open(str(hf_dir / shard), framework="pt", device="cpu") as f:
        return f.get_tensor(name)


def dequant_ref(fp8: torch.Tensor, sf_e8m0: torch.Tensor, block: int = 32) -> torch.Tensor:
    """Reference layout: (out, in) FP8, scale (out/32, in/32) e8m0 on a 32x32 grid."""
    s = ue8m0_to_fp32(sf_e8m0)
    s = s.repeat_interleave(block, 0).repeat_interleave(block, 1)[: fp8.shape[0], : fp8.shape[1]]
    return fp8.float() * s


def dequant_ours(fp8: torch.Tensor, sf_fp32: torch.Tensor, block: int = 128) -> torch.Tensor:
    """SM90 layout: (out, in) FP8, scales (in/block, out/block) fp32, contraction-major."""
    s = sf_fp32.t().repeat_interleave(block, 0).repeat_interleave(block, 1)[: fp8.shape[0], : fp8.shape[1]]
    return fp8.float() * s


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--hf-dir", default="/data/models/DeepSeek-V4.1-Flash-1017")
    parser.add_argument("--layer", type=int, default=0)
    parser.add_argument("--expert", type=int, default=370)
    args = parser.parse_args()

    hf_dir = Path(args.hf_dir)
    index = json.loads((hf_dir / "model.safetensors.index.json").read_text())
    prefix = f"layers.{args.layer}.ffn.experts.{args.expert}."

    for proj in ("w1", "w3", "w2"):
        packed = read_tensor(hf_dir, index, prefix + proj + ".weight")
        scale = read_tensor(hf_dir, index, prefix + proj + ".scale")
        print(f"\n{proj}: packed {tuple(packed.shape)} {packed.dtype}, scale {tuple(scale.shape)} {scale.dtype}")

        ref_fp8, ref_sf = cast_e2m1fn_to_e4m3fn(packed, scale)
        ours_fp8, ours_sf = requant_expert_fp4(packed, scale, out_block=128)

        # `requant_expert_fp4` returns (n_out, k) FP8 and a contraction-major scale; the reference
        # returns (out, in) too, so the value grids line up directly.
        ref_vals = dequant_ref(ref_fp8, ref_sf)
        ours_vals = dequant_ours(ours_fp8, ours_sf)

        assert ref_vals.shape == ours_vals.shape, f"{ref_vals.shape} vs {ours_vals.shape}"
        num = (ours_vals - ref_vals).norm().item()
        den = ref_vals.norm().item()
        cos = torch.nn.functional.cosine_similarity(ours_vals.reshape(-1), ref_vals.reshape(-1), dim=0).item()
        print(f"  ref scale: e8m0 {tuple(ref_sf.shape)}   ours: fp32 {tuple(ours_sf.shape)}")
        print(f"  |ref|={den:.4f}  |ours|={ours_vals.norm().item():.4f}  ratio={ours_vals.norm().item() / den:.5f}")
        print(f"  rel_err={num / den:.5f}  cos={cos:.7f}")

    print("\nEXPERT_WEIGHT_CHECK_DONE")


if __name__ == "__main__":
    main()