"""
Stage B: Two MMAs (Q@K^T then P@V) with a_source=TMEM for MMA2.

Architecture (following NVIDIA's fmha.py reference):
  MMA1: Q @ K^T → tmem_scores (a_source=SMEM, accumulate=False)
  MMA2: P @ V → tmem_output (a_source=TMEM, accumulate=True)
  P fragment: constructed with P TMEM layout from make_smem_layout_A(pv_tiled_mma, ...)
  P TMEM address: fragment.iterator + (acc_width/a_width) * tmem_p_offset

Reference: output = Q @ K^T @ V (no softmax, P = raw scores)
"""
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
        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

    def _setup(self, qk_mma, pv_mma, a_dtype, b_dtype, c_dtype, a_major, b_major, c_layout):
        self.a_dtype = a_dtype
        self.b_dtype = b_dtype
        self.c_dtype = c_dtype
        self.c_layout = c_layout
        self.use_2cta_instrs = False

        # QK MMA tiler
        qk_inst_k = cute.size(qk_mma.shape_mnk, mode=[2])
        self.qk_mma_tiler = (*self.mma_tiler_mn, qk_inst_k * 4)
        # PV MMA tiler (same M,N but potentially different K)
        pv_inst_k = cute.size(pv_mma.shape_mnk, mode=[2])
        self.pv_mma_tiler = (*self.mma_tiler_mn, pv_inst_k * 4)
        # Use QK tiler for the overall tiler (A/B come from QK layout)
        self.mma_tiler = self.qk_mma_tiler

        self.cta_tile_shape_mnk = (self.qk_mma_tiler[0], self.qk_mma_tiler[1], self.qk_mma_tiler[2])
        self.cluster_layout_vmnk = cute.tiled_divide(cute.make_layout((1,1,1)), (qk_mma.thr_id.shape,))

        # Epilogue tile from PV MMA (output is O, same shape as PV MMA's C)
        self.epi_tile = utils.sm100.compute_epilogue_tile_shape(
            self.cta_tile_shape_mnk, False, c_layout, c_dtype)

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

        # SMEM layouts for QK MMA (Q and K)
        self.q_smem_layout = utils.sm100.make_smem_layout_a(qk_mma, self.qk_mma_tiler, a_dtype, 1)
        self.k_smem_layout = utils.sm100.make_smem_layout_b(qk_mma, self.qk_mma_tiler, b_dtype, 1)
        # SMEM layout for V (from PV MMA)
        self.v_smem_layout = utils.sm100.make_smem_layout_b(pv_mma, self.pv_mma_tiler, b_dtype, 1)
        # TMEM layout for P (from PV MMA's A operand — this is the KEY)
        self.p_tmem_layout = utils.sm100.make_smem_layout_a(pv_mma, self.pv_mma_tiler, a_dtype, 1)
        # C/Output SMEM layout
        self.c_smem_layout = utils.sm100.make_smem_layout_epi(c_dtype, c_layout, self.epi_tile, 2)

        # TMEM allocation: two regions
        # Region 0: scores (Q@K^T accumulator, QK MMA's C layout)
        acc_shape = qk_mma.partition_shape_C(self.mma_tiler_mn)
        tCtAcc_fake = qk_mma.make_fragment_C(cute.append(acc_shape, 1))
        self.num_tmem_cols_scores = utils.get_num_tmem_alloc_cols(tCtAcc_fake, arch="sm_100")
        # Region 1: output (P@V accumulator, PV MMA's C layout)
        acc_shape_pv = pv_mma.partition_shape_C(self.mma_tiler_mn)
        tCtO_fake = pv_mma.make_fragment_C(cute.append(acc_shape_pv, 1))
        self.num_tmem_cols_output = utils.get_num_tmem_alloc_cols(tCtO_fake, arch="sm_100")
        # Total
        self.total_tmem_cols = max(self.num_tmem_cols_scores + self.num_tmem_cols_output, 256)
        # tmem_p_offset: P TMEM starts at 0 (same as scores, P replaces scores in TMEM)
        self.tmem_p_offset = 0

        # TMA load bytes
        q_smem = cute.slice_(self.q_smem_layout, (None, None, None, 0))
        k_smem = cute.slice_(self.k_smem_layout, (None, None, None, 0))
        self.num_tma_bytes = (
            cute.size_in_bytes(a_dtype, q_smem) +
            cute.size_in_bytes(b_dtype, k_smem)
        ) * cute.size(qk_mma.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)

        # QK MMA: Q @ K^T, A from SMEM
        qk_mma = 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,
        )

        # PV MMA: P @ V, A from TMEM, P is K-major
        # Following NVIDIA fmha.py: P (intermediate) is K-major and from TMEM
        p_major = cute.nvgpu.OperandMajorMode.K
        pv_mma = utils.sm100.make_trivial_tiled_mma(
            a_dtype, b_dtype, p_major, b_major,
            self.acc_dtype, self.cta_group, self.mma_tiler_mn,
            tcgen05.OperandSource.TMEM,
        )

        self._setup(qk_mma, pv_mma, a_dtype, b_dtype, c_dtype, a_major, b_major, c_layout)

        q_smem = cute.slice_(self.q_smem_layout, (None, None, None, 0))
        k_smem = cute.slice_(self.k_smem_layout, (None, None, None, 0))
        v_smem = cute.slice_(self.v_smem_layout, (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, qk_mma.thr_id),
            a, q_smem, self.qk_mma_tiler, qk_mma, 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, qk_mma.thr_id),
            b, k_smem, self.qk_mma_tiler, qk_mma, self.cluster_layout_vmnk.shape)
        tma_v, tma_tensor_v = cute.nvgpu.make_tiled_tma_atom_B(
            utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, pv_mma.thr_id),
            b, v_smem, self.pv_mma_tiler, pv_mma, self.cluster_layout_vmnk.shape)
        epi_smem = cute.select(self.c_smem_layout, mode=[0, 1])
        tma_c, tma_tensor_c = cpasync.make_tiled_tma_atom(
            cpasync.CopyBulkTensorTileS2GOp(), c, epi_smem, self.epi_tile)

        self._kernel(
            qk_mma, pv_mma,
            tma_a, tma_tensor_a, tma_b, tma_tensor_b,
            tma_v, tma_tensor_v,
            tma_c, tma_tensor_c, self.cluster_layout_vmnk,
            self.q_smem_layout, self.k_smem_layout,
            self.v_smem_layout, self.p_tmem_layout,
            self.c_smem_layout, self.epi_tile,
        ).launch(grid=(1,1,1), block=[self.threads_per_cta, 1, 1], stream=stream)

    @cute.kernel
    def _kernel(self, qk_mma, pv_mma,
                tma_q, mQ, tma_k, mK,
                tma_v, mV,
                tma_c, mC, cl_vmnk,
                q_smem_layout, k_smem_layout,
                v_smem_layout, p_tmem_layout,
                c_smem_layout, epi_tile):
        warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx())
        tidx, _, _ = cute.arch.thread_idx()

        if warp_idx == self.tma_warp_id:
            cpasync.prefetch_descriptor(tma_q)
            cpasync.prefetch_descriptor(tma_k)
            cpasync.prefetch_descriptor(tma_v)
            cpasync.prefetch_descriptor(tma_c)

        @cute.struct
        class SS:
            ab_bar: cute.struct.MemRange[cutlass.Int64, 2]
            acc_bar: cute.struct.MemRange[cutlass.Int64, 2]
            dealloc_bar: cutlass.Int64
            holding: cutlass.Int32

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

        ab_prod, ab_cons = pipeline.PipelineTmaUmma.create(
            barrier_storage=storage.ab_bar.data_ptr(), num_stages=1,
            producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),
            consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1),
            tx_count=self.num_tma_bytes, cta_layout_vmnk=cl_vmnk,
            defer_sync=True).make_participants()

        acc_pipe = pipeline.PipelineUmmaAsync.create(
            barrier_storage=storage.acc_bar.data_ptr(), num_stages=1,
            producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),
            consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 128),
            cta_layout_vmnk=cl_vmnk, defer_sync=True)

        tmem_bar = pipeline.NamedBarrier(barrier_id=2, num_threads=160)
        tmem = utils.TmemAllocator(storage.holding.ptr, barrier_for_retrieve=tmem_bar,
                                   allocator_warp_id=0, is_two_cta=False,
                                   two_cta_tmem_dealloc_mbar_ptr=storage.dealloc_bar.ptr)

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

        sQ = smem.allocate_tensor(element_type=BFloat16, layout=q_smem_layout.outer,
                                   byte_alignment=128, swizzle=q_smem_layout.inner)
        sK = smem.allocate_tensor(element_type=BFloat16, layout=k_smem_layout.outer,
                                   byte_alignment=128, swizzle=k_smem_layout.inner)
        sV = smem.allocate_tensor(element_type=BFloat16, layout=v_smem_layout.outer,
                                   byte_alignment=128, swizzle=v_smem_layout.inner)
        sC = smem.allocate_tensor(element_type=BFloat16, layout=c_smem_layout.outer,
                                   byte_alignment=128, swizzle=c_smem_layout.inner)

        # Q and K TMA partition
        gQ = cute.local_tile(mQ, cute.slice_(self.qk_mma_tiler, (None,0,None)), (None,None,None))
        gK = cute.local_tile(mK, cute.slice_(self.qk_mma_tiler, (0,None,None)), (None,None,None))
        gC = cute.local_tile(mC, cute.slice_(self.qk_mma_tiler, (None,None,0)), (None,None,None))
        k_cnt = cute.size(gQ, mode=[3])

        qk_thr = qk_mma.get_slice(0)
        tCgQ = qk_thr.partition_A(gQ)
        tCgK = qk_thr.partition_B(gK)
        tCgC = qk_thr.partition_C(gC)

        a_cta = cute.make_layout(cute.slice_(cl_vmnk, (0,0,None,0)).shape)
        tAsQ, tAgQ = cpasync.tma_partition(tma_q, 0, a_cta, cute.group_modes(sQ,0,3), cute.group_modes(tCgQ,0,3))
        b_cta = cute.make_layout(cute.slice_(cl_vmnk, (0,None,0,0)).shape)
        tAsK, tAgK = cpasync.tma_partition(tma_k, 0, b_cta, cute.group_modes(sK,0,3), cute.group_modes(tCgK,0,3))
        tAgQ = tAgQ[(None,0,None,0)]
        tAgK = tAgK[(None,0,None,0)]

        # QK MMA fragments
        tCrQ = qk_mma.make_fragment_A(sQ)
        tCrK = qk_mma.make_fragment_B(sK)

        # PV MMA fragments
        tCrV = pv_mma.make_fragment_B(sV)

        # TMEM accumulator fake for QK (scores)
        acc_shape_qk = qk_mma.partition_shape_C(self.mma_tiler_mn)
        tCtS_fake = qk_mma.make_fragment_C(cute.append(acc_shape_qk, 1))

        # TMEM accumulator fake for PV (output)
        acc_shape_pv = pv_mma.partition_shape_C(self.mma_tiler_mn)
        tCtO_fake = pv_mma.make_fragment_C(cute.append(acc_shape_pv, 1))

        pipeline.pipeline_init_wait(cluster_shape_mn=cl_vmnk)

        # ── TMA LOAD WARP ──
        if warp_idx == self.tma_warp_id:
            ab_prod.reset()
            peek = ab_prod.try_acquire()
            for kt in cutlass.range(k_cnt, unroll=1):
                h = ab_prod.acquire_and_advance(peek)
                cute.copy(tma_q, tAgQ[(None,h.count)], tAsQ[(None,h.index)], tma_bar_ptr=h.barrier)
                cute.copy(tma_k, tAgK[(None,h.count)], tAsK[(None,h.index)], tma_bar_ptr=h.barrier)
                peek = cutlass.Boolean(1)
                if h.count+1<k_cnt: peek = ab_prod.try_acquire()
            ab_prod.tail()

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

            # TMEM region 0: scores (QK MMA accumulator)
            tCtS_base = cute.make_tensor(tmem_ptr, tCtS_fake.layout)
            tCtS = tCtS_base[(None,None,None,0)]

            # TMEM region 1: output (PV MMA accumulator)
            out_ptr = cute.recast_ptr(
                tmem_ptr + self.num_tmem_cols_scores, dtype=self.acc_dtype)
            tCtO_base = cute.make_tensor(out_ptr, tCtO_fake.layout)
            tCtO = tCtO_base[(None,None,None,0)]

            # ── P fragment for PV MMA (a_source=TMEM) ──
            # Following NVIDIA fmha.py pattern:
            # 1. Create tP tensor with P TMEM layout (from make_smem_layout_A for PV MMA)
            # 2. make_fragment_A(tP) creates the A fragment with the right layout
            # 3. Shift the fragment's iterator to the P TMEM offset
            tP = cute.make_tensor(tmem_ptr, p_tmem_layout.outer)
            tOrP_base = pv_mma.make_fragment_A(tP)
            tOrP = tOrP_base[(None,None,None,0)]
            # Shift iterator to the P offset (scores region = offset 0, so no shift needed)
            # In general: tOrP_shifted = cute.make_tensor(
            #     tOrP.iterator + (self.acc_dtype.width // self.a_dtype.width) * self.tmem_p_offset,
            #     tOrP.layout)

            ab_cons.reset()
            peek = cutlass.Boolean(1)
            peek = ab_cons.try_wait()

            # ── QK MMA: Q @ K^T → tmem_scores ──
            qk_mma.set(tcgen05.Field.ACCUMULATE, False)
            for kt in range(k_cnt):
                h = ab_cons.wait_and_advance(peek)
                nblk = cute.size(tCrQ, mode=[2])
                for kb in cutlass.range(nblk, unroll_full=True):
                    crd = (None, None, kb, h.index)
                    cute.gemm(qk_mma, tCtS, tCrQ[crd], tCrK[crd], tCtS)
                h.release()
                peek = cutlass.Boolean(1)
                if h.count+1<k_cnt: peek = ab_cons.try_wait()

            # Fence TMEM writes (scores are ready)
            cute.arch.fence_view_async_tmem_store()

            # ── PV MMA: P @ V → tmem_output ──
            # A = P from TMEM (a_source=TMEM), B = V from SMEM
            pv_mma.set(tcgen05.Field.ACCUMULATE, True)
            # V fragment: slice to remove stage dimension
            tCrV_s = tCrV[(None, None, None, 0)]
            nblk_pv = cute.size(tOrP, mode=[2])
            for kb in cutlass.range(nblk_pv, unroll_full=True):
                cute.gemm(pv_mma, tCtO, tOrP[(None, None, kb)], tCrV_s[(None, None, kb)], tCtO)

            # Signal output ready
            acc_st = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, 1)
            acc_pipe.producer_acquire(acc_st)
            acc_pipe.producer_commit(acc_st)
            acc_st.advance()
            acc_pipe.producer_tail(acc_st)

        # ── EPILOGUE WARPS ──
        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)

            out_ptr = cute.recast_ptr(
                tmem_ptr + self.num_tmem_cols_scores, dtype=self.acc_dtype)
            tCtO_base = cute.make_tensor(out_ptr, tCtO_fake.layout)

            cons = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, 1)
            c_grp = pipeline.CooperativeGroup(pipeline.Agent.Thread, 128)
            c_pipe = pipeline.PipelineTmaStore.create(num_stages=2, producer_group=c_grp)

            cons = utils.gemm.sm100.epilogue_tma_store(
                self, tidx, warp_idx, tma_c, tCtO_base, sC, tCgC,
                epi_tile, 0, const_expr(lambda x: x), (0,0,0), cons, acc_pipe, c_pipe)
            c_pipe.producer_tail()
            tmem.relinquish_alloc_permit()
            tmem.free(tmem_ptr)


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

    q = torch.randn(m, k, 1, dtype=torch.bfloat16, device="cuda")
    k = 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 (K=V for now)
    qf = q[:,:,0].float()
    kf = k[:,:,0].float()
    ref = qf @ kf.T @ kf  # Q @ K^T @ V (with V=K)

    import cutlass.torch as ct
    mQ = ct.from_dlpack(q).mark_layout_dynamic(leading_dim=ct.get_leading_dim(q))
    mK = ct.from_dlpack(k).mark_layout_dynamic(leading_dim=ct.get_leading_dim(k))
    mC = ct.from_dlpack(c).mark_layout_dynamic(leading_dim=ct.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, mQ, mK, mC, stream)
    print("Running...", flush=True)
    compiled(mQ, mK, mC, stream)
    torch.cuda.synchronize()

    out = c[:,:,0].float()
    cos = torch.nn.functional.cosine_similarity(
        out.flatten().unsqueeze(0), ref.flatten().unsqueeze(0)).item()
    max_err = (out - ref).abs().max().item()
    print("Stage B: Q @ K^T @ V (a_source=TMEM, P TMEM layout from PV MMA)")
    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()
