"""
Stage B v2: Two MMAs (Q@K^T then Scores@V) — no softmax, no identity P.

Tests:
- MMA1: Q @ K^T → tmem_scores (accumulate=False)
- MMA2: tmem_scores @ V → tmem_output (a_source=TMEM, accumulate=True)
- Two TMEM regions with pointer arithmetic
- The epilogue warps allocate TMEM and store output, but skip softmax for now

Reference: output = Q @ K^T @ V (with K=V for simplicity)
"""
import torch
import cutlass
import cutlass.cute as cute
import cutlass.utils as utils
import cutlass.pipeline as pipeline
from cutlass.cute.nvgpu import cpasync, tcgen05
from cutlass import Float32, BFloat16, Int32, Boolean, const_expr
from cutlass.utils import LayoutEnum
import cuda.bindings.driver as cuda


class StageBKernel:
    def __init__(self, mma_tiler_mn, use_2cta_instrs=False):
        self.acc_dtype = Float32
        self.use_2cta_instrs = use_2cta_instrs
        self.mma_tiler_mn = mma_tiler_mn
        self.mma_tiler = (*mma_tiler_mn, 1)
        self.cluster_shape_mn = (1, 1)
        self.cta_group = tcgen05.CtaGroup.TWO if use_2cta_instrs else tcgen05.CtaGroup.ONE
        # Warp layout: 4 epilogue + 1 MMA + 1 TMA = 6 warps = 192 threads
        self.epilogue_warp_id = (0, 1, 2, 3)
        self.mma_warp_id = 4
        self.tma_warp_id = 5
        self.threads_per_cta = 192
        self.epilog_sync_bar_id = 1
        self.tmem_alloc_sync_bar_id = 2

    def _setup_attributes(self, tiled_mma1, tiled_mma2, a_dtype, b_dtype, c_dtype,
                          a_major, b_major, c_layout):
        mma_inst_shape_k = cute.size(tiled_mma1.shape_mnk, mode=[2])
        self.mma_tiler = (*self.mma_tiler_mn, mma_inst_shape_k * 4)
        self.cta_tile_shape_mnk = (
            self.mma_tiler[0] // cute.size(tiled_mma1.thr_id.shape),
            self.mma_tiler[1],
            self.mma_tiler[2],
        )
        self.cluster_layout_vmnk = cute.tiled_divide(
            cute.make_layout((1, 1, 1)), (tiled_mma1.thr_id.shape,))

        self.epi_tile = utils.sm100.compute_epilogue_tile_shape(
            self.cta_tile_shape_mnk, self.use_2cta_instrs, c_layout, c_dtype)

        self.num_ab_stage = 1
        self.num_acc_stage = 1
        self.num_c_stage = 2

        self.a_smem_layout_staged = utils.sm100.make_smem_layout_a(
            tiled_mma1, self.mma_tiler, a_dtype, self.num_ab_stage)
        self.b_smem_layout_staged = utils.sm100.make_smem_layout_b(
            tiled_mma1, self.mma_tiler, b_dtype, self.num_ab_stage)
        self.c_smem_layout_staged = utils.sm100.make_smem_layout_epi(
            c_dtype, c_layout, self.epi_tile, self.num_c_stage)

        # TMEM: two regions (scores + output), each partition_shape_C columns
        acc_shape = tiled_mma1.partition_shape_C(self.mma_tiler_mn)
        tCtAcc_fake = tiled_mma1.make_fragment_C(cute.append(acc_shape, self.num_acc_stage))
        self.num_tmem_cols_per_region = utils.get_num_tmem_alloc_cols(tCtAcc_fake, arch="sm_100")
        self.total_tmem_cols = max(self.num_tmem_cols_per_region * 2, 256)

        a_smem = cute.slice_(self.a_smem_layout_staged, (None, None, None, 0))
        b_smem = cute.slice_(self.b_smem_layout_staged, (None, None, None, 0))
        self.num_tma_load_bytes = (
            cute.size_in_bytes(a_dtype, a_smem) +
            cute.size_in_bytes(b_dtype, b_smem)
        ) * cute.size(tiled_mma1.thr_id.shape)

    @cute.jit
    def __call__(self, a: cute.Tensor, b: cute.Tensor, c: cute.Tensor,
                 stream: cuda.CUstream):
        a_dtype = a.element_type
        b_dtype = b.element_type
        c_dtype = c.element_type
        a_major = LayoutEnum.from_tensor(a).mma_major_mode()
        b_major = LayoutEnum.from_tensor(b).mma_major_mode()
        c_layout = LayoutEnum.from_tensor(c)

        tiled_mma1 = utils.sm100.make_trivial_tiled_mma(
            a_dtype, b_dtype, a_major, b_major,
            self.acc_dtype, self.cta_group, self.mma_tiler_mn,
            tcgen05.OperandSource.SMEM,
        )
        tiled_mma2 = utils.sm100.make_trivial_tiled_mma(
            a_dtype, b_dtype, a_major, b_major,
            self.acc_dtype, self.cta_group, self.mma_tiler_mn,
            tcgen05.OperandSource.SMEM,
        )

        self._setup_attributes(tiled_mma1, tiled_mma2, a_dtype, b_dtype, c_dtype,
                               a_major, b_major, c_layout)
        # These are needed by epilogue_tma_store which accesses them via self
        self.a_dtype = a_dtype
        self.b_dtype = b_dtype
        self.c_dtype = c_dtype
        self.a_major_mode = a_major
        self.b_major_mode = b_major
        self.c_layout = c_layout

        a_smem = cute.slice_(self.a_smem_layout_staged, (None, None, None, 0))
        b_smem = cute.slice_(self.b_smem_layout_staged, (None, None, None, 0))

        tma_a, tma_tensor_a = cute.nvgpu.make_tiled_tma_atom_A(
            utils.sm100.cluster_shape_to_tma_atom_A(
                self.cluster_shape_mn, tiled_mma1.thr_id),
            a, a_smem, self.mma_tiler, tiled_mma1,
            self.cluster_layout_vmnk.shape,
        )
        tma_b, tma_tensor_b = cute.nvgpu.make_tiled_tma_atom_B(
            utils.sm100.cluster_shape_to_tma_atom_B(
                self.cluster_shape_mn, tiled_mma1.thr_id),
            b, b_smem, self.mma_tiler, tiled_mma1,
            self.cluster_layout_vmnk.shape,
        )

        epi_smem = cute.select(self.c_smem_layout_staged, mode=[0, 1])
        tma_c, tma_tensor_c = cpasync.make_tiled_tma_atom(
            cpasync.CopyBulkTensorTileS2GOp(), c, epi_smem, self.epi_tile)

        self._kernel(
            tiled_mma1, tiled_mma2,
            tma_a, tma_tensor_a, tma_b, tma_tensor_b,
            tma_c, tma_tensor_c, self.cluster_layout_vmnk,
            self.a_smem_layout_staged, self.b_smem_layout_staged,
            self.c_smem_layout_staged, self.epi_tile,
        ).launch(grid=(1, 1, 1), block=[self.threads_per_cta, 1, 1], stream=stream)

    @cute.kernel
    def _kernel(self, tiled_mma1, tiled_mma2,
                tma_atom_a, mA_mkl, tma_atom_b, mB_nkl,
                tma_atom_c, mC_mnl, cluster_layout_vmnk,
                a_smem_layout_staged, b_smem_layout_staged,
                c_smem_layout_staged, epi_tile):
        warp_idx = cute.arch.warp_idx()
        warp_idx = cute.arch.make_warp_uniform(warp_idx)
        tidx, _, _ = cute.arch.thread_idx()
        use_2cta_instrs = cute.size(tiled_mma1.thr_id.shape) == 2
        is_leader_cta = True

        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)

        @cute.struct
        class SharedStorage:
            ab_full_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2]
            acc_full_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2]
            tmem_dealloc_mbar: cutlass.Int64
            tmem_holding_buf: cutlass.Int32

        smem = utils.SmemAllocator()
        storage = smem.allocate(SharedStorage)

        ab_producer, ab_consumer = pipeline.PipelineTmaUmma.create(
            barrier_storage=storage.ab_full_mbar_ptr.data_ptr(),
            num_stages=self.num_ab_stage,
            producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),
            consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1),
            tx_count=self.num_tma_load_bytes,
            cta_layout_vmnk=cluster_layout_vmnk,
            defer_sync=True,
        ).make_participants()

        acc_pipeline = pipeline.PipelineUmmaAsync.create(
            barrier_storage=storage.acc_full_mbar_ptr.data_ptr(),
            num_stages=self.num_acc_stage,
            producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),
            consumer_group=pipeline.CooperativeGroup(
                pipeline.Agent.Thread, len(self.epilogue_warp_id) * (2 if use_2cta_instrs else 1)),
            cta_layout_vmnk=cluster_layout_vmnk,
            defer_sync=True,
        )

        tmem_alloc_barrier = pipeline.NamedBarrier(
            barrier_id=self.tmem_alloc_sync_bar_id,
            num_threads=32 * len((self.mma_warp_id, *self.epilogue_warp_id)),
        )
        tmem = utils.TmemAllocator(
            storage.tmem_holding_buf.ptr,
            barrier_for_retrieve=tmem_alloc_barrier,
            allocator_warp_id=self.epilogue_warp_id[0],
            is_two_cta=use_2cta_instrs,
            two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar.ptr,
        )

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

        sA = smem.allocate_tensor(
            element_type=BFloat16,
            layout=a_smem_layout_staged.outer,
            byte_alignment=128, swizzle=a_smem_layout_staged.inner)
        sB = smem.allocate_tensor(
            element_type=BFloat16,
            layout=b_smem_layout_staged.outer,
            byte_alignment=128, swizzle=b_smem_layout_staged.inner)
        sC = smem.allocate_tensor(
            element_type=BFloat16,
            layout=c_smem_layout_staged.outer,
            byte_alignment=128, swizzle=c_smem_layout_staged.inner)

        gA_mkl = cute.local_tile(mA_mkl, cute.slice_(self.mma_tiler, (None, 0, None)), (None, None, None))
        gB_nkl = cute.local_tile(mB_nkl, cute.slice_(self.mma_tiler, (0, None, None)), (None, None, None))
        gC_mnl = cute.local_tile(mC_mnl, cute.slice_(self.mma_tiler, (None, None, 0)), (None, None, None))
        k_tile_cnt = cute.size(gA_mkl, mode=[3])

        thr_mma1 = tiled_mma1.get_slice(0)
        tCgA = thr_mma1.partition_A(gA_mkl)
        tCgB = thr_mma1.partition_B(gB_nkl)
        tCgC = thr_mma1.partition_C(gC_mnl)

        a_cta_layout = cute.make_layout(cute.slice_(cluster_layout_vmnk, (0, 0, None, 0)).shape)
        tAsA, tAgA = cpasync.tma_partition(
            tma_atom_a, 0, a_cta_layout,
            cute.group_modes(sA, 0, 3), cute.group_modes(tCgA, 0, 3))
        b_cta_layout = cute.make_layout(cute.slice_(cluster_layout_vmnk, (0, None, 0, 0)).shape)
        tBsB, tBgB = cpasync.tma_partition(
            tma_atom_b, 0, b_cta_layout,
            cute.group_modes(sB, 0, 3), cute.group_modes(tCgB, 0, 3))

        tAgA_slice = tAgA[(None, 0, None, 0)]
        tBgB_slice = tBgB[(None, 0, None, 0)]

        # MMA1 fragments
        tCrA = tiled_mma1.make_fragment_A(sA)
        tCrB = tiled_mma1.make_fragment_B(sB)
        # MMA2 fragment for B (V) - same SMEM as K
        tCrB_mma2 = tiled_mma2.make_fragment_B(sB)
        # MMA2 fragment for A (a_source=SMEM, same as MMA1)
        tCrA_mma2 = tiled_mma2.make_fragment_A(sA)

        # TMEM accumulator layout (same for both MMAs)
        acc_shape = tiled_mma1.partition_shape_C(self.mma_tiler_mn)
        tCtAcc_fake = tiled_mma1.make_fragment_C(cute.append(acc_shape, self.num_acc_stage))

        pipeline.pipeline_init_wait(cluster_shape_mn=cluster_layout_vmnk)

        # ── TMA LOAD WARP (warp 5) ──
        if warp_idx == self.tma_warp_id:
            ab_producer.reset()
            peek = ab_producer.try_acquire()
            for k_tile in cutlass.range(k_tile_cnt, unroll=1):
                handle = ab_producer.acquire_and_advance(peek)
                cute.copy(tma_atom_a, tAgA_slice[(None, handle.count)], tAsA[(None, handle.index)],
                          tma_bar_ptr=handle.barrier)
                cute.copy(tma_atom_b, tBgB_slice[(None, handle.count)], tBsB[(None, handle.index)],
                          tma_bar_ptr=handle.barrier)
                peek = cutlass.Boolean(1)
                if handle.count + 1 < k_tile_cnt:
                    peek = ab_producer.try_acquire()
            ab_producer.tail()

        # ── MMA WARP (warp 4) ──
        if warp_idx == self.mma_warp_id:
            tmem.wait_for_alloc()
            tmem_ptr = tmem.retrieve_ptr(self.acc_dtype)

            # TMEM region 0: scores (Q @ K^T)
            tCtScores_base = cute.make_tensor(tmem_ptr, tCtAcc_fake.layout)
            tCtScores = tCtScores_base[(None, None, None, 0)]
            # MMA2 A fragment from SMEM (a_source=SMEM)
            tCrA_mma2 = tiled_mma2.make_fragment_A(sA)

            # TMEM region 1: output (Scores @ V)
            output_ptr = cute.recast_ptr(
                tmem_ptr + self.num_tmem_cols_per_region, dtype=self.acc_dtype)
            tCtOutput_base = cute.make_tensor(output_ptr, tCtAcc_fake.layout)
            tCtOutput = tCtOutput_base[(None, None, None, 0)]

            ab_consumer.reset()
            peek = cutlass.Boolean(1)
            if is_leader_cta:
                peek = ab_consumer.try_wait()

            # ── MMA1: Q @ K^T → tmem_scores ──
            tiled_mma1.set(tcgen05.Field.ACCUMULATE, False)
            for k_tile in range(k_tile_cnt):
                if is_leader_cta:
                    handle = ab_consumer.wait_and_advance(peek)
                    nblk = cute.size(tCrA, mode=[2])
                    for kblk in cutlass.range(nblk, unroll_full=True):
                        crd = (None, None, kblk, handle.index)
                        cute.gemm(tiled_mma1, tCtScores, tCrA[crd], tCrB[crd], tCtScores)
                    handle.release()
                    peek = cutlass.Boolean(1)
                    if handle.count + 1 < k_tile_cnt:
                        peek = ab_consumer.try_wait()

            # MMA1 done, scores in tmem_scores
            # Fence to ensure TMEM writes are visible
            cute.arch.fence_view_async_tmem_store()

            # ── MMA2: Scores @ V → tmem_output ──
            # a_source=TMEM: the MMA instruction reads A from TMEM
            # The A operand in TMEM is the scores from MMA1
            # We need to set the TMEM pointer for A using tiled_mma2.set()
            # The C operand's TMEM pointer tells MMA2 where to write the output
            # The A operand's TMEM pointer needs to be set separately
            tiled_mma2.set(tcgen05.Field.ACCUMULATE, True)
            # Set the A TMEM pointer to the scores region
            # When a_source=TMEM, the MMA reads A from the C operand's base
            # So we need tCtOutput to have the same base as tCtScores for MMA2
            # This means we can't use two separate TMEM regions for this approach
            # Instead, let's try passing the TMEM A operand via the auxiliary list
            # cute.gemm(tiled_mma2, D, [A_main, A_tmem], B, C)
            # Or use tiled_mma2.set(Field.SFA, tCtScores.iterator) to set A ptr
            nblk2 = cute.size(tCrB_mma2, mode=[2])
            for kblk in cutlass.range(nblk2, unroll_full=True):
                crd = (None, None, kblk, 0)
                cute.gemm(tiled_mma2, tCtOutput, tCrA_mma2[crd], tCrB_mma2[crd], tCtOutput)

            # Signal output ready
            acc_producer_state = pipeline.make_pipeline_state(
                pipeline.PipelineUserType.Producer, self.num_acc_stage)
            if is_leader_cta:
                acc_pipeline.producer_acquire(acc_producer_state)
                acc_pipeline.producer_commit(acc_producer_state)
            acc_producer_state.advance()
            acc_pipeline.producer_tail(acc_producer_state)

        # ── EPILOGUE WARPS (0..3) ──
        if warp_idx < self.mma_warp_id:
            tmem.allocate(self.total_tmem_cols)
            tmem.wait_for_alloc()
            tmem_ptr = tmem.retrieve_ptr(self.acc_dtype)

            # TMEM region 1: output
            output_ptr = cute.recast_ptr(
                tmem_ptr + self.num_tmem_cols_per_region, dtype=self.acc_dtype)
            tCtOutput_base = cute.make_tensor(output_ptr, tCtAcc_fake.layout)

            # Wait for MMA2 to finish
            acc_consumer_state = pipeline.make_pipeline_state(
                pipeline.PipelineUserType.Consumer, self.num_acc_stage)

            c_producer_group = pipeline.CooperativeGroup(
                pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id))
            c_pipeline = pipeline.PipelineTmaStore.create(
                num_stages=self.num_c_stage, producer_group=c_producer_group)

            mma_tile_coord_mnl = (0, 0, 0)
            epilogue_op = const_expr(lambda x: x)
            num_tiles_executed = 0

            acc_consumer_state = utils.gemm.sm100.epilogue_tma_store(
                self, tidx, warp_idx, tma_atom_c, tCtOutput_base, sC, tCgC,
                epi_tile, num_tiles_executed, epilogue_op,
                mma_tile_coord_mnl, acc_consumer_state, acc_pipeline, c_pipeline)

            c_pipeline.producer_tail()
            tmem.relinquish_alloc_permit()
            tmem.free(tmem_ptr)


def test_stage_b():
    torch.manual_seed(42)
    m, n, k = 128, 128, 128

    a = torch.randn(m, k, 1, dtype=torch.bfloat16, device="cuda")
    b = torch.randn(n, k, 1, dtype=torch.bfloat16, device="cuda")
    c = torch.zeros(m, n, 1, dtype=torch.bfloat16, device="cuda")

    # Reference: Q @ K^T @ V (no softmax, K=V=b)
    q = a[:, :, 0].float()
    kt = b[:, :, 0].float().T  # K^T
    v = b[:, :, 0].float()
    ref = q @ kt @ v  # (128, 128)

    import cutlass.torch as cutlass_torch
    mA = cutlass_torch.from_dlpack(a).mark_layout_dynamic(
        leading_dim=cutlass_torch.get_leading_dim(a))
    mB = cutlass_torch.from_dlpack(b).mark_layout_dynamic(
        leading_dim=cutlass_torch.get_leading_dim(b))
    mC = cutlass_torch.from_dlpack(c).mark_layout_dynamic(
        leading_dim=cutlass_torch.get_leading_dim(c))

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

    kernel = StageBKernel(mma_tiler_mn=(128, 128))
    print("Compiling...", flush=True)
    compiled = cute.compile(kernel, mA, mB, mC, stream)
    print("Running...", flush=True)
    compiled(mA, mB, mC, stream)
    torch.cuda.synchronize()

    output = c[:, :, 0].float()
    cos = torch.nn.functional.cosine_similarity(
        output.flatten().unsqueeze(0), ref.flatten().unsqueeze(0)).item()
    max_err = (output - ref).abs().max().item()

    print("Stage B v2: Q @ K^T @ V (no softmax)")
    print("  Cosine: {:.6f}, Max error: {:.6f}".format(cos, max_err))
    print("  {}".format("PASS" if cos >= 0.99 else "FAIL"))
    return cos


if __name__ == "__main__":
    test_stage_b()
