"""
Stage A Minimal: Just tcgen05.mma with TMEM accumulator, epilogue writes to GMEM.
No TMA loads (use regular SMEM loads). No pipeline. Just validate the MMA path.
"""

import torch
import cutlass
import cutlass.cute as cute
import cutlass.utils as utils
import cutlass.pipeline as pipeline
from cutlass.cute.nvgpu import tcgen05
from cutlass import BFloat16, Float32
import cuda.bindings.driver as cuda

# Q: (128, 512) BF16, K^T: (64, 512) BF16 -> S: (128, 64) FP32
# These are the minimum valid tile sizes for tcgen05.mma
M = 128
N = 64
K = 512


class MinimalMMAKernel:
    def __init__(self):
        self.cta_group = tcgen05.CtaGroup.ONE
        self.mma_tiler_mn = (M, N)
        self.threads_per_cta = 192  # 6 warps
        self.epilog_warp_ids = [0, 1, 2, 3]
        self.mma_warp_id = 4
        self.tma_warp_id = 5

    @cute.jit
    def __call__(self, a_ptr, b_ptr, c_ptr, problem_m, problem_n, problem_k, stream):
        a_dtype = a_ptr.value_type
        b_dtype = b_ptr.value_type
        c_dtype = c_ptr.value_type
        acc_dtype = Float32

        m, n, k = problem_m, problem_n, problem_k

        # TiledMMA
        tiled_mma = sm100_utils.make_trivial_tiled_mma(
            a_dtype, b_dtype,
            tcgen05.OperandMajorMode.K,
            tcgen05.OperandMajorMode.K,
            acc_dtype,
            self.cta_group,
            self.mma_tiler_mn,
        )

        atom_thr_size = cute.size(tiled_mma.thr_id.shape)
        self.atom_thr_size = atom_thr_size
        mma_tiler = (self.mma_tiler_mn[0], self.mma_tiler_mn[1], cute.size(tiled_mma.shape_mnk, mode=[2]) * 4)
        self.mma_tiler = mma_tiler
        cta_tile_shape_mnk = (mma_tiler[0] // atom_thr_size, mma_tiler[1], mma_tiler[2])
        self.cta_tile_shape_mnk = cta_tile_shape_mnk

        # SMEM layouts
        num_ab_stages = 1
        a_smem_layout = sm100_utils.make_smem_layout_a(tiled_mma, mma_tiler, a_dtype, num_ab_stages)
        b_smem_layout = sm100_utils.make_smem_layout_b(tiled_mma, mma_tiler, b_dtype, num_ab_stages)

        # Epilogue tile
        c_layout_enum = utils.LayoutEnum.ROW_MAJOR
        epi_tile = sm100_utils.compute_epilogue_tile_shape(
            cta_tile_shape_mnk, False, c_layout_enum, c_dtype)
        self.epi_tile = epi_tile

        c_smem_layout = sm100_utils.make_smem_layout_epi(c_dtype, c_layout_enum, epi_tile, 2)
        self.c_smem_layout = c_smem_layout

        # TMEM columns
        self.num_accumulator_tmem_cols = cta_tile_shape_mnk[1]

        # GMEM tensors
        a_gmem_layout = cute.make_ordered_layout((m, k), order=(1, 0))
        b_gmem_layout = cute.make_ordered_layout((n, k), order=(1, 0))
        c_gmem_layout = cute.make_ordered_layout((m, n), order=(1, 0))
        gA = cute.make_tensor(a_ptr, a_gmem_layout)
        gB = cute.make_tensor(b_ptr, b_gmem_layout)
        gC = cute.make_tensor(c_ptr, c_gmem_layout)

        # TMA descriptors
        a_smem_layout_one = cute.slice_(a_smem_layout, (None, None, None, 0))
        b_smem_layout_one = cute.slice_(b_smem_layout, (None, None, None, 0))
        c_smem_layout_one = cute.slice_(c_smem_layout, (None, None, 0))

        tma_atom_a, tma_tensor_a = cute.nvgpu.make_tiled_tma_atom_A(
            sm100_utils.cluster_shape_to_tma_atom_A((1, 1), tiled_mma.thr_id),
            gA, a_smem_layout_one, mma_tiler, tiled_mma,
            (1, 1, 1),
        )
        tma_atom_b, tma_tensor_b = cute.nvgpu.make_tiled_tma_atom_B(
            sm100_utils.cluster_shape_to_tma_atom_B((1, 1), tiled_mma.thr_id),
            gB, b_smem_layout_one, mma_tiler, tiled_mma,
            (1, 1, 1),
        )
        tma_atom_c, tma_tensor_c = cpasync.make_tiled_tma_atom(
            cpasync.CopyBulkTensorTileS2GOp(),
            gC, c_smem_layout_one, epi_tile,
        )

        # Pipeline barriers
        a_copy_size = cute.size_in_bytes(a_dtype, a_smem_layout_one)
        b_copy_size = cute.size_in_bytes(b_dtype, b_smem_layout_one)
        tma_load_bytes = (a_copy_size + b_copy_size) * atom_thr_size

        self.tma_load_bytes = tma_load_bytes

        # Named barriers
        self.epilog_sync_barrier = pipeline.NamedBarrier(
            barrier_id=1,
            num_threads=32 * len(self.epilog_warp_ids),
        )
        self.tmem_alloc_barrier = pipeline.NamedBarrier(
            barrier_id=2,
            num_threads=32 * (1 + len(self.epilog_warp_ids)),
        )

        @cute.struct
        class SharedStorage:
            ab_full_mbar: cute.struct.MemRange[cutlass.Int64, num_ab_stages]
            ab_empty_mbar: cute.struct.MemRange[cutlass.Int64, num_ab_stages]
            acc_full_mbar: cute.struct.MemRange[cutlass.Int64, 1]
            acc_empty_mbar: cute.struct.MemRange[cutlass.Int64, 1]
            tmem_dealloc_mbar: cutlass.Int64
            tmem_holding_buf: cutlass.Int32
            sA: cute.struct.Align[cute.struct.MemRange[a_dtype, cute.cosize(a_smem_layout.outer)], 1024]
            sB: cute.struct.Align[cute.struct.MemRange[b_dtype, cute.cosize(b_smem_layout.outer)], 1024]
            sC: cute.struct.Align[cute.struct.MemRange[c_dtype, cute.cosize(c_smem_layout.outer)], 1024]

        self.shared_storage = SharedStorage

        # Cluster
        self.cluster_shape_mn = (1, 1)
        cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1, 1, 1)), (tiled_mma.thr_id.shape,))
        self.cluster_layout_vmnk = cluster_layout_vmnk

        self._kernel(
            tiled_mma, tma_atom_a, tma_tensor_a, tma_atom_b, tma_tensor_b,
            tma_atom_c, tma_tensor_c, cluster_layout_vmnk,
            a_smem_layout, b_smem_layout, c_smem_layout, epi_tile,
            gA, gB, gC, mma_tiler,
        ).launch(grid=(1, 1, 1), block=[self.threads_per_cta, 1, 1], stream=stream)

    @cute.kernel
    def _kernel(self, tiled_mma, tma_atom_a, mA, tma_atom_b, mB,
                tma_atom_c, mC, cluster_layout_vmnk,
                a_smem_layout, b_smem_layout, c_smem_layout, epi_tile,
                gA, gB, gC, mma_tiler):
        warp_idx = cute.arch.warp_idx()
        warp_idx = cute.arch.make_warp_uniform(warp_idx)
        tidx, _, _ = cute.arch.thread_idx()

        use_2cta = cute.size(tiled_mma.thr_id.shape) == 2

        smem = utils.SmemAllocator()
        storage = smem.allocate(self.shared_storage)

        # AB pipeline
        ab_pipeline = pipeline.PipelineTmaUmma.create(
            barrier_storage=storage.ab_full_mbar.data_ptr(),
            num_stages=1,
            producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),
            consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1),
            tx_count=self.tma_load_bytes,
            cta_layout_vmnk=cluster_layout_vmnk,
            defer_sync=True,
        )

        # Accumulator pipeline
        acc_pipeline = pipeline.PipelineUmmaAsync.create(
            barrier_storage=storage.acc_full_mbar.data_ptr(),
            num_stages=1,
            producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),
            consumer_group=pipeline.CooperativeGroup(
                pipeline.Agent.Thread,
                32 * len(self.epilog_warp_ids) * (2 if use_2cta else 1)),
            cta_layout_vmnk=cluster_layout_vmnk,
            defer_sync=True,
        )

        # TMEM allocator
        tmem = utils.TmemAllocator(
            storage.tmem_holding_buf.ptr,
            barrier_for_retrieve=self.tmem_alloc_barrier,
            allocator_warp_id=self.epilog_warp_ids[0],
            is_two_cta=use_2cta,
            two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar.ptr,
        )

        pipeline.pipeline_init_arrive(cluster_shape_mn=self.cluster_shape_mn, is_relaxed=True)

        # SMEM tensors
        sA = storage.sA.get_tensor(a_smem_layout.outer, swizzle=a_smem_layout.inner)
        sB = storage.sB.get_tensor(b_smem_layout.outer, swizzle=b_smem_layout.inner)
        sC = storage.sC.get_tensor(c_smem_layout.outer, swizzle=c_smem_layout.inner)

        # gC tiled for epilogue partition
        gC_tiled = cute.local_tile(gC, (mma_tiler[0], mma_tiler[1]), (0, 0))

        # Partition for TMA — use TMA tensors directly (they have static shape)
        thr_mma = tiled_mma.get_slice(0)
        tCgC = thr_mma.partition_C(gC_tiled)

        tAsA, tAgA = cpasync.tma_partition(
            tma_atom_a, 0, cute.make_layout(1),
            cute.group_modes(sA, 0, 3),
            cute.group_modes(mA, 0, 3),
        )
        tBsB, tBgB = cpasync.tma_partition(
            tma_atom_b, 0, cute.make_layout(1),
            cute.group_modes(sB, 0, 3),
            cute.group_modes(mB, 0, 3),
        )

        # MMA fragments
        tCrA = tiled_mma.make_fragment_A(sA)
        tCrB = tiled_mma.make_fragment_B(sB)

        # TMEM accumulator
        acc_shape = tiled_mma.partition_shape_C(mma_tiler[:2])
        tCtAcc_all = tiled_mma.make_fragment_C(cute.append(acc_shape, 1))

        k_tile_cnt = cute.size(mA, mode=[3])
        
        # gC tiled for epilogue partition
        gC_tiled = cute.local_tile(gC, (mma_tiler[0], mma_tiler[1]), (0, 0))

        pipeline.pipeline_init_wait(cluster_shape_mn=self.cluster_shape_mn)

        # ══════════════════════════════════════════════════════════
        # TMA LOAD WARP (warp 5)
        # ══════════════════════════════════════════════════════════
        if warp_idx == self.tma_warp_id:
            cpasync.prefetch_descriptor(tma_atom_a)
            cpasync.prefetch_descriptor(tma_atom_b)
            cpasync.prefetch_descriptor(tma_atom_c)

            ab_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, 1)
            for k_tile in cutlass.range(k_tile_cnt, unroll=1):
                ab_pipeline.producer_acquire(ab_state)
                cute.copy(tma_atom_a, tAgA[(None, ab_state.count)], tAsA[(None, ab_state.index)],
                          tma_bar_ptr=ab_pipeline.producer_get_barrier(ab_state))
                cute.copy(tma_atom_b, tBgB[(None, ab_state.count)], tBsB[(None, ab_state.index)],
                          tma_bar_ptr=ab_pipeline.producer_get_barrier(ab_state))
                ab_state.advance()
            ab_pipeline.producer_tail(ab_state)

        # ══════════════════════════════════════════════════════════
        # MMA WARP (warp 4)
        # ══════════════════════════════════════════════════════════
        if warp_idx == self.mma_warp_id:
            tmem.wait_for_alloc()
            acc_ptr = tmem.retrieve_ptr(Float32)
            tCtAcc_base = cute.make_tensor(acc_ptr, tCtAcc_all.layout)
            tCtAcc = tCtAcc_base[(None, None, None, 0)]

            ab_cstate = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, 1)
            acc_pstate = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, 1)

            acc_pipeline.producer_acquire(acc_pstate)
            tiled_mma.set(tcgen05.Field.ACCUMULATE, False)

            for k_tile in range(k_tile_cnt):
                ab_pipeline.consumer_wait(ab_cstate, cutlass.Boolean(1))
                for kblock in cutlass.range(cute.size(tCrA, mode=[2]), unroll_full=True):
                    coord = (None, None, kblock, ab_cstate.index)
                    cute.gemm(tiled_mma, tCtAcc, tCrA[coord], tCrB[coord], tCtAcc)
                    tiled_mma.set(tcgen05.Field.ACCUMULATE, True)
                ab_pipeline.consumer_release(ab_cstate)
                ab_cstate.advance()

            acc_pipeline.producer_commit(acc_pstate)
            acc_pstate.advance()
            acc_pipeline.producer_tail(acc_pstate)

        # ══════════════════════════════════════════════════════════
        # EPILOGUE WARPS (0..3)
        # ══════════════════════════════════════════════════════════
        if warp_idx < self.mma_warp_id:
            tmem.allocate(self.num_accumulator_tmem_cols)
            tmem.wait_for_alloc()

            acc_ptr = tmem.retrieve_ptr(Float32)
            tCtAcc_base = cute.make_tensor(acc_ptr, tCtAcc_all.layout)

            c_layout_enum = utils.LayoutEnum.ROW_MAJOR
            c_dtype = BFloat16

            # TMEM→reg
            copy_atom_t2r = sm100_utils.get_tmem_load_op(
                self.cta_tile_shape_mnk, c_layout_enum, c_dtype, Float32, epi_tile, False)
            tAcc_epi = cute.flat_divide(tCtAcc_base[((None, None), 0, 0, None)], epi_tile)
            tiled_copy_t2r = tcgen05.make_tmem_copy(copy_atom_t2r, tAcc_epi[(None, None, 0, 0, 0)])
            thr_t2r = tiled_copy_t2r.get_slice(tidx)
            tTR_tAcc = thr_t2r.partition_S(tAcc_epi)
            tTR_rAcc = cute.make_rmem_tensor(
                thr_t2r.partition_D(
                    cute.flat_divide(tCgC[((None, None), 0, 0, None, None)], epi_tile)
                )[(None, None, None, 0, 0, 0, 0)].shape, Float32)
            tTR_rC = cute.make_rmem_tensor(tTR_rAcc.shape, c_dtype)

            # reg→SMEM
            copy_atom_r2s = sm100_utils.get_smem_store_op(c_layout_enum, c_dtype, Float32, tiled_copy_t2r)
            tiled_copy_r2s = cute.make_tiled_copy_D(copy_atom_r2s, tiled_copy_t2r)
            thr_r2s = tiled_copy_r2s.get_slice(tidx)
            tRS_sC = thr_r2s.partition_D(sC)
            tRS_rC = tiled_copy_r2s.retile(tTR_rC)

            # SMEM→GMEM (TMA)
            gC_epi = cute.flat_divide(tCgC[((None, None), 0, 0, None, None)], epi_tile)
            bSG_sC, bSG_gC = cpasync.tma_partition(
                tma_atom_c, 0, cute.make_layout(1),
                cute.group_modes(sC, 0, 2),
                cute.group_modes(gC_epi, 0, 2))

            acc_cstate = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, 1)
            c_pipeline = pipeline.PipelineTmaStore.create(
                num_stages=2,
                producer_group=pipeline.CooperativeGroup(
                    pipeline.Agent.Thread, 32 * len(self.epilog_warp_ids)))

            acc_pipeline.consumer_wait(acc_cstate)

            tTR_tAcc_g = cute.group_modes(tTR_tAcc, 3, cute.rank(tTR_tAcc))
            bSG_gC_g = cute.group_modes(bSG_gC, 1, cute.rank(bSG_gC))

            for subtile in cutlass.range(cute.size(tTR_tAcc_g.shape, mode=[3])):
                cute.copy(tiled_copy_t2r, tTR_tAcc_g[(None, None, None, subtile)], tTR_rAcc)
                acc_vec = tiled_copy_r2s.retile(tTR_rAcc).load()
                tRS_rC.store(acc_vec.to(c_dtype))

                c_buf = subtile % 2
                cute.copy(tiled_copy_r2s, tRS_rC, tRS_sC[(None, None, None, c_buf)])
                cute.arch.fence_proxy("async.shared", space="cta")
                self.epilog_sync_barrier.arrive_and_wait()

                if warp_idx == self.epilog_warp_ids[0]:
                    cute.copy(tma_atom_c, bSG_sC[(None, c_buf)], bSG_gC_g[(None, subtile)])
                    c_pipeline.producer_commit()
                    c_pipeline.producer_acquire()
                self.epilog_sync_barrier.arrive_and_wait()

            acc_pipeline.consumer_release(acc_cstate)
            tmem.relinquish_alloc_permit()
            self.epilog_sync_barrier.arrive_and_wait()
            tmem.free(acc_ptr)
            c_pipeline.producer_tail()


from cutlass.cute.runtime import make_ptr
from cutlass.cute.nvgpu import cpasync
import cutlass.utils.blackwell_helpers as sm100_utils

def test_minimal_mma():
    device = torch.device("cuda")
    torch.manual_seed(42)

    prob_m, prob_n, prob_k = M, N, K

    tA = torch.randn(prob_m, prob_k, dtype=torch.bfloat16, device=device)
    tB = torch.randn(prob_n, prob_k, dtype=torch.bfloat16, device=device)
    ref = torch.matmul(tA.to(torch.float32), tB.to(torch.float32).T)
    tC = torch.zeros(prob_m, prob_n, dtype=torch.bfloat16, device=device)

    a_ptr = make_ptr(BFloat16, 0, cute.AddressSpace.gmem, assumed_align=16)
    b_ptr = make_ptr(BFloat16, 0, cute.AddressSpace.gmem, assumed_align=16)
    c_ptr = make_ptr(BFloat16, 0, cute.AddressSpace.gmem, assumed_align=16)

    stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream)

    kernel = MinimalMMAKernel()
    compiled = cute.compile(kernel, a_ptr, b_ptr, c_ptr,
                            cutlass.Int32(prob_m), cutlass.Int32(prob_n), cutlass.Int32(prob_k), stream)

    a_ptr_r = make_ptr(BFloat16, tA.data_ptr(), cute.AddressSpace.gmem, assumed_align=16)
    b_ptr_r = make_ptr(BFloat16, tB.data_ptr(), cute.AddressSpace.gmem, assumed_align=16)
    c_ptr_r = make_ptr(BFloat16, tC.data_ptr(), cute.AddressSpace.gmem, assumed_align=16)

    compiled(a_ptr_r, b_ptr_r, c_ptr_r, prob_m, prob_n, prob_k, stream)
    torch.cuda.synchronize()

    output = tC.to(torch.float32)
    cos = torch.nn.functional.cosine_similarity(
        output.flatten().unsqueeze(0), ref.flatten().unsqueeze(0)).item()
    max_err = (output - ref).abs().max().item()

    print(f"Minimal MMA: A({prob_m},{prob_k}) @ B^T({prob_k},{prob_n}) → C({prob_m},{prob_n})")
    print(f"  Cosine: {cos:.6f}, Max error: {max_err:.6f}")
    print(f"  {'✅ PASS' if cos >= 0.99 else '❌ FAIL'}")
    return cos


if __name__ == "__main__":
    test_minimal_mma()
