# SPDX-License-Identifier: MIT
import torch
import torch.nn.functional as F
import sys
import os
from aiter import dtypes
from aiter.test_common import checkAllclose, perftest, benchmark
from einops import rearrange
from einops import repeat as eirp
from aiter import logger
import pandas as pd
from aiter.blaslt_scale_mm import scale_mm
# from aiter.ops.triton.gemm_a8w8 import gemm_a8w8
from aiter.ops.triton.gemm_w8a8 import gemm_w8a8

from vllm.model_executor.layers.quantization.utils.int8_utils import per_token_group_quant_int8
from vllm.model_executor.layers.quantization.utils.fp8_utils import per_token_group_quant_fp8

from aiter import pertoken_quant

import math

if 1:
    _path = os.path.abspath(os.path.dirname(__file__))
    sys.path.insert(0, f"{_path}/../../")
    from aiter.tuned_gemm import tgemm

block_shape = (128, 128)

@perftest(num_iters=5)
def run_torch(x, weight, x_scale, w_scale, dtype=dtypes.bf16):
    block_shape_n, block_shape_k = block_shape
    m, k = x.shape
    n = weight.shape[0]
    scale_n = (n + block_shape_n - 1) // block_shape_n
    scale_k = (k + block_shape_k - 1) // block_shape_k
    x = x.to(x_scale.dtype).view(
        m, k // block_shape[1], block_shape[1]
    ) * x_scale.unsqueeze(-1)
    x = x.view(m, k)

    w_scale = rearrange(
        w_scale.view(-1, 1)
        .repeat(1, block_shape_n * block_shape_k)
        .view(scale_n, scale_k, block_shape_n, block_shape_k),
        "num_blk_n num_blk_k blk_n blk_k -> (num_blk_n blk_n) (num_blk_k blk_k)",
    )
    w_scale = w_scale[:n, :k]
    weight = weight.to(w_scale.dtype) * w_scale

    out = F.linear(x.to(dtypes.fp32), weight.to(dtypes.fp32))
    return out.to(dtype)

@perftest()
def run_gemm_b(x, weight, scaleA=None, scaleB=None, otype=None, bias=None, scaleType=None):
    #交换AB 位置
    fn = torch.compile(scale_mm, backend="inductor", fullgraph=True)
    return fn(x, weight, bias, otype, scaleA, scaleB, scale_type=scaleType)

@perftest()
def run_triton(x, weight, x_scale, w_scale, block_size, dtype=torch.bfloat16):
    # return gemm_a8w8(x, weight, x_scale, w_scale, bias, dtype, y)
    return gemm_w8a8(x, weight, x_scale, w_scale, block_size, dtype)

@benchmark()
def test_gemm(quant_dtype, dtype, m, n, k):
    dim = (m, n, k)
    block_shape_n, block_shape_k = block_shape
    scale_n = (n + block_shape_n - 1) // block_shape_n
    scale_k = (k + block_shape_k - 1) // block_shape_k

    x = torch.randn((m, k), dtype=torch.float16, device="cuda")
    if quant_dtype == dtypes.fp8:
        x, x_scale = per_token_group_quant_fp8(x, block_shape_k)
    elif quant_dtype == dtypes.i8:
        x, x_scale = per_token_group_quant_int8(x, block_shape_k)
    else:
        assert False, f"{quant_dtype=} not supported"
    # w = torch.randint(-128, 127, (n, k), dtype=torch.int8, device="cuda")
    # w_scale = torch.rand([scale_n, scale_k], dtype=dtypes.fp32, device="cuda")

    block_shape_n_aligned = math.ceil(n / block_shape_n) * block_shape_n
    w2 = torch.rand((1, block_shape_n_aligned, k), dtype=dtype, device="cuda")
    tmp = rearrange(
        w2.view(
            -1,
            w2.shape[1] // block_shape_n,
            block_shape_n,
            math.ceil(w2.shape[2] / block_shape_k),
            block_shape_k,
        ),
        "e num_blk_n blk_n num_blk_k blk_k -> e num_blk_n num_blk_k (blk_n blk_k)",
    ).contiguous()
    w2_qweight, w2_scales = pertoken_quant(tmp, quant_dtype=quant_dtype)
    w = rearrange(
        w2_qweight.view(
            -1,
            w2.shape[1] // block_shape_n,
            w2.shape[2] // block_shape_k,
            block_shape_n,
            block_shape_k,
        ),
        "e num_blk_n num_blk_k blk_n blk_k -> e (num_blk_n blk_n) (num_blk_k blk_k)",
    ).contiguous()

    w = w.view(w.shape[1], w.shape[2])
    # cut to orignal shape
    w = w[:n]
    w_scale = w2_scales.view(w2_scales.shape[1], w2_scales.shape[2])

    # print(f'{x.shape=} {x.dtype=} {w.shape=} {w_scale.shape=} {w_scale.dtype=}')

    # a, avg_a = run_torch(x, w, x_scale, w_scale, dtype)
    #这里替换ck/blaslt 实现
    b, avg_b = run_gemm_b(x, w, x_scale, w_scale, dtype, scaleType=2)

    # triton
    c, avg_c = run_triton(x, w, x_scale, w_scale, block_size=block_shape, dtype=dtype)

    # msg = f"[perf] dim: {str(dim):<20} dtype: {dtype}, torch avg: {avg_a:<8.2f} us, blaslt avg: {avg_b:<8.2f} us, uplift: {avg_a/avg_b -1:<5.1%}"
    # checkAllclose(a, b, msg="a,b: " + msg, rtol=1e-2, atol=0.01)

    tflops = lambda us: 2 * m * n * k * 1e-12 / (us * 1e-6)

    # check triton correctness
    # msg = f"[perf] dim: {str(dim):<20} dtype: {dtype}, triton avg: {avg_c:<8.2f} us, blaslt avg: {avg_b:<8.2f} us"
    msg = f"[perf] dim: {str(dim):<20} dtype: {dtype}, triton avg: {tflops(avg_c):<8.2f} TFLOPS, blaslt avg: {tflops(avg_b):<8.2f} TFLOPS"
    checkAllclose(b, c, msg="b,c: " + msg, rtol=1e-2, atol=0.01)
    return {"blaslt": avg_b, "triton": avg_c}


df = []
for dtype in [dtypes.bf16]:
# for dtype in [dtypes.fp16]:
    # deepseek-r1
    for m in [1, 2, 4, 8, 16, 32, 64]:
    # for m in [16]:
        for n, k in [
            # (1536, 7168),
            # (3072, 1536),
            # (576, 7168),
            # (7168, 256),
            # (7168, 2048),
            # (4608, 7168),
            # (7168, 2304),
            (7168, 2048),
            # (7168, 2304),
            # (512, 7168),
            # (4096, 512),
        ]:
            for quant_dtype in [dtypes.i8]:
            # for quant_dtype in [dtypes.i8]:
                ret = test_gemm(quant_dtype, dtype, m, n, k)
                df.append(ret)
df = pd.DataFrame(df)
logger.info(f"summary:\n{df}")