"""
D1.3 SMEM-P: Direct SMEM write test.
Write known values to sP via coordinate indexing,
then copy sP to GMEM and verify.
"""
import torch, math
import cutlass, cutlass.cute as cute, cutlass.utils as utils
from cutlass.cute.nvgpu import tcgen05, cpasync
from cutlass import Float32, BFloat16, Int32, const_expr
from cutlass.utils import LayoutEnum
import cutlass.torch as ct
import cuda.bindings.driver as cuda


@cute.jit
def smem_write_test(q, k, v, c, stream):
    """Write known values to sP, copy to GMEM, verify."""
    a_major = LayoutEnum.from_tensor(q).mma_major_mode()
    b_major = LayoutEnum.from_tensor(k).mma_major_mode()
    v_fmha = cute.make_tensor(
        v.iterator,
        cute.make_layout((64, 128, 1), stride=(1, 64, 64 * 128)),
    )
    v_major = LayoutEnum.from_tensor(v_fmha).mma_major_mode()

    qk_mma = utils.sm100.make_trivial_tiled_mma(
        BFloat16, BFloat16, a_major, b_major, Float32,
        tcgen05.CtaGroup.ONE, (128, 128), tcgen05.OperandSource.SMEM
    )
    pv_mma = utils.sm100.make_trivial_tiled_mma(
        BFloat16, BFloat16, cute.nvgpu.OperandMajorMode.K, v_major, Float32,
        tcgen05.CtaGroup.ONE, (128, 64), tcgen05.OperandSource.SMEM
    )

    pv_mma_tiler = (128, 64, 128)
    qk_mma_tiler = (128, 128, 128 * 4)

    p_smem_s = utils.sm100.make_smem_layout_a(pv_mma, pv_mma_tiler, BFloat16, 1)
    q_smem_s = utils.sm100.make_smem_layout_a(qk_mma, qk_mma_tiler, BFloat16, 1)
    k_smem_s = utils.sm100.make_smem_layout_b(qk_mma, qk_mma_tiler, BFloat16, 2)
    v_smem_s = utils.sm100.make_smem_layout_b(pv_mma, pv_mma_tiler, BFloat16, 2)

    # TMA for reading back sP
    epi_s = cute.select(p_smem_s, mode=[0, 1])  # 2D view of sP for TMA
    # Actually, let's just use a simple output tensor
    # We'll write known values to sP, then copy to output tensor c

    q_s = cute.slice_(q_smem_s, (None, None, None, 0))
    k_s = cute.slice_(k_smem_s, (None, None, None, 0))
    v_s = cute.slice_(v_smem_s, (None, None, None, 0))
    cta = cute.size(qk_mma.thr_id.shape)

    tma_q, mQ = cute.nvgpu.make_tiled_tma_atom_A(
        utils.sm100.cluster_shape_to_tma_atom_A((1, 1), qk_mma.thr_id),
        q, q_s, qk_mma_tiler, qk_mma, (1, 1, 1, 1)
    )
    tma_k, mK = cute.nvgpu.make_tiled_tma_atom_B(
        utils.sm100.cluster_shape_to_tma_atom_B((1, 1), qk_mma.thr_id),
        k, k_s, qk_mma_tiler, qk_mma, (1, 1, 1, 1)
    )
    tma_v, mV = cute.nvgpu.make_tiled_tma_atom_B(
        utils.sm100.cluster_shape_to_tma_atom_B((1, 1), pv_mma.thr_id),
        v_fmha, v_s, pv_mma_tiler, pv_mma, (1, 1, 1, 1)
    )

    # Output: use c as a flat buffer to read back sP values
    # c has shape (128, 64, 1) — same as sP's logical size
    # We'll TMA-store sP to c
    c_smem_s = utils.sm100.make_smem_layout_a(pv_mma, pv_mma_tiler, BFloat16, 1)
    epi_tile = (128, 64)
    epi_s2 = cute.select(c_smem_s, mode=[0, 1])
    tma_c, mC = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(), c, epi_s2, epi_tile)

    # Just use 128 threads (4 warps) for simplicity
    _kernel(qk_mma, pv_mma, tma_q, mQ, tma_k, mK, tma_v, mV, tma_c, mC,
            p_smem_s, q_smem_s, k_smem_s, v_smem_s, c_smem_s, epi_tile).launch(
        grid=(1, 1, 1), block=[128, 1, 1], stream=stream
    )


@cute.kernel
def _kernel(qk_mma, pv_mma, tma_q, mQ, tma_k, mK, tma_v, mV, tma_c, mC,
            p_smem_s, q_smem_s, k_smem_s, v_smem_s, c_smem_s, epi_tile):
    tidx, _, _ = cute.arch.thread_idx()
    warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx())

    @cute.struct
    class SS:
        q_bar: cute.struct.MemRange[cutlass.Int64, 2]
        tmem_dealloc: cutlass.Int64
        holding: cutlass.Int32
    smem = utils.SmemAllocator()
    st = smem.allocate(SS)

    sP = smem.allocate_tensor(element_type=BFloat16, layout=p_smem_s.outer, byte_alignment=128, swizzle=p_smem_s.inner)
    sC = smem.allocate_tensor(element_type=BFloat16, layout=c_smem_s.outer, byte_alignment=128, swizzle=c_smem_s.inner)

    # All 128 threads write known values to sP
    # Strategy: each thread writes its own portion of sP
    # using a simple pattern: value = thread_id (cast to BF16)
    # This tests that the SMEM write addressing works correctly

    # For this test, each thread writes to sP using the same coordinate mapping
    # as the FMHA kernel. But we don't have tTMEM_LOADcS here.
    # Instead, let's use a simpler approach: directly write sequential values.

    # Actually, let's just write sP using the MMA fragment A partition
    # to verify that write-then-read works.
    pv_thr = pv_mma.get_slice(0)
    sP_stage = sP[(None, None, None, 0)]

    # Write: each thread writes its portion using MMA's A-operand partition
    tCrP = pv_mma.make_fragment_A(sP)
    # tCrP is the MMA warp's register fragment for reading sP.
    # For writing, we need the "store" side.
    # Actually, make_fragment_A creates a load fragment, not a store fragment.

    # Simpler test: just have each thread write a known value to sP directly
    # using coordinate indexing with a simple loop.
    # Each thread writes 128 values (one row) to sP.
    # Thread t writes to row t (for t in 0..127).
    if tidx < 128:
        m = tidx
        for k in range(128):
            k0 = k % 16
            k1 = (k // 16) % 4
            k2 = k // 64
            # Write the linear index as BF16: value = (m * 128 + k) % 256
            val = BFloat16(float((m * 128 + k) % 256))
            sP_stage[(m, k0), 0, (k1, k2)] = val

    cute.arch.fence_proxy("async.shared", space="cta")
    # Barrier to ensure all writes are visible
    bar = pipeline.NamedBarrier(barrier_id=5, num_threads=128)
    bar.arrive_and_wait()

    # Now copy sP to sC (same layout), then TMA store to GMEM
    # sP and sC have the same layout, so we can copy directly
    # Use the TMA store path
    gC = cute.local_tile(mC, cute.slice_((128, 64), (None, 0)), (None, None))
    tCgC = pv_thr.partition_C(gC)

    # Copy sP to sC
    sC_stage = sC[(None, None, None, 0)]
    for m in range(128):
        for k in range(128):
            k0 = k % 16
            k1 = (k // 16) % 4
            k2 = k // 64
            # Only thread 0 does the copy (simple but slow)
            if tidx == 0:
                sC_stage[(m, k0), 0, (k1, k2)] = sP_stage[(m, k0), 0, (k1, k2)]

    cute.arch.fence_proxy("async.shared", space="cta")
    bar.arrive_and_wait()

    # TMA store sC to GMEM
    if tidx == 0:
        cpasync.copy_tma_g2s(tma_c, sC, gC)  # Wrong direction, need s2g
        # Actually, for TMA store (SMEM→GMEM), we need cpasync.copy
        # Let me just use a direct store instead

    # Actually this is getting too complicated. Let me use a simpler approach.
    # Write the sP values to GMEM directly using a simple loop from thread 0.


def test_smem_write():
    print("=== SMEM-P Direct Write Test ===\n")
    hd = 64; s_k = 128
    q = torch.randn(128, hd, 1, dtype=torch.bfloat16, device='cuda')
    k = torch.randn(s_k, hd, 1, dtype=torch.bfloat16, device='cuda')
    v = torch.randn(s_k, hd, dtype=torch.bfloat16, device='cuda')
    c = torch.zeros(128, hd, 1, dtype=torch.bfloat16, device='cuda')

    stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream)
    v_tile = v.unsqueeze(-1)
    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))
    mV = ct.from_dlpack(v_tile).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v_tile))
    mC = ct.from_dlpack(c).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c))

    print('This test is too complex. Let me take a different approach.', flush=True)


if __name__ == '__main__':
    test_smem_write()