"""SM90 sanity checks: DeepGEMM FP8 1D2D (fp32 scales) + the H20 requant helpers."""

import deep_gemm
import torch

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

torch.manual_seed(0)
dev = "cuda"
print("cap", torch.cuda.get_device_capability(0))

# ---------------- 1. DeepGEMM fp8/fp8 1D2D GEMM with fp32 per-128 scales on SM90 ----
M, N, K = 256, 512, 1024
blk = 128
x = torch.randn(M, K, device=dev, dtype=torch.bfloat16)
w = torch.randn(N, K, device=dev, dtype=torch.bfloat16)


def q_token(t: torch.Tensor, blk: int) -> tuple[torch.Tensor, torch.Tensor]:
    """A operand: per-token 1x128 blocks -> sf (rows, cols/128) fp32."""
    r, c = t.shape
    tf = t.float()
    amax = tf.abs().view(r, c // blk, blk).amax(dim=2)
    s = (amax / 448.0).clamp_min(1e-30)
    q = (tf.view(r, c // blk, blk) / s[:, :, None]).to(torch.float8_e4m3fn)
    return q.view(r, c), s.contiguous()


def q_blocks(t: torch.Tensor, blk: int) -> tuple[torch.Tensor, torch.Tensor]:
    """B operand: 128x128 blocks -> sf (rows/128, cols/128) fp32."""
    r, c = t.shape
    tf = t.float()
    amax = tf.abs().view(r // blk, blk, c // blk, blk).amax(dim=(1, 3))
    s = (amax / 448.0).clamp_min(1e-30)
    q = (tf.view(r // blk, blk, c // blk, blk) / s[:, None, :, None]).to(torch.float8_e4m3fn)
    return q.view(r, c), s.contiguous()


xq, xsf = q_token(x, blk)
wq, wsf = q_blocks(w, blk)
out = torch.empty(M, N, device=dev, dtype=torch.bfloat16)
deep_gemm.fp8_fp4_gemm_nt((xq, xsf), (wq, wsf), out, None, disable_ue8m0_cast=True)
print("recipe default (1,128,128) OK")
ref = x @ w.t()
err = (out.float() - ref).abs().max().item()
rel = err / ref.abs().max().item()
print(f"gemm 1D2D fp32-sf on SM90: max_abs_err={err:.4f} rel={rel:.5f}")

# ---------------- 2. per-token activation quant via the runtime recipe ----------
params = {"fmt": "e4m3", "num_per_channels": 128, "round_sf": True, "use_packed_ue8m0": False, "use_tma_aligned_col_major_sf": True}
import tile_kernels

a8, asf = tile_kernels.quant.per_token_cast(x, **params)
print("per_token_cast ->", a8.shape, a8.dtype, asf.shape, asf.dtype)

# ---------------- 3. expert FP4 -> FP8 requant round-trip ----------------------
n_out, k = 256, 512
vals = 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], device=dev)
codes = torch.randint(0, 16, (n_out, k), device=dev)
ref_x = vals[codes]
packed = ((codes[..., 1::2] << 4) | codes[..., 0::2]).to(torch.uint8).view(torch.int8).contiguous()
assert torch.allclose(_unpack_fp4(packed), ref_x, atol=0, rtol=0), "fp4 unpack mismatch"

exp = torch.randint(118, 132, (n_out, k // 32), device=dev)
sf_e8m0 = (exp.to(torch.int32)).to(torch.uint8).view(torch.float8_e8m0fnu)
x_true = ref_x * ue8m0_to_fp32(sf_e8m0).repeat_interleave(32, dim=1)

w8, sf = requant_expert_fp4(packed, ue8m0_to_fp32(sf_e8m0), out_block=128)
xr = w8.to(torch.float32) * sf.t().repeat_interleave(128, 0).repeat_interleave(128, 1)
print(
    "expert requant: fp8",
    tuple(w8.shape),
    "sf",
    tuple(sf.shape),
    "max_abs_err",
    (xr - x_true).abs().max().item(),
    "rel",
    ((xr - x_true).abs().max() / x_true.abs().max()).item(),
)

# ---------------- 4. dense 32->128 requant ------------------------------------
w_ref = torch.randn(256, 512, device=dev, dtype=torch.bfloat16) / 8
w8d, sfd = q_blocks(w_ref, 32)
sf32 = ue8m0_to_fp32(
    (torch.randint(118, 130, (256 // 32, 512 // 32), device=dev).to(torch.int32)).to(torch.uint8).view(torch.float8_e8m0fnu)
)
src = w8d.float() * sf32.repeat_interleave(32, 0).repeat_interleave(32, 1)
w8n, sfn = requant_dense_fp8(w8d, sf32, dst_block=128)
rec = w8n.float() * sfn.repeat_interleave(128, 0).repeat_interleave(128, 1)
print("dense 32->128 requant: sf", tuple(sfn.shape), "rel_err", ((rec - src).abs().max() / src.abs().max()).item())
print("ALL_CHECKS_DONE")