"""Probe which FP8 scale recipes DeepGEMM accepts on this GPU, via its own layout helper.

The SM90 path re-quantizes everything to fp32 per-128 because `get_gemm_recipe()` assumes Hopper
has nothing finer. The reference converter instead keeps an e8m0 scale per 32x32 block and folds
the relative offset into the FP8 values, which is lossless -- worth adopting only if the grouped
GEMM actually accepts it.

`transform_sf_into_required_layout` is the single place DeepGEMM validates a recipe (the engine
calls it for every MoE weight), so probing it answers the question directly and in milliseconds,
without materializing weights or launching a kernel.

    python scripts/h20_recipe_probe.py
"""

import deep_gemm
import torch

E8M0 = torch.float8_e8m0fnu


def make_sf(m: int, k: int, block_m: int, block_k: int, num_groups: int, e8m0: bool) -> torch.Tensor:
    """Scales in the plainest layout: one value per (block_m x block_k) tile, per group."""
    shape = (num_groups * (m // block_m), k // block_k)
    if e8m0:
        return torch.full(shape, 127, dtype=torch.uint8, device="cuda").view(E8M0)
    return torch.ones(shape, dtype=torch.float32, device="cuda")


def probe(recipe_b: tuple[int, int], block_a: int, e8m0: bool, n: int, k: int, num_groups: int) -> str:
    block_n, block_k = recipe_b
    sf = make_sf(n, k, block_n, block_k, num_groups, e8m0)
    try:
        out = deep_gemm.transform_sf_into_required_layout(
            sf,
            n,
            k,
            recipe=(1, block_n, block_k),
            num_groups=num_groups,
            is_sfa=False,
            disable_ue8m0_cast=e8m0,
        )
        got = tuple(out.shape)
        # The helper silently re-lays-out whatever it is given, so acceptance alone is not proof;
        # report the shape so an obviously wrong result is visible too.
        return f"ACCEPTED  out={got} dtype={out.dtype}"
    except Exception as exc:  # noqa: BLE001 - report whatever it raises
        first = (str(exc).strip().splitlines() or [type(exc).__name__])[0]
        return f"REJECTED  {first[:100]}"


def main() -> None:
    print(f"device: {torch.cuda.get_device_name(0)}")
    print(f"deep_gemm: {deep_gemm.__version__}\n")

    n, k, num_groups = 2304, 2560, 8

    print("B-operand recipes (what expert weights use, is_sfa=False):")
    for recipe, block_a, e8m0, label in [
        ((128, 128), 128, False, "fp32 per-128   (current SM90)"),
        ((128, 128), 128, True, "e8m0 per-128"),
        ((32, 32), 128, False, "fp32 per-32    (ref-converter grid)"),
        ((32, 32), 128, True, "e8m0 per-32    (ref-converter, lossless)"),
        ((1, 32), 128, True, "e8m0 per-row/32 (MXFP8)"),
        ((1, 128), 128, True, "e8m0 per-row/128"),
    ]:
        print(f"  {label:34} recipe={recipe} e8m0={e8m0:<5} -> {probe(recipe, block_a, e8m0, n, k, num_groups)}")

    print("\nA-operand recipes (activations, is_sfa=True):")
    for block_k, e8m0, label in [
        (128, False, "fp32 per-128   (current SM90)"),
        (32, True, "e8m0 per-32    (MXFP8)"),
        (32, False, "fp32 per-32"),
    ]:
        sf = make_sf(n, k, 1, block_k, 1, e8m0)
        try:
            out = deep_gemm.transform_sf_into_required_layout(
                sf, n, k, recipe=(1, 1, block_k), num_groups=1, is_sfa=True, disable_ue8m0_cast=e8m0
            )
            print(f"  {label:34} block_k={block_k:<4} e8m0={e8m0:<5} -> ACCEPTED  out={tuple(out.shape)}")
        except Exception as exc:  # noqa: BLE001
            first = (str(exc).strip().splitlines() or [type(exc).__name__])[0]
            print(f"  {label:34} block_k={block_k:<4} e8m0={e8m0:<5} -> REJECTED  {first[:100]}")

    print("\nPROBE_DONE")


if __name__ == "__main__":
    main()