import torch
import time
from typing import Optional, Type,Any, Dict, List, Tuple
import pandas as pd
import os
import json
import triton
import triton.language as tl
import pandas as pd
from transformer_engine.pytorch.triton.per_token_group_quant import _int8_gemm_helper_b
import functools
import logging

logger = logging.getLogger(__name__)

device_name=torch.cuda.get_device_properties('cuda').name.replace(" ","_")


tuning_full_space = False
# tuning_full_space = True

def get_full_tuning_space():
    configs = []
    if not tuning_full_space:
        return configs

    block_m_range = [16, 32, 64]
    block_n_range = [16, 32, 64, 128]
    block_k_range = [32, 64, 128]
    num_warps_range = [4, 8]
    group_m_range = [2, 4, 8]

    # For now we see better perf with num_stages=0 for all gemm configs we care
    # But keep this explicit so that we do not forget we may need to set it to
    # other values in the future
    num_stage_range = [0, 1, 2]

    for block_m in block_m_range:
        for block_n in block_n_range:
            for block_k in block_k_range:
                for num_warps in num_warps_range:
                    for group_m in group_m_range:
                            for num_stages in num_stage_range:
                                configs.append(triton.Config({'BLOCK_SIZE_M': block_m, 'BLOCK_SIZE_N': block_n, 'BLOCK_SIZE_K': block_k, 'GROUP_SIZE_M': group_m}, num_stages=num_stages, num_warps=num_warps, enable_mmacfuse=2))

    return configs


@triton.autotune(
    configs= get_full_tuning_space() if tuning_full_space else [
        # triton.Config({'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 128, 'GROUP_SIZE_M': 2, 'kpack':2}, num_stages=2, num_warps=8),
        triton.Config({'BLOCK_SIZE_M': 16, 'BLOCK_SIZE_N': 64, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 2,}, num_stages=1, num_warps=4, enable_mmacfuse=2),
        triton.Config({'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 128, 'GROUP_SIZE_M': 8,}, num_stages=1, num_warps=4, enable_mmacfuse=2),
    ],
    key=['M', 'N', 'K'],
    # reset_to_zero=['c_ptr']
)

@triton.jit
def _w8a8_block_int8_matmul(
    # Pointers to inputs and output
    A,
    B,
    C,
    As,
    Bs,
    # Shape for matmul
    M,
    N,
    K,
    # Block size for block-wise quantization
    group_n,
    group_k,
    # Stride for inputs and output
    stride_am,
    stride_ak,
    stride_bk,
    stride_bn,
    stride_cm,
    stride_cn,
    stride_As_m,
    stride_As_k,
    stride_Bs_k,
    stride_Bs_n,
    # Meta-parameters
    BLOCK_SIZE_M: tl.constexpr,
    BLOCK_SIZE_N: tl.constexpr,
    BLOCK_SIZE_K: tl.constexpr,
    GROUP_SIZE_M: tl.constexpr,
):
    """Triton-accelerated function used to perform linear operations (dot
    product) on input tensors `A` and `B` with block-wise quantization,
    and store the result in output tensor `C`.
    """

    pid = tl.program_id(axis=0)
    num_pid_m = tl.cdiv(M, BLOCK_SIZE_M)
    num_pid_n = tl.cdiv(N, BLOCK_SIZE_N)
    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

    offs_am = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M
    offs_bn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)) % N

    offs_bsn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)) % N
    # offs_bsn = pid_n * BLOCK_SIZE_N // group_n

    offs_k = tl.arange(0, BLOCK_SIZE_K)
    a_ptrs = A + (offs_am[:, None] * stride_am + offs_k[None, :] * stride_ak)
    b_ptrs = B + (offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn)


    As_ptrs = As + offs_am * stride_As_m
    # offs_bsn = offs_bn // group_n
    Bs_ptrs = Bs + offs_bsn * stride_Bs_n

    accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
    for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)):

        k_start = k * BLOCK_SIZE_K
        offs_ks = k_start // group_k
        a_s = tl.load(As_ptrs + offs_ks * stride_As_k)
        b_s = tl.load(Bs_ptrs + offs_ks * stride_Bs_k)

        a = tl.load(a_ptrs,
                    mask=offs_k[None, :] < K - k * BLOCK_SIZE_K,
                    other=0.0)
        b = tl.load(b_ptrs,
                    mask=offs_k[:, None] < K - k * BLOCK_SIZE_K,
                    other=0.0)


        accumulator += tl.dot(a, b).to(tl.float32) * a_s[:, None] * b_s[None, :]
        
        a_ptrs += BLOCK_SIZE_K * stride_ak
        b_ptrs += BLOCK_SIZE_K * stride_bk


    if C.dtype.element_ty == tl.bfloat16:
        c = accumulator.to(tl.bfloat16)
    elif C.dtype.element_ty == tl.float16:
        c = accumulator.to(tl.float16)
    else:
        c = accumulator.to(tl.float32)

    offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
    offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
    c_ptrs = C + stride_cm * offs_cm[:, None] + stride_cn * offs_cn[None, :]
    c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N)
    tl.store(c_ptrs, c, mask=c_mask)


@functools.lru_cache
def get_w8a8_block_int8_configs(
    N: int, K: int, block_n: int, block_k: int
) -> Optional[Dict[int, Any]]:
    """
    Return optimized configurations for the w8a8 block fp8 kernel.
    The return value will be a dictionary that maps an irregular grid of
    batch sizes to configurations of the w8a8 block fp8 kernel. To evaluate the
    kernel on a given batch size bs, the closest batch size in the grid should
    be picked and the associated configuration chosen to invoke the kernel.
    """

    # First look up if an optimized configuration is available in the configs
    # directory
    device_name=torch.cuda.get_device_properties('cuda').name.replace(" ","_")
    json_file_name = f"N={N},K={K},device_name={device_name},dtype=int8_w8a8,block_shape=[{block_n}, {block_k}].json" # noqa: E501

    config_file_path = os.path.join(
        os.path.dirname(os.path.realpath(__file__)), "configs", json_file_name
    )
    if os.path.exists(config_file_path):
        with open(config_file_path) as f:
            logger.info(
                "Using configuration from %s for W8A8 Block INT8 kernel.",
                config_file_path,
            )
            # If a configuration has been found, return it
            return {int(key): val for key, val in json.load(f).items()}

    # If no optimized configuration is available, we will use the default
    # configuration
    logger.warning(
        (
            "Using default W8A8 Block INT8 kernel config. Performance might "
            "be sub-optimal! Config file not found at %s"
        ),
        config_file_path,
    )
    return None


def w8a8_block_int8_matmul_wgrad(
    A: torch.Tensor,
    B: torch.Tensor,
    As: torch.Tensor,
    Bs: torch.Tensor,
    block_size: List[int],
    output_dtype: torch.dtype = torch.float16,
    best_config:Optional[dict]=None
) -> torch.Tensor:
    """matrix multiplication with block-wise quantization.
    It takes two input tensors `A` and `B` with scales `As` and `Bs`.
    The output is returned in the specified `output_dtype`.
    Args:
        A: The input tensor, e.g., activation.
        B: The input tensor, e.g., weight.
        As: The per-token-group quantization scale for `A`.
        Bs: The per-block quantization scale for `B`.
        block_size: The block size for per-block quantization. It should
        be 2-dim, e.g., [128, 128].
        output_dytpe: The dtype of the returned tensor.
    Returns:
        torch.Tensor: The result of matmul.
    """
    assert len(block_size) == 2
    block_n, block_k = block_size[0], block_size[1]

    assert A.shape[-1] == B.shape[-1]
    # print(f"A.shape[:-1] : {A.shape[:-1]}, As.shape[:-1]: {As.shape[:-1]}")
    # assert A.shape[:-1] == As.shape[:-1] and A.is_contiguous()
    assert triton.cdiv(A.shape[-1], block_k) == As.shape[0]
    M = A.numel() // A.shape[-1]

    assert B.ndim == 2 and B.is_contiguous() and Bs.ndim == 2
    N, K = B.shape
    # assert triton.cdiv(N, block_n) == Bs.shape[0]
    # assert triton.cdiv(K, block_k) == Bs.shape[1]

    C_shape = A.shape[:-1] + (N,)
    C = A.new_empty(C_shape, dtype=output_dtype)

    # configs = get_w8a8_block_int8_configs(N, K, block_size[0], block_size[1])
    # if configs:
    #     # If an optimal configuration map has been found, look up the
    #     # optimal config
    #     config = configs[min(configs.keys(), key=lambda x: abs(x - M))]
    # else:
    if best_config:
        config=best_config
        
    else: 
        #print("best config has not found!")
        # config = {
        #         "BLOCK_SIZE_M": 32, #64
        #         "BLOCK_SIZE_N": block_size[0],
        #         "BLOCK_SIZE_K": block_size[1],
        #         "GROUP_SIZE_M": 32,
        #         "num_warps": 4,
        #         "num_stages": 3,
        #     }
        # Default config
        # Block-wise quant: BLOCK_SIZE_K must be divisable by block_size[1]
        #print("block_size[0]:{},block_size[1]:{}".format(block_size[0],block_size[1]))
        if M<=64:
            config = {
                "BLOCK_SIZE_M": 16, #64
                "BLOCK_SIZE_N":block_size[0],
                "BLOCK_SIZE_K": block_size[1],
                "GROUP_SIZE_M": 2,
                "num_warps": 4,
                "num_stages": 0,
            }
        elif M<128:
            config = {
                "BLOCK_SIZE_M": 32, #64
                "BLOCK_SIZE_N": block_size[0],
                "BLOCK_SIZE_K": block_size[1],
                "GROUP_SIZE_M": 2,
                "num_warps": 4,
                "num_stages": 0,
            }   
        elif M<=256:
            config = {
                "BLOCK_SIZE_M": 64, #64
                "BLOCK_SIZE_N": block_size[0],
                "BLOCK_SIZE_K": block_size[1],
                "GROUP_SIZE_M": 2,
                "num_warps": 4,
                "num_stages": 0,
            }                     
        else :
            config = {
                "BLOCK_SIZE_M": 64, #64
                "BLOCK_SIZE_N": block_size[0],
                "BLOCK_SIZE_K": block_size[1],
                "GROUP_SIZE_M": 8,
                "num_warps": 8,
                "num_stages": 0,
            } 

    def grid(META):
        return (
            triton.cdiv(M, META["BLOCK_SIZE_M"]) *
            triton.cdiv(N, META["BLOCK_SIZE_N"]),
        )
        
    #print("config:",config)
    # print(f"zhenggf, A.shape:{A.shape}, B.shape:{B.shape}")
    # print(f"zhenggf, A.stride(-2):{A.stride(-2)}, A.stride(-1):{A.stride(-1)}, B.stride(1):{B.stride(1)}, B.stride(0):{B.stride(0)}")
    # print(f"zhenggf, As.stride(-2):{As.stride(-2)}, As.stride(-1):{As.stride(-1)}, Bs.stride(1):{Bs.stride(1)}, Bs.stride(0):{Bs.stride(0)}")
    # print(f"zhenggf, As.stride(-2):{As.stride(-2)}, As.stride(-1):{As.stride(-1)}, Bs.stride(1):{Bs.stride(1)}, Bs.stride(0):{Bs.stride(0)}")
    # As = As.permute(1, 0).contiguous()
    _w8a8_block_int8_matmul[grid](
        A,
        B,
        C,
        As,
        Bs,
        M,
        N,
        K,
        block_n,
        block_k,
        A.stride(-2),
        A.stride(-1),
        B.stride(1),
        B.stride(0),
        C.stride(-2),
        C.stride(-1),
        # As.stride(-2),
        # As.stride(-1),
        As.stride(1),
        As.stride(0),
        Bs.stride(-2),
        Bs.stride(-1),
        # Bs.stride(1),
        # Bs.stride(0),
        # **config,
    )
    config = _w8a8_block_int8_matmul.best_config

    return C,config 




def apply_w8a8_block_int8_linear_helper(m: int,
                                            n: int,
                                            k: int,
                                            out_dtype: Type[torch.dtype] = torch.float16,
                                            device: str = "cuda",
                                            block_size: List[int]=[128,128],
                                            bias: Optional[torch.Tensor] = None,
                                            best_config:Optional[dict]=None):
    
    
    q_input, x_scale,weight,weight_scale=_int8_gemm_helper_b(m=m,n=n,k=k,out_dtype=out_dtype,device=device,block_size=block_size)
    print(f"zhenggf, q_input:{q_input.shape}, x_scale:{x_scale.shape}, weight:{weight.shape}, weight_scale:{weight_scale.shape}")
    
    torch_output = native_w8a8_block_int8_matmul(q_input, weight, x_scale, weight_scale, block_size)

    x_scale = x_scale.permute(1, 0).contiguous()
    weight_scale = weight_scale.permute(1, 0).contiguous()

    print(f"zhenggf 转置后传递给triton kernel, q_input:{q_input.shape}, x_scale:{x_scale.shape}, weight:{weight.shape}, weight_scale:{weight_scale.shape}")
    output,config = w8a8_block_int8_matmul_wgrad(
        q_input, weight, x_scale, weight_scale, block_size,
        output_dtype=out_dtype,
        best_config=best_config
    )


    
    if not torch.allclose(output, torch_output, rtol=1e-2, atol=5e-2):
        print("triton 精度检查不合格！！！") 
        
    else:
        print("triton 精度检查合格") 
    #  unit test end
    g = torch.cuda.CUDAGraph()
    with torch.cuda.graph(g):
        for it in range(1000):
            output,_ = w8a8_block_int8_matmul_wgrad(
                    q_input, weight, x_scale, weight_scale, block_size,
                    output_dtype=out_dtype,
                    best_config=best_config
                )
    torch.cuda.synchronize()            
    start_time_ = time.time()  # 开始计时
    g.replay()
    torch.cuda.synchronize()    
    end_time_ = time.time()  # 结束计时
    
    elapsed_time = round((end_time_ - start_time_) *1000 ,7)# 计算耗时  
    print("_time:{} us\n".format(elapsed_time))
    quantiles = [0.5, 0.2, 0.8]
    gpu_costtime = triton.testing.do_bench(lambda:w8a8_block_int8_matmul_wgrad(q_input, weight, x_scale, weight_scale, block_size,output_dtype=out_dtype,best_config=best_config),quantiles=None, return_mode="mean")*1000
    
    if bias is not None:
        output = output + bias
    return output.to(dtype=out_dtype),elapsed_time,gpu_costtime,config

def get_triton_cache(file_path,n,k,block_n,block_k):
    #会将所报错的json文件以字典的形式return出来
    #先读取指定的文件，该文件地址不存在则会读默认路径

    cache_json_file=file_path
    
    if os.path.exists(file_path):
    #try:
        with open(cache_json_file, 'r') as file:
            cachedata = json.load(file)
    else:
        cachedata = {}
        # 写入空数据到新的JSON文件
        with open(file_path, 'w') as file:
            json.dump(cachedata, file)   
                 
    #把所有的cache解析成key:config的形式：[M_N_K]:[config]
    configs_dict={}
    for key, value in cachedata.items():
        for sub_key, sub_value in value.items():
            configs_key= f"{sub_key}_{key}"
            configs_value={
                'BLOCK_SIZE_M': int(sub_value["BLOCK_SIZE_M"]),
                'BLOCK_SIZE_N': int(sub_value["BLOCK_SIZE_N"]),
                'BLOCK_SIZE_K': int(sub_value["BLOCK_SIZE_K"]),
                'GROUP_SIZE_M': int(sub_value["GROUP_SIZE_M"]),
                'num_stages':int(sub_value['num_stages']),
                'num_warps':int(sub_value['num_warps']),
                # 'kpack':int(sub_value['kpack']),
                'enable_mmacfuse':int(2),
            }
            configs_dict[configs_key]=configs_value
    return configs_dict

def getspec_config(configs_dict,m,n,k,block_n,block_k):
    
    if f"{m}_{n}_{k}_block[{block_n},{block_k}]" in configs_dict:
        return configs_dict[f"{m}_{n}_{k}_block[{block_n},{block_k}]"]
    else:
        return None  

# For test
def native_w8a8_block_int8_matmul(A, B, As, Bs, block_size, output_dtype=torch.bfloat16):
    """This function performs matrix multiplication with block-wise quantization using native torch.

    It takes two input tensors `A` and `B` with scales `As` and `Bs`.
    The output is returned in the specified `output_dtype`.
    """

    A = A.to(torch.float32)
    B = B.to(torch.float32)
    assert A.shape[-1] == B.shape[-1]
    assert B.ndim == 2 and B.is_contiguous() and Bs.ndim == 2
    assert len(block_size) == 2
    block_n, block_k = block_size[0], block_size[1]
    assert (A.shape[-1] + block_k - 1) // block_k == As.shape[-1]
    assert A.shape[:-1] == As.shape[:-1]

    M = A.numel() // A.shape[-1]
    N, K = B.shape
    origin_C_shape = A.shape[:-1] + (N,)
    A = A.reshape(M, A.shape[-1])
    As = As.reshape(M, As.shape[-1])
    n_tiles = (N + block_n - 1) // block_n
    k_tiles = (K + block_k - 1) // block_k
    # assert n_tiles == Bs.shape[0]
    assert k_tiles == Bs.shape[1]

    C_shape = (M, N)
    C = torch.zeros(C_shape, dtype=torch.float32, device=A.device)

    A_tiles = [A[:, i * block_k : min((i + 1) * block_k, K)] for i in range(k_tiles)]
    B_tiles = [
        [
            B[
                j * block_n : min((j + 1) * block_n, N),
                i * block_k : min((i + 1) * block_k, K),
            ]
            for i in range(k_tiles)
        ]
        for j in range(n_tiles)
    ]
    C_tiles = [C[:, j * block_n : min((j + 1) * block_n, N)] for j in range(n_tiles)]
    As_tiles = [As[:, i : i + 1] for i in range(k_tiles)]
    Bs_tiles = [Bs[:, i : i + 1] for i in range(k_tiles)]


    for i in range(k_tiles):
        for j in range(n_tiles):
            a = A_tiles[i]
            b = B_tiles[j][i]
            c = C_tiles[j]
            s = As_tiles[i] * Bs_tiles[i].t()[:,j*block_n:min((j + 1) * block_n, N)]
            c[:, :] += torch.matmul(a, b.t()) * s

    C = C.reshape(origin_C_shape).to(output_dtype)
    return C

def main():
    m1=[item if item < 17 else  1 << (item - 27) for item in range(1, 17)]
    m2=[item<<2  if item <17 else (item - 8)<<3  for item in range(5, 29)]
    m3=[2<<(item)  for item in range(7, 13)]
        
    m_list = m1+m2+m3
    n_list=[576,2048,7168,256,7168,1536,1536,2304,7168]
    k_list=[7168,512,1024,7168,128,7168,1536,7168,1152]

    m_list = [8192]
    n_list=[7168]
    k_list=[1152]

    block_size=[128, 128]
    
    out_dtype=torch.bfloat16
    
    _n=[]
    _k=[]
    _m=[]
    
    _configs_block_m=[]
    _configs_block_n=[]
    _configs_block_k=[]
    _configs_block_group_m=[]
    _configs_block_num_warps=[]
    _configs_block_num_stages=[]
    _configs_kpack=[]
    
    cost_times=[]
    gpu_costtimes=[]
    device_name=torch.cuda.get_device_properties('cuda').name.replace(" ", "_")
    
    for i in range(0,len(k_list),1):
        
        for m in m_list:
            print("m:{} n:{} k:{} ".format(m,n_list[i],k_list[i]))

            best_config = []
                
                       
            output,elapsed_time,gpu_costtime,config=apply_w8a8_block_int8_linear_helper(m=m,n=n_list[i],k=k_list[i],block_size=block_size,out_dtype=out_dtype,best_config=best_config)

            cost_times.append(elapsed_time)
            gpu_costtimes.append(gpu_costtime)
            _n.append(n_list[i])
            _k.append(k_list[i])
            _m.append(m)
            print(f"zhenggf, {config}")
            print(f"zhenggf, {config.kwargs}")
            _configs_block_m.append(config.kwargs['BLOCK_SIZE_M'])
            _configs_block_n.append(config.kwargs['BLOCK_SIZE_N'])
            _configs_block_k.append(config.kwargs['BLOCK_SIZE_K'])
            _configs_block_group_m.append(config.kwargs['GROUP_SIZE_M'])
            _configs_block_num_warps.append(config.num_warps)
            _configs_block_num_stages.append(config.num_stages)
            # _configs_kpack.append(config['kpack'])
            
    
    # 创建一个包含这三个列表的 DataFrame
    df = pd.DataFrame({'m':_m,'n':_n,'k':_k,'线性层gemm量化算子耗时': cost_times,'GPU算子耗时':gpu_costtimes,
                        'BLOCK_SIZE_M':_configs_block_m,'BLOCK_SIZE_N':_configs_block_n,'BLOCK_SIZE_K':_configs_block_k,
                        'GROUP_SIZE_M':_configs_block_group_m,'num_warps':_configs_block_num_warps,'num_stages':_configs_block_num_stages,#'kpack':_configs_kpack
                        })
    # 将 DataFrame 写入 Excel 文件
    df.to_excel('gemmoutput.xlsx', index=False)

    print("表格已保存到 gemmoutput.xlsx 文件中。")    

if __name__ == "__main__":
    main()

