# SPDX-License-Identifier: MIT
# Copyright (C) 2024-2025, Advanced Micro Devices, Inc. All rights reserved.

import functools
from typing import Optional

import pandas as pd
import torch
from torch import Tensor

from aiter import logger
from aiter.jit.utils.torch_guard import torch_compile_guard

from ..jit.core import AITER_CONFIGS, AITER_LOG_TUNED_CONFIG, compile_ops
from ..jit.utils.chip_info import get_cu_num, get_gfx
from ..ops.gemm_op_common import get_padded_m
from ..utility import dtypes


@functools.lru_cache(maxsize=1024)
def compute_gemm_SplitK(M: int, N: int, K: int, tile_m: int, tile_n: int, tile_k: int):
    cu_num = get_cu_num()
    tile_num = ((M + tile_m - 1) // tile_m) * ((N + tile_n - 1) // tile_n)
    cusPerTile = cu_num / tile_num
    splitK = 0
    while cusPerTile >= pow(2, splitK + 1) and (pow(2, splitK + 1) * tile_k) < 2 * K:
        splitK += 1
    ## to make sure the precision is not lost, max is 4
    # return min(splitK, 4)
    return 3


@functools.lru_cache(maxsize=1024)
def get_GEMM_config(M: int, N: int, K: int):
    if not hasattr(get_GEMM_config, "gemm_dict"):
        gemm_dict = pd.read_csv(
            AITER_CONFIGS.AITER_CONFIG_GEMM_A4W4_FILE
        ).drop_duplicates()
        get_GEMM_config.gemm_dict = gemm_dict.set_index(
            ["cu_num", "M", "N", "K"]
        ).to_dict("index")
    cu_num = get_cu_num()
    padded_M = M
    config = None
    for gl in [None, 0, 1]:
        padded_M = M if gl is None else get_padded_m(M, N, K, gl)
        config = get_GEMM_config.gemm_dict.get((cu_num, padded_M, N, K), None)
        if config is not None:
            if AITER_LOG_TUNED_CONFIG:
                logger.info(
                    f"shape is M:{M}, N:{N}, K:{K}, found padded_M: {padded_M}, N:{N}, K:{K} is tuned on cu_num = {cu_num} in {AITER_CONFIGS.AITER_CONFIG_GEMM_A4W4_FILE}, kernel name is {config['kernelName']}, splitK is {config['splitK']}!"
                )
            break
    else:
        logger.info(
            f"shape is M:{M}, N:{N}, K:{K}, not found tuned config in CKGEMM or asmGEMM, will use default config!"
        )
    return config


def gemm_a4w4_fake(
    A: Tensor,  # A:[M, K/2] f4x2
    B: Tensor,  # B:[N, K/2] f4x2
    A_scale: Tensor,  # A_scale:[M, K/32] e8m0 paded
    B_scale: Tensor,  # B_scale:[N, K/32] e8m0 paded
    bias: Optional[Tensor] = None,  # bias:[1, N] f32
    dtype: torch.dtype = dtypes.bf16,
    alpha: Optional[float] = 1.0,
    beta: Optional[float] = 0.0,
    bpreshuffle: Optional[bool] = True,
) -> torch.Tensor:
    m = A.numel() // A.shape[-1]
    n = B.shape[0]
    out = torch.empty((m, n), dtype=dtype, device=A.device)
    return out


@torch_compile_guard(gen_fake=gemm_a4w4_fake)
def gemm_a4w4(
    A: Tensor,  # A:[M, K/2] f4x2
    B: Tensor,  # B:[N, K/2] f4x2
    A_scale: Tensor,  # A_scale:[M, K/32] e8m0 paded
    B_scale: Tensor,  # B_scale:[N, K/32] e8m0 paded
    bias: Optional[Tensor] = None,  # bias:[1, N] f32
    dtype: torch.dtype = dtypes.bf16,
    alpha: Optional[float] = 1.0,
    beta: Optional[float] = 0.0,
    bpreshuffle: Optional[bool] = True,
) -> torch.Tensor:
    """
    A4W4 GEMM kernel for AMD GPUs.
    This function is a wrapper for the A4W4 GEMM kernel.
    It is used to perform matrix multiplication with 4-bit quantization.
    """
    # Load the A4W4 GEMM kernel
    m = A.numel() // A.shape[-1]
    n = B.shape[0]
    k = A.shape[-1] * 2
    out = torch.empty(((m + 31) // 32 * 32, n), dtype=dtype, device=A.device)
    gfx_arch = get_gfx()
    if gfx_arch in ["gfx942"]:
        raise RuntimeError(
            f"A4W4 GEMM kernel is not supported on gfx942, but got {gfx_arch}!"
        )
    ck_config = get_GEMM_config(m, n, k)
    # splitK = None
    splitK = 0
    kernelName = ""
    if ck_config is not None:
        splitK = ck_config.get("splitK", None)
        kernelName = ck_config["kernelName"]
    if (
        ck_config is not None
        and kernelName.find("_ZN") == -1
        # or bias is None
    ):
        splitK = 0 if splitK is None else splitK
        return gemm_a4w4_blockscale(
            A.view(m, k // 2), B, A_scale, B_scale, out, splitK=splitK
        )[:m]
    assert (
        out.shape[0] % 32 == 0
    ), "Dim0 of gemm_a4w4_asm output needs to be padded to multiples of 32!"
    gemm_a4w4_asm(
        A.view(m, k // 2),
        B,
        A_scale,
        B_scale,
        out,
        kernelName,
        bias,
        alpha,
        beta,
        bpreshuffle,
        log2_k_split=splitK,
    )
    return out[:m].view(*A.shape[:-1], n)


def gen_gemm_a4w4_asm_fake_tensors(
    A: Tensor,  # A:[M, K/2] f4x2
    B: Tensor,  # B:[N, K/2] f4x2
    A_scale: Tensor,  # A_scale:[M, K/32] e8m0 paded
    B_scale: Tensor,  # B_scale:[N, K/32] e8m0 paded
    out: Tensor,  # Out:[M, N] bf16
    kernelName: str,
    bias: Optional[Tensor] = None,  # bias:[1, N] f32
    alpha: Optional[float] = 1.0,
    beta: Optional[float] = 0.0,
    bpreshuffle: Optional[bool] = True,
    log2_k_split: Optional[int] = None,
) -> Tensor:
    return out


@compile_ops("module_gemm_a4w4_asm", gen_fake=gen_gemm_a4w4_asm_fake_tensors)
def gemm_a4w4_asm(
    A: Tensor,  # A:[M, K/2] f4x2
    B: Tensor,  # B:[N, K/2] f4x2
    A_scale: Tensor,  # A_scale:[M, K/32] e8m0 paded
    B_scale: Tensor,  # B_scale:[N, K/32] e8m0 paded
    out: Tensor,  # Out:[M, N] bf16
    kernelName: str,
    bias: Optional[Tensor] = None,  # bias:[1, N] f32
    alpha: Optional[float] = 1.0,
    beta: Optional[float] = 0.0,
    bpreshuffle: Optional[bool] = True,
    log2_k_split: Optional[int] = None,
) -> Tensor: ...


def gen_gemm_a4w4_blockscale_fake_tensors(
    XQ: torch.Tensor,
    WQ: torch.Tensor,
    x_scale: torch.Tensor,
    w_scale: torch.Tensor,
    Out: torch.Tensor,
    splitK: int = 0,
) -> torch.Tensor:
    return Out


@compile_ops(
    "module_gemm_a4w4_blockscale", gen_fake=gen_gemm_a4w4_blockscale_fake_tensors
)
def gemm_a4w4_blockscale(
    XQ: torch.Tensor,
    WQ: torch.Tensor,
    x_scale: torch.Tensor,
    w_scale: torch.Tensor,
    Out: torch.Tensor,
    splitK: int = 0,
) -> Tensor: ...


@compile_ops(
    "module_gemm_a4w4_blockscale_tune",
    fc_name="gemm_a4w4_blockscale_tune",
    gen_fake=gen_gemm_a4w4_blockscale_fake_tensors,
)
def gemm_a4w4_blockscale_tune(
    XQ: torch.Tensor,
    WQ: torch.Tensor,
    x_scale: torch.Tensor,
    w_scale: torch.Tensor,
    Out: torch.Tensor,
    kernelId: int,
    splitK: int = 0,
) -> Tensor: ...
