import torch
from task import input_t, output_t

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

DEVICE = triton.runtime.driver.active.get_active_torch_device()

# 针对MI300优化的配置，增加更多针对Strassen的配置
fp8_gemm_configs = [
    # 原始配置
    Config({'BLOCK_SIZE_M': block_m, 'BLOCK_SIZE_N': block_n, 'BLOCK_SIZE_K': 128}, num_stages=num_stages, num_warps=8)
    for block_m in [32, 64, 128] for block_n in [32, 64, 128] for num_stages in [3, 4, 5]
] + [
    # Strassen优化配置 - 更大的块尺寸以便递归分解
    Config({'BLOCK_SIZE_M': 256, 'BLOCK_SIZE_N': 256, 'BLOCK_SIZE_K': 128}, num_stages=4, num_warps=8),
    Config({'BLOCK_SIZE_M': 512, 'BLOCK_SIZE_N': 512, 'BLOCK_SIZE_K': 128}, num_stages=3, num_warps=8),
    # 针对非方阵的优化配置
    Config({'BLOCK_SIZE_M': 256, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 128}, num_stages=4, num_warps=8),
    Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 256, 'BLOCK_SIZE_K': 128}, num_stages=4, num_warps=8),
]

@triton.autotune(configs=fp8_gemm_configs, key=['N', 'K'])
@triton.jit
def fp8_gemm_kernel(a_ptr, b_ptr, c_ptr,
                    a_s_ptr, b_s_ptr,
                    M, N: tl.constexpr, K: tl.constexpr,
                    BLOCK_SIZE_M: tl.constexpr,
                    BLOCK_SIZE_N: tl.constexpr,
                    BLOCK_SIZE_K: tl.constexpr):

    pid_m = tl.program_id(axis=0)
    pid_n = tl.program_id(axis=1)
    k = tl.cdiv(K, BLOCK_SIZE_K)
    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
    offs_k = tl.arange(0, BLOCK_SIZE_K)
    a_ptrs = a_ptr + offs_m[:, None] * K + offs_k[None, :]
    b_ptrs = b_ptr + offs_n[None, :] * K + offs_k[:, None]
    a_s_ptrs = a_s_ptr + offs_m * k
    b_s_ptrs = b_s_ptr + (offs_n // BLOCK_SIZE_K) * k

    accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
    for i in range(k):
        a = tl.load(a_ptrs, mask=offs_k[None, :] < K - i * BLOCK_SIZE_K, other=0.0)
        b = tl.load(b_ptrs, mask=offs_k[:, None] < K - i * BLOCK_SIZE_K, other=0.0)
        a_s = tl.load(a_s_ptrs)
        b_s = tl.load(b_s_ptrs)
        # 使用MI300的mfma指令优化的点积运算
        accumulator += tl.dot(a, b) * a_s[:, None] * b_s[None, :]
        a_ptrs += BLOCK_SIZE_K
        b_ptrs += BLOCK_SIZE_K
        a_s_ptrs += 1
        b_s_ptrs += 1
    c = accumulator.to(c_ptr.dtype.element_ty)
    offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
    offs_n = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
    c_ptrs = c_ptr + offs_m[:, None] * N + offs_n[None, :]
    mask = (offs_m[:, None] < M) & (offs_n[None, :] < N)
    tl.store(c_ptrs, c, mask=mask)


def pad_to_power_of_2(tensor, dim, target_size):
    """将张量在指定维度填充到2的幂次大小"""
    current_size = tensor.size(dim)
    if current_size >= target_size:
        return tensor
    
    pad_size = target_size - current_size
    pad_dims = [0] * (2 * tensor.dim())
    pad_dims[2 * (tensor.dim() - 1 - dim)] = pad_size
    
    return torch.nn.functional.pad(tensor, pad_dims, mode='constant', value=0)


def pad_scale_tensor(scale_tensor, original_size, target_size, scale_dim_ratio):
    """为scale张量进行相应的填充"""
    if original_size >= target_size:
        return scale_tensor
    
    scale_factor = target_size // scale_dim_ratio
    current_scale_size = original_size // scale_dim_ratio
    pad_scale_size = scale_factor - current_scale_size
    
    if pad_scale_size <= 0:
        return scale_tensor
        
    pad_dims = [0] * (2 * scale_tensor.dim())
    if scale_tensor.dim() == 2:
        pad_dims[2] = pad_scale_size  # 填充最后一个维度
    else:
        pad_dims[0] = pad_scale_size  # 填充第一个维度
    
    return torch.nn.functional.pad(scale_tensor, pad_dims, mode='constant', value=1.0)


def strassen_threshold_check(m, n, k):
    """判断是否应该使用Strassen算法"""
    # 由于FP8运算的限制，提高阈值并更保守地使用Strassen
    min_size = min(m, n, k)
    max_size = max(m, n, k)
    
    # 更高的阈值：最小维度 >= 512 且矩阵接近方阵
    return min_size >= 512 and max_size / min_size <= 2


def fp8_add_sub_safe(a, b, operation='add'):
    """安全的FP8加减运算，转换到支持的数据类型"""
    # 转换为bfloat16进行运算
    a_bf16 = a.to(torch.bfloat16)
    b_bf16 = b.to(torch.bfloat16)
    
    if operation == 'add':
        result = a_bf16 + b_bf16
    else:  # subtract
        result = a_bf16 - b_bf16
    
    # 转换回原始dtype，但保持为可计算的类型
    return result.to(torch.float8_e4m3fnuz)


def strassen_gemm_recursive(a, b, a_scale, b_scale, depth=0, max_depth=2):
    """
    递归Strassen矩阵乘法实现
    depth: 当前递归深度
    max_depth: 降低最大递归深度以避免FP8运算问题
    """
    m, k1 = a.shape[-2:]
    n, k2 = b.shape
    assert k1 == k2, f"矩阵维度不匹配: {k1} vs {k2}"
    k = k1
    
    # 递归终止条件 - 更保守的阈值
    if depth >= max_depth or not strassen_threshold_check(m, n, k) or min(m, n) <= 256:
        # 使用标准GEMM
        c = a.new_empty(*a.shape[:-1], n, dtype=torch.get_default_dtype())
        grid = lambda META: (triton.cdiv(m, META['BLOCK_SIZE_M']), triton.cdiv(n, META['BLOCK_SIZE_N']))
        
        a_flat = a.view(-1, k)
        M_flat = a_flat.size(0)
        fp8_gemm_kernel[grid](a_flat, b, c.view(-1, n), a_scale.view(-1, a_scale.size(-1)), b_scale, M_flat, n, k)
        return c.to(torch.bfloat16)
    
    # 确保维度是偶数，便于分割
    m_pad = (m + 1) // 2 * 2
    n_pad = (n + 1) // 2 * 2  
    k_pad = (k + 1) // 2 * 2
    
    # 填充矩阵
    if m < m_pad or k < k_pad:
        a_padded = pad_to_power_of_2(a, -2, m_pad)
        a_padded = pad_to_power_of_2(a_padded, -1, k_pad)
        a_scale_padded = pad_scale_tensor(a_scale, m, m_pad, 128)
        a_scale_padded = pad_scale_tensor(a_scale_padded, k, k_pad, 128)
    else:
        a_padded = a
        a_scale_padded = a_scale
        
    if n < n_pad or k < k_pad:
        b_padded = pad_to_power_of_2(b, 0, n_pad)
        b_padded = pad_to_power_of_2(b_padded, 1, k_pad)
        b_scale_padded = pad_scale_tensor(b_scale, n, n_pad, 128)
        b_scale_padded = pad_scale_tensor(b_scale_padded, k, k_pad, 128)
    else:
        b_padded = b
        b_scale_padded = b_scale
    
    # 分割矩阵
    m_half = m_pad // 2
    n_half = n_pad // 2
    k_half = k_pad // 2
    
    # 分割A矩阵
    A11 = a_padded[..., :m_half, :k_half]
    A12 = a_padded[..., :m_half, k_half:]
    A21 = a_padded[..., m_half:, :k_half]
    A22 = a_padded[..., m_half:, k_half:]
    
    # 分割B矩阵
    B11 = b_padded[:n_half, :k_half]
    B12 = b_padded[:n_half, k_half:]
    B21 = b_padded[n_half:, :k_half]
    B22 = b_padded[n_half:, k_half:]
    
    # 分割scale
    a_scale_11 = a_scale_padded[..., :m_half, :triton.cdiv(k_half, 128)]
    a_scale_12 = a_scale_padded[..., :m_half, triton.cdiv(k_half, 128):]
    a_scale_21 = a_scale_padded[..., m_half:, :triton.cdiv(k_half, 128)]
    a_scale_22 = a_scale_padded[..., m_half:, triton.cdiv(k_half, 128):]
    
    b_scale_11 = b_scale_padded[:triton.cdiv(n_half, 128), :triton.cdiv(k_half, 128)]
    b_scale_12 = b_scale_padded[:triton.cdiv(n_half, 128), triton.cdiv(k_half, 128):]
    b_scale_21 = b_scale_padded[triton.cdiv(n_half, 128):, :triton.cdiv(k_half, 128)]
    b_scale_22 = b_scale_padded[triton.cdiv(n_half, 128):, triton.cdiv(k_half, 128):]
    
    # Strassen的7个乘法运算，使用安全的加减法
    # M1 = (A11 + A22) * (B11 + B22)
    M1 = strassen_gemm_recursive(
        fp8_add_sub_safe(A11, A22, 'add'), 
        fp8_add_sub_safe(B11, B22, 'add'), 
        a_scale_11, b_scale_11, depth + 1, max_depth
    )
    
    # M2 = (A21 + A22) * B11  
    M2 = strassen_gemm_recursive(
        fp8_add_sub_safe(A21, A22, 'add'), 
        B11, 
        a_scale_21, b_scale_11, depth + 1, max_depth
    )
    
    # M3 = A11 * (B12 - B22)
    M3 = strassen_gemm_recursive(
        A11, 
        fp8_add_sub_safe(B12, B22, 'sub'), 
        a_scale_11, b_scale_12, depth + 1, max_depth
    )
    
    # M4 = A22 * (B21 - B11)
    M4 = strassen_gemm_recursive(
        A22, 
        fp8_add_sub_safe(B21, B11, 'sub'), 
        a_scale_22, b_scale_21, depth + 1, max_depth
    )
    
    # M5 = (A11 + A12) * B22
    M5 = strassen_gemm_recursive(
        fp8_add_sub_safe(A11, A12, 'add'), 
        B22, 
        a_scale_11, b_scale_22, depth + 1, max_depth
    )
    
    # M6 = (A21 - A11) * (B11 + B12)
    M6 = strassen_gemm_recursive(
        fp8_add_sub_safe(A21, A11, 'sub'), 
        fp8_add_sub_safe(B11, B12, 'add'), 
        a_scale_21, b_scale_11, depth + 1, max_depth
    )
    
    # M7 = (A12 - A22) * (B21 + B22)
    M7 = strassen_gemm_recursive(
        fp8_add_sub_safe(A12, A22, 'sub'), 
        fp8_add_sub_safe(B21, B22, 'add'), 
        a_scale_12, b_scale_21, depth + 1, max_depth
    )
    
    # 组合结果 - 结果已经是bfloat16，可以直接运算
    C11 = M1 + M4 - M5 + M7
    C12 = M3 + M5
    C21 = M2 + M4
    C22 = M1 - M2 + M3 + M6
    
    # 合并结果矩阵
    C_upper = torch.cat([C11, C12], dim=-1)
    C_lower = torch.cat([C21, C22], dim=-1)
    C_full = torch.cat([C_upper, C_lower], dim=-2)
    
    # 裁剪到原始大小
    return C_full[..., :m, :n]


def custom_kernel(data: input_t) -> output_t:
    """
    Strassen算法优化的block-scale FP8 GEMM实现
    Args:
        data: Tuple that expands to:
            a: torch.Tensor[float8_e4m3fnuz] of shape [m, k],
            b: torch.Tensor[float8_e4m3fnuz] of shape [n, k],
            a_scale: torch.Tensor[float32] of shape [m, k // 128],
            b_scale: torch.Tensor[float32] of shape [n // 128, k // 128],
            c: torch.Tensor[bfloat16] of shape [m, n]
    Returns:
        Tensor containing output in bf16
    """
    a, b, a_scale, b_scale, c = data

    # 确保张量连续性
    a_scale = a_scale.contiguous()
    b_scale = b_scale.contiguous()
    a = a.contiguous()
    b = b.contiguous()

    K = a.size(-1)
    M = a.numel() // K
    N = b.size(0)
    
    # reshape a to proper 2D matrix for batch processing
    a_2d = a.view(-1, K)  # [M, K]
    
    # 判断是否使用Strassen算法
    if strassen_threshold_check(M, N, K):
        # 使用Strassen算法
        result = strassen_gemm_recursive(a_2d, b, a_scale, b_scale)
    else:
        # 使用标准GEMM（原始实现）
        result = a.new_empty(M, N, dtype=torch.get_default_dtype())
        grid = lambda META: (triton.cdiv(M, META['BLOCK_SIZE_M']), triton.cdiv(N, META['BLOCK_SIZE_N']))
        fp8_gemm_kernel[grid](a_2d, b, result, a_scale.view(-1, a_scale.size(-1)), b_scale, M, N, K)
    
    # 恢复原始形状
    original_shape = list(a.shape[:-1]) + [N]
    result = result.view(original_shape)
    
    return result.to(torch.bfloat16)
python eval.py test /tmp/tmptcyn5odk