"""Test: Verify NVFP4 CuTeDSL compilation with MmaMXF4NVF4Op (sf_vec_size=16).

This test does NOT run the kernel — it only verifies that the CuTeDSL JIT
compiler can handle the NVF4 block-scaled GEMM with proper pipeline abstractions.
If this compiles, we can add the custom epilogue.
"""

import torch
import cutlass
import cutlass.cute as cute
from cutlass.cute.nvgpu import cpasync, tcgen05
import cutlass.utils as utils
import cutlass.pipeline as pipeline
import cutlass.utils.blackwell_helpers as sm100_utils
import cutlass.utils.blockscaled_layout as blockscaled_utils
import cutlass.torch as cutlass_torch

from dsv4.ops.quantize import quantize_weight_to_nvfp4, quantize_activation_nvfp4
from dsv4.ops.layouts import make_b_k_major, assemble_raw_scales_2d3d_3d_side


def test_nvfp4_cutedsl_compilation():
    """Test that NVFP4 block-scaled GEMM compiles with CuTeDSL."""
    device = "cuda:0"
    M, N, K = 1, 384, 7168
    top_k = 6

    # Quantize
    gsa = 1.0 / (6.0 * 448.0)
    hs = torch.randn(M, K, dtype=torch.bfloat16, device=device)
    x_fp4, x_sf = quantize_activation_nvfp4(hs, gsa)

    W = torch.randn(K, N, dtype=torch.bfloat16, device=device)
    w_fp4, w_sf, w_gs = quantize_weight_to_nvfp4(W)
    stacked = torch.stack([w_fp4]).permute(0, 2, 1).contiguous()
    mat_b = make_b_k_major(stacked)
    scale_b = assemble_raw_scales_2d3d_3d_side([w_sf.T.contiguous()])

    print(f"x_fp4: {x_fp4.shape}, dtype={x_fp4.dtype}")
    print(f"x_sf: {x_sf.shape}, dtype={x_sf.dtype}")
    print(f"mat_b: {mat_b.shape}, dtype={mat_b.dtype}")
    print(f"scale_b: {scale_b.shape}, dtype={scale_b.dtype}")

    # Convert to CuTe tensors
    a_tensor = cutlass_torch.from_dlpack(x_fp4)
    a_tensor = a_tensor.mark_layout_dynamic(leading_dim=cutlass_torch.get_leading_dim(x_fp4))

    b_tensor = cutlass_torch.from_dlpack(mat_b)
    b_tensor = b_tensor.mark_layout_dynamic(leading_dim=cutlass_torch.get_leading_dim(mat_b))

    sfa_tensor = cutlass_torch.from_dlpack(x_sf)
    sfa_tensor = sfa_tensor.mark_layout_dynamic(leading_dim=cutlass_torch.get_leading_dim(x_sf))

    sfb_tensor = cutlass_torch.from_dlpack(scale_b)
    sfb_tensor = sfb_tensor.mark_layout_dynamic(leading_dim=cutlass_torch.get_leading_dim(scale_b))

    c_tensor = cutlass_torch.from_dlpack(
        torch.empty(M, N, dtype=torch.bfloat16, device=device))
    c_tensor = c_tensor.mark_layout_dynamic(leading_dim=cutlass_torch.get_leading_dim(
        torch.empty(M, N, dtype=torch.bfloat16, device=device)))

    print("CuTe tensors created OK")

    # ---- Setup exactly like dense.py ----
    sf_vec_size = 16  # NVF4
    a_dtype = cutlass.Float4E2M1FN
    b_dtype = cutlass.Float4E2M1FN
    sf_dtype = cutlass.Float8E4M3FN
    c_dtype = cutlass.BFloat16

    mma_tiler_mn = (128, 128)
    cluster_shape_mn = (1, 1)
    use_2cta = False
    cta_group = tcgen05.CtaGroup.ONE

    a_major = utils.LayoutEnum.from_tensor(a_tensor).mma_major_mode()
    b_major = utils.LayoutEnum.from_tensor(b_tensor).mma_major_mode()

    mma_inst_shape_mn_sfb = (
        mma_tiler_mn[0] // (2 if use_2cta else 1),
        cute.round_up(mma_tiler_mn[1], 128),
    )

    print(f"Creating tiled_mma with sf_vec_size={sf_vec_size}...", flush=True)
    tiled_mma = sm100_utils.make_blockscaled_trivial_tiled_mma(
        a_dtype, a_major, b_major, sf_dtype, sf_vec_size,
        cta_group, mma_tiler_mn)
    print(f"tiled_mma OK: shape_mnk={tiled_mma.shape_mnk}", flush=True)

    tiled_mma_sfb = sm100_utils.make_blockscaled_trivial_tiled_mma(
        a_dtype, a_major, b_major, sf_dtype, sf_vec_size,
        tcgen05.CtaGroup.ONE, mma_inst_shape_mn_sfb)
    print(f"tiled_mma_sfb OK", flush=True)

    # MMA tiler
    inst_shape_k = cute.size(tiled_mma.shape_mnk, mode=[2])
    inst_tile_k = 4
    k_tile = inst_shape_k * inst_tile_k
    mma_tiler = (cutlass.Int32(mma_tiler_mn[0]),
                 cutlass.Int32(mma_tiler_mn[1]),
                 cutlass.Int32(k_tile))

    cta_tile_shape_mnk = (
        mma_tiler[0] // cute.size(tiled_mma.thr_id.shape),
        mma_tiler[1],
        mma_tiler[2],
    )

    cluster_layout_vmnk = cute.tiled_divide(
        cute.make_layout((*cluster_shape_mn, 1)),
        (tiled_mma.thr_id.shape,))

    # SMEM layouts
    num_ab_stages = 2
    print("Creating SMEM layouts...", flush=True)
    a_smem_staged = sm100_utils.make_smem_layout_a(tiled_mma, mma_tiler, a_dtype, num_ab_stages)
    b_smem_staged = sm100_utils.make_smem_layout_b(tiled_mma, mma_tiler, b_dtype, num_ab_stages)
    sfa_smem_staged = blockscaled_utils.make_smem_layout_sfa(tiled_mma, mma_tiler, sf_vec_size, num_ab_stages)
    sfb_smem_staged = blockscaled_utils.make_smem_layout_sfb(tiled_mma, mma_tiler, sf_vec_size, num_ab_stages)
    print("SMEM layouts OK", flush=True)

    # TMA
    a_smem0 = cute.slice_(a_smem_staged, (None, None, None, 0))
    b_smem0 = cute.slice_(b_smem_staged, (None, None, None, 0))
    sfa_smem0 = cute.slice_(sfa_smem_staged, (None, None, None, 0))
    sfb_smem0 = cute.slice_(sfb_smem_staged, (None, None, None, 0))

    print("Creating TMA atoms...", flush=True)
    a_op = sm100_utils.cluster_shape_to_tma_atom_A(cluster_shape_mn, tiled_mma.thr_id)
    tma_a, gA = cute.nvgpu.make_tiled_tma_atom_A(a_op, a_tensor, a_smem0, mma_tiler, tiled_mma, cluster_layout_vmnk.shape)
    print("TMA A OK", flush=True)

    b_op = sm100_utils.cluster_shape_to_tma_atom_B(cluster_shape_mn, tiled_mma.thr_id)
    tma_b, gB = cute.nvgpu.make_tiled_tma_atom_B(b_op, b_tensor, b_smem0, mma_tiler, tiled_mma, cluster_layout_vmnk.shape)
    print("TMA B OK", flush=True)

    tma_sfa, gSFA = cute.nvgpu.make_tiled_tma_atom_A(
        a_op, sfa_tensor, sfa_smem0, mma_tiler, tiled_mma,
        cluster_layout_vmnk.shape, internal_type=cutlass.Int16)
    print("TMA SFA OK", flush=True)

    mma_tiler_sfb = (cutlass.Int32(mma_inst_shape_mn_sfb[0]),
                     cutlass.Int32(mma_inst_shape_mn_sfb[1]),
                     cutlass.Int32(k_tile))
    cluster_layout_sfb_vmnk = cute.tiled_divide(
        cute.make_layout((*cluster_shape_mn, 1)),
        (tiled_mma_sfb.thr_id.shape,))
    sfb_op = sm100_utils.cluster_shape_to_tma_atom_SFB(cluster_shape_mn, tiled_mma.thr_id)
    tma_sfb, gSFB = cute.nvgpu.make_tiled_tma_atom_B(
        sfb_op, sfb_tensor, sfb_smem0, mma_tiler_sfb, tiled_mma_sfb,
        cluster_layout_sfb_vmnk.shape, internal_type=cutlass.Int16)
    print("TMA SFB OK", flush=True)

    # Now try compiling the dense GEMM kernel (no custom epilogue)
    print("Compiling dense_blockscaled GEMM with NVF4...", flush=True)
    kernel = sm100_utils.Sm100BlockScaledPersistentDenseGemmKernel(
        a_tensor, b_tensor, c_tensor, sfa_tensor, sfb_tensor,
        acc_dtype=cutlass.Float32,
        mma_tiler_mn=mma_tiler_mn,
        cluster_shape_mn=cluster_shape_mn,
        sf_vec_size=sf_vec_size,
    )
    print("COMPILATION SUCCEEDED! NVF4 CuTeDSL path works.", flush=True)


if __name__ == "__main__":
    test_nvfp4_cutedsl_compilation()
