from task import input_t, output_t
import triton
import triton.language as tl
import torch
from itertools import product
from functools import lru_cache
from typing import List


@triton.jit
def pid2mn(
    pid: int,
    m: int,
    n: int,
    pid2mn_mode: tl.constexpr,
    BLOCK_M: tl.constexpr,
    BLOCK_N: tl.constexpr,
):
    if pid2mn_mode == 0:
        # M Major
        # grid: (cdiv(m, BLOCK_M) * cdiv(n, BLOCK_N), 1, 1)
        tile_id = pid
        
        m_block = tl.cdiv(m, BLOCK_M)
        n_block = tl.cdiv(n, BLOCK_N)
        
        start_m = (tile_id % m_block) * BLOCK_M
        start_n = (tile_id // m_block) * BLOCK_N
        return start_m, start_n
    elif pid2mn_mode == 1:
        # N Major
        # grid: (cdiv(m, BLOCK_M) * cdiv(n, BLOCK_N), 1, 1)
        tile_id = pid
        
        m_block = tl.cdiv(m, BLOCK_M)
        n_block = tl.cdiv(n, BLOCK_N)
        
        start_n = (tile_id % n_block) * BLOCK_N
        start_m = (tile_id // n_block) * BLOCK_M
        return start_m, start_n
    else:
        return 0, 0

@triton.jit
def scale_product_matmul(
    a,
    b,
    a_scale,
    b_scale,
    c,
    m,
    n,
    k,
    BLOCK_M: tl.constexpr,
    BLOCK_N: tl.constexpr,
    BLOCK_K: tl.constexpr,
):
    # Layout
    # a: (m, k) : (k, 1)
    # b: (n, k) : (k, 1)
    # a_scale: (m, k // 128) : (1, m)
    # b_scale: (n // 128, k // 128) : (1, n // 128)
    # c: (m, n) : (n, 1)
    # simple_pid2_tileid grid: (cdiv(m, BLOCK_M) * cdiv(n, BLOCK_N), 1, 1)
    
    pid = tl.program_id(axis=0)
    pid2mn_mode: tl.constexpr = 0
    # 0: M Major
    # 1: N Major
    start_m, start_n = pid2mn(pid, m, n, pid2mn_mode, BLOCK_M, BLOCK_N)  
    
    # a: (m, k) : (k, 1)
    a_stride_m = k
    a_stride_k: tl.constexpr = 1
    
    # b: (n, k) : (k, 1)
    b_stride_n = k
    b_stride_k: tl.constexpr = 1
    
    # a_scale: (m, k // 128) : (1, m)
    a_scale_stride_m: tl.constexpr = 1
    a_scale_stride_k = m
    
    # b_scale: (n // 128, k // 128) : (1, n // 128)
    b_scale_stride_n: tl.constexpr = 1
    b_scale_stride_k = tl.cdiv(n, 128)
    
    # c: (m, n) : (n, 1)
    c_stride_m = n
    c_stride_n: tl.constexpr = 1
    
    # a_block_ptr: (BLOCK_M, BLOCK_K)
    a_block_ptr = tl.make_block_ptr(
        a,
        shape=(m, k),
        strides=(a_stride_m, a_stride_k),
        offsets=(start_m, 0),
        block_shape=(BLOCK_M, BLOCK_K),
        order=(1, 0),
    )
    
    # b_block_ptr: (BLOCK_N, BLOCK_K)
    b_block_ptr = tl.make_block_ptr(
        b,
        shape=(n, k),
        strides=(b_stride_n, b_stride_k),
        offsets=(start_n, 0),
        block_shape=(BLOCK_N, BLOCK_K),
        order=(1, 0),
    )
    
    # c_block_ptr: (BLOCK_M, BLOCK_N)
    c_block_ptr = tl.make_block_ptr(
        c,
        shape=(m, n),
        strides=(c_stride_m, c_stride_n),
        offsets=(start_m, start_n),
        block_shape=(BLOCK_M, BLOCK_N),
        order=(1, 0),
        # order2=(0, 1),
    )
    acc_type: tl.constexpr = tl.float32
    acc_block = tl.zeros((BLOCK_M, BLOCK_N), dtype=acc_type)
    for start_k in tl.range(0, k, BLOCK_K):
        tl.static_assert(BLOCK_K <= 128)
        tl.static_assert(BLOCK_N % 128 == 0 or 128 % BLOCK_N == 0)
        # SCALE_BLOCK_K = tl.constexpr(BLOCK_K // 128 if BLOCK_K >= 128 else 1)
        SCALE_BLOCK_K: tl.constexpr = 1
        SCALE_BLOCK_N: tl.constexpr = (BLOCK_N + 127) // 128
        # a_scale_block_ptr: (BLOCK_M, BLOCK_K // 128)
        a_scale_block_ptr = tl.make_block_ptr(
            a_scale,
            shape=(m, tl.cdiv(k, 128)),
            strides=(a_scale_stride_m, a_scale_stride_k),
            offsets=(start_m, start_k // 128),
            block_shape=(BLOCK_M, SCALE_BLOCK_K),
            order=(0, 1),
        )
        # b_scale_block_ptr: (BLOCK_N // 128, BLOCK_K // 128)
        b_scale_block_ptr = tl.make_block_ptr(
            b_scale,
            shape=(tl.cdiv(n, 128), tl.cdiv(k, 128)),
            strides=(b_scale_stride_n, b_scale_stride_k),
            offsets=(start_n // 128, start_k // 128),
            block_shape=(SCALE_BLOCK_N, SCALE_BLOCK_K),
            order=(0, 1),
        )
        
        a_block = tl.load(a_block_ptr, boundary_check=(0, 1)) # (BLOCK_M, BLOCK_K)
        b_block = tl.load(b_block_ptr, boundary_check=(0, 1)) # (BLOCK_N, BLOCK_K)
        a_scale_block = tl.load(a_scale_block_ptr, boundary_check=(0, 1)) # (BLOCK_M, SCALE_BLOCK_K)
        b_scale_block = tl.load(b_scale_block_ptr, boundary_check=(0, 1)) # (SCALE_BLOCK_N, SCALE_BLOCK_K)

        cur_acc_block = tl.zeros((BLOCK_M, BLOCK_N), dtype=acc_type)
        cur_acc_block = tl.dot(a_block, b_block.T) # (BLOCK_M, BLOCK_N)
        
        # scale cur_acc_block
        scale_block = a_scale_block.reshape(BLOCK_M, 1) * b_scale_block.reshape(1, SCALE_BLOCK_N) # (BLOCK_M, SCALE_BLOCK_N)
        scale_block = scale_block.reshape((BLOCK_M, SCALE_BLOCK_N, 1)) # (BLOCK_M, SCALE_BLOCK_N, 1)

        scaled_acc = cur_acc_block.reshape((BLOCK_M, SCALE_BLOCK_N, BLOCK_N // SCALE_BLOCK_N)) * scale_block # (BLOCK_M, SCALE_BLOCK_N, BLOCK_N // SCALE_BLOCK_N)
        scaled_acc = scaled_acc.reshape((BLOCK_M, BLOCK_N)) # (BLOCK_M, BLOCK_N)
        #  
        acc_block += scaled_acc # (BLOCK_M, BLOCK_N)

        a_block_ptr = tl.advance(a_block_ptr, (0, BLOCK_K))
        b_block_ptr = tl.advance(b_block_ptr, (0, BLOCK_K))
    
    tl.store(c_block_ptr, acc_block.to(tl.bfloat16), boundary_check=(0, 1))
    
    

def custom_kernel(data: input_t) -> output_t:
    """
    Reference implementation of block-scale fp8 gemm
    Args:
        data: Tuple that expands to:
            a: torch.Tensor[float8_e4m3fnuz] of shape [m, k], stride [1, m]
            b: torch.Tensor[float8_e4m3fnuz] of shape [n, k], stride [1, n]
            a_scale: torch.Tensor[float32] of shape [m, k // 128], stride [1, m]
            b_scale: torch.Tensor[float32] of shape [n // 128, k // 128], stride [1, n // 128]
            c: torch.Tensor[bfloat16] of shape [m, n]
    Returns:
        Tensor containing output in bf16
    """
    
    a, b, a_scale, b_scale, c = data
    a = a.contiguous()
    b = b.contiguous()
    # Your implementation here
    m,n,k = a.shape[0], b.shape[0], a.shape[1]
    grid = lambda meta: (
        triton.cdiv(m, meta["BLOCK_M"]) * triton.cdiv(n, meta["BLOCK_N"]),
    )
    kernel = scale_product_matmul[grid](
        a, b, a_scale, b_scale, c, m, n, k,
        BLOCK_M=128,
        BLOCK_N=128,
        BLOCK_K=32,
        # num_stages=2,
        # num_warps=4,
    )
    
    return c


python eval.py test /tmp/tmp23j4oprp