import torch
import triton
import triton.language as tl
from triton import Config

from task import input_t, output_t


g_num_warps = 8
g_num_stages = 2
g_waves_per_eu = 0
g_BLOCK_SIZE_M = 32
g_BLOCK_SIZE_N = 32
g_BLOCK_SIZE_K = 128
g_GROUP_SIZE_M = 8

g_plot_name = "transpose_persistent"


def setup_meta(num_warps, num_stages, waves_per_eu, BLOCK_SIZE_M, BLOCK_SIZE_N, BLOCK_SIZE_K, GROUP_SIZE_M):
    global g_num_warps
    global g_num_stages
    global g_waves_per_eu
    global g_BLOCK_SIZE_M
    global g_BLOCK_SIZE_N
    global g_BLOCK_SIZE_K
    global g_GROUP_SIZE_M

    g_num_warps = num_warps
    g_num_stages = num_stages
    g_waves_per_eu = waves_per_eu
    g_BLOCK_SIZE_M = BLOCK_SIZE_M
    g_BLOCK_SIZE_N = BLOCK_SIZE_N
    g_BLOCK_SIZE_K = BLOCK_SIZE_K
    g_GROUP_SIZE_M = GROUP_SIZE_M


@triton.jit
def remap_xcd(pid, GRID_MN, NUM_XCDS: tl.constexpr = 8):
    ## pid remapping on xcds
    # Number of pids per XCD in the new arrangement
    pids_per_xcd = (GRID_MN + NUM_XCDS - 1) // NUM_XCDS
    # When GRID_MN cannot divide NUM_XCDS, some xcds will have
    # pids_per_xcd pids, the other will have pids_per_xcd - 1 pids.
    # We calculate the number of xcds that have pids_per_xcd pids as
    # tall_xcds
    tall_xcds = GRID_MN % NUM_XCDS
    tall_xcds = NUM_XCDS if tall_xcds == 0 else tall_xcds
    # Compute current XCD and local pid within the XCD
    xcd = pid % NUM_XCDS
    local_pid = pid // NUM_XCDS
    # Calculate new pid based on the new grouping
    # Note that we need to consider the following two cases:
    # 1. the current pid is on a tall xcd
    # 2. the current pid is on a short xcd
    if xcd < tall_xcds:
        pid = xcd * pids_per_xcd + local_pid
    else:
        pid = tall_xcds * pids_per_xcd + (xcd - tall_xcds) * (pids_per_xcd - 1) + local_pid

    return pid


@triton.jit
def pid_grid(pid, num_pid_m, num_pid_n, GROUP_SIZE_M: tl.constexpr = 1):
    if GROUP_SIZE_M == 1:
        pid_m = pid // num_pid_n
        pid_n = pid % num_pid_n
    else:
        num_pid_in_group = GROUP_SIZE_M * num_pid_n
        group_id = pid // num_pid_in_group
        first_pid_m = group_id * GROUP_SIZE_M
        group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M)
        pid_m = first_pid_m + (pid % group_size_m)
        pid_n = (pid % num_pid_in_group) // group_size_m

    return pid_m, pid_n


@triton.jit
def fp8_gemm_kernel(a_ptr, b_ptr, c_ptr,
                    a_s_ptr, b_s_ptr,
                    M: tl.constexpr, N: tl.constexpr, K: tl.constexpr,
                    BLOCK_SIZE_M: tl.constexpr,
                    BLOCK_SIZE_N: tl.constexpr,
                    BLOCK_SIZE_K: tl.constexpr,
                    GROUP_SIZE_M: tl.constexpr,
                    NUM_SMS: tl.constexpr,):
    start_pid = tl.program_id(axis=0)
    NUM_XCDS: tl.constexpr = 8

    num_pid_m = tl.cdiv(M, BLOCK_SIZE_M)
    num_pid_n = tl.cdiv(N, BLOCK_SIZE_N)
    num_tiles = num_pid_m * num_pid_n
    tile_id = start_pid

    offs_k = tl.arange(0, BLOCK_SIZE_K)
    num_block_k = tl.cdiv(K, BLOCK_SIZE_K)

    b_s_step = tl.cdiv(N, 128)

    num_valid_tiles = tl.cdiv((num_tiles - tile_id), NUM_SMS)
    for _ in range(0, num_valid_tiles):
        tile_id_remapped = remap_xcd(tile_id, num_tiles, NUM_XCDS)
        pid_m, pid_n = pid_grid(tile_id_remapped, num_pid_m, num_pid_n, GROUP_SIZE_M)

        offs_m = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M
        offs_n = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)) % N
        a_s_ptrs = a_s_ptr + offs_m
        b_s_ptrs = b_s_ptr + (offs_n // 128)

        a_block_ptr = tl.make_block_ptr(
            base=a_ptr,
            shape=(M, K),
            strides=(1, M),
            offsets=(pid_m * BLOCK_SIZE_M, 0),
            block_shape=(BLOCK_SIZE_M, BLOCK_SIZE_K),
            order=(0, 1)
        )
        b_block_ptr = tl.make_block_ptr(
            base=b_ptr,
            shape=(K, N),
            strides=(N, 1),
            offsets=(0, pid_n * BLOCK_SIZE_N),
            block_shape=(BLOCK_SIZE_K, BLOCK_SIZE_N),
            order=(1, 0)
        )

        accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
        for k in range(num_block_k):
            a = tl.load(a_block_ptr)
            b = tl.load(b_block_ptr)
            a_s = tl.load(a_s_ptrs)
            b_s = tl.load(b_s_ptrs)
            accumulator += tl.dot(a, b) * a_s[:, None] * b_s[None, :]
            a_block_ptr = tl.advance(a_block_ptr, (0, BLOCK_SIZE_K))
            b_block_ptr = tl.advance(b_block_ptr, (BLOCK_SIZE_K, 0))
            a_s_ptrs += M
            b_s_ptrs += b_s_step
        acc = accumulator.to(c_ptr.dtype.element_ty)
        c_ptrs = c_ptr + offs_m[:, None] * N + offs_n[None, :]
        mask = (offs_m[:, None] < M) & (offs_n[None, :] < N)
        tl.store(c_ptrs, acc, mask=mask)

        tile_id += NUM_SMS


def custom_kernel(data: input_t) -> output_t:
    a, b, a_scale, b_scale, c = data
    K = a.size(-1)
    M = a.size(0)
    N = b.size(0)

    # BLOCK_SIZE_K = 128
    # BLOCK_SIZE_M = 128
    # BLOCK_SIZE_N = 64
    # GROUP_SIZE_M = 8
    # num_stages = 2
    # num_warps = 8
    # waves_per_eu = 2

    num_warps = g_num_warps
    num_stages = g_num_stages
    waves_per_eu = g_waves_per_eu
    BLOCK_SIZE_M = g_BLOCK_SIZE_M
    BLOCK_SIZE_N = g_BLOCK_SIZE_N
    BLOCK_SIZE_K = g_BLOCK_SIZE_K
    GROUP_SIZE_M = g_GROUP_SIZE_M

    NUM_SMS = torch.cuda.get_device_properties("cuda").multi_processor_count * 2
    grid = lambda META: (
        min(NUM_SMS,
            triton.cdiv(M, META["BLOCK_SIZE_M"]) * triton.cdiv(N, META["BLOCK_SIZE_N"])), )
    fp8_gemm_kernel[grid](a, b, c, a_scale, b_scale, M, N, K,
        BLOCK_SIZE_M=BLOCK_SIZE_M,
        BLOCK_SIZE_N=BLOCK_SIZE_N,
        BLOCK_SIZE_K=BLOCK_SIZE_K,
        GROUP_SIZE_M=GROUP_SIZE_M,
        NUM_SMS=NUM_SMS,
        num_stages=num_stages,
        num_warps=num_warps,
        waves_per_eu=waves_per_eu,
        kpack=2,
        matrix_instr_nonkdim=16,)

    return c

python eval.py test /tmp/tmpctc6f2m3