from tvm import DataType
from tvm.runtime import convert
from tvm.tirx import const
import tilelang.language as T


def shared_16x4_to_local_64x1_layout_A(i, j):
    thread_id = j * 16 + i
    return thread_id, const(0)


def thread_id_shared_access_64x1_to_16x4_layout_A(thread_id, local_id):
    i = thread_id % 16
    j = thread_id // 16
    return i, j


def shared_4x16_to_local_64x1_layout_B(i, j):
    thread_id = i * 16 + j
    return thread_id, const(0)


def thread_id_shared_access_64x1_to_4x16_layout_B(thread_id, local_id):
    i = thread_id // 16
    j = thread_id % 16
    return i, j


def shared_16x16_to_local_64x4_layout_C(i, j):
    thread_id = j + (i // 4) * 16
    local = i % 4
    return thread_id, local


def shared_16x16_to_ldmatrix_64x4_layout(ind):
    i, j = ind[0], ind[1]
    thread_id, local_id = shared_16x16_to_local_64x4_layout_C(i, j)
    return convert([thread_id, local_id])


def thread_id_shared_access_64x4_to_16x16_layout_A(thread_id, local_id):
    i = thread_id % 16
    j = (thread_id // 16) * 4 + local_id
    return i, j


def shared_16x16_to_local_64x4_layout_A(i, j):
    thread_id = i + 16 * (j // 4)
    local = j % 4
    return thread_id, local


def thread_id_shared_access_64x4_to_16x16_layout_B(thread_id, local_id):
    i = local_id + (thread_id // 16) * 4
    j = thread_id % 16
    return i, j


def shared_16x16_to_local_64x4_layout_B(i, j):
    thread_id = j + (i // 4) * 16
    local = i % 4
    return thread_id, local


shared_16x16_to_local_64x4_layout_m_n = shared_16x16_to_local_64x4_layout_A
shared_16x16_to_local_64x4_layout_n_k = shared_16x16_to_local_64x4_layout_A
shared_16x16_to_local_64x4_layout_n_m = shared_16x16_to_local_64x4_layout_B
shared_16x16_to_local_64x4_layout_k_n = shared_16x16_to_local_64x4_layout_B


def thread_id_shared_access_64x4_to_16x16_layout_C_m_n(thread_id, local_id):
    i = local_id + (thread_id // 16) * 4
    j = thread_id % 16
    return i, j


def thread_id_shared_access_64x4_to_16x16_layout_C_n_m(thread_id, local_id):
    i = thread_id % 16
    j = local_id + (thread_id // 16) * 4
    return i, j


def thread_id_shared_access_64x8_to_16x32_layout_A(thread_id, local_id):
    i = thread_id % 16
    j = (thread_id // 16) * 8 + local_id
    return i, j


def shared_16x32_to_local_64x8_layout_A(i, j):
    thread_id = i + 16 * (j // 8)
    local = j % 8
    return thread_id, local


def thread_id_shared_access_64x8_to_16x32_layout_B(thread_id, local_id):
    i = local_id + (thread_id // 16) * 8
    j = thread_id % 16
    return i, j


def shared_16x32_to_local_64x8_layout_B(i, j):
    thread_id = j + (i // 8) * 16
    local = i % 8
    return thread_id, local


def thread_id_shared_access_64x16_to_16x64_layout_A(thread_id, local_id):
    i = thread_id % 16
    j = local_id + (thread_id // 16) * 16
    return i, j


def shared_16x64_to_local_64x16_layout_A(i, j):
    thread_id = i + 16 * (j // 16)
    local = j % 16
    return thread_id, local


def thread_id_shared_access_64x16_to_16x64_layout_B(thread_id, local_id):
    i = local_id + (thread_id // 16) * 16
    j = thread_id % 16
    return i, j


def shared_16x64_to_local_64x16_layout_B(i, j):
    thread_id = i + 16 * (j // 16)
    local = j % 16
    return thread_id, local


def shared_32x32_to_local_64x16_layout_C(i, j):
    # i=M_row, j=N_col. For v_mfma_i32_32x32x32_i8: tid%32=M_row, tid//32 selects N group.
    tid_high = (j // 4) % 2  # 0 for j in {0-3,8-11,...}, 1 for j in {4-7,12-15,...}
    thread_id = i + 32 * tid_high
    k = j - 4 * tid_high
    local_id = (k // 8) * 4 + k % 4
    return thread_id, local_id


def thread_id_shared_access_64x16_to_32x32_layout_C_n_m(thread_id, local_id):
    # Returns (row=M, col=N) for v_mfma_i32_32x32x32_i8 output layout.
    # tid%32 = M_row, (tid//32)*4 + lid%4 + (lid//4)*8 = N_col.
    i = thread_id % 32
    j = (thread_id // 32) * 4 + local_id % 4 + (local_id // 4) * 8
    return i, j


def thread_id_shared_access_64x16_to_32x32_layout_C_m_n(thread_id, local_id):
    """Return (m, n) = (row, col) for the 32x32 MFMA output register layout.

    For v_mfma_i32_32x32x32_i8 (gfx950), each wave-64 lane holds 16 output
    i32 values.  The row (M-dimension) is indexed by ``thread_id % 32``
    and the column (N-dimension) is given by the interleaved formula below.
    This function returns ``(m_idx, n_idx)`` matching the ``(row, col)``
    convention expected by ``stmatrix``.
    """
    m = thread_id % 32
    n = (thread_id // 32) * 4 + local_id % 4 + (local_id // 4) * 8
    return m, n


def shared_32x32_to_local_64x16_layout_A(i, j):
    thread_id = i + 32 * (j // 16)
    local_id = j % 16
    return thread_id, local_id


def thread_id_shared_access_64x16_to_32x32_layout_A(thread_id, local_id):
    i = thread_id % 32
    j = (thread_id // 32) * 16 + local_id
    return i, j


def shared_32x32_to_local_64x16_layout_B(i, j):
    thread_id = j + 32 * (i // 16)
    local_id = i % 16
    return thread_id, local_id


def thread_id_shared_access_64x16_to_32x32_layout_B(thread_id, local_id):
    i = (thread_id // 32) * 16 + local_id
    j = thread_id % 32
    return i, j


def make_mfma_swizzle_layout(shared_buf, vecSize=8):
    dtype = shared_buf.dtype
    shape = shared_buf.shape

    numBanks = 32
    bankBitWidth = 32
    SIMDWidth = 16

    innerDimLength = shape[-1]
    typeWidthInBit = DataType(dtype).bits

    elemsPerOneBanksRow = (numBanks * bankBitWidth) // typeWidthInBit
    perPhase = max(1, elemsPerOneBanksRow // innerDimLength)
    maxPhase = min(SIMDWidth // perPhase, innerDimLength // vecSize)

    def transform(row, col):
        phase = (row // perPhase) % maxPhase
        colOffSwizzled = ((col // vecSize) ^ phase) * vecSize
        colOffOrdered = col % vecSize
        colOff = colOffSwizzled + colOffOrdered
        return row, colOff

    return T.Layout(shape, transform)
