from .dense_blockscaled_gemm_persistent_amax import (
    Sm100BlockScaledPersistentDenseGemmKernel,
)

from cuda.bindings import driver as cuda
import os
import torch
from typing import Tuple, Optional

import cutlass
import cutlass.cute as cute
from cutlass.cute.runtime import make_fake_stream

from cudnn.datatypes import _convert_to_cutlass_data_type
from cudnn.api_base import APIBase, TensorDesc, TupleDict, is_power_of_2, ceil_div


class GemmAmaxSm100(APIBase):
    def __init__(
        self,
        sample_a: torch.Tensor,
        sample_b: torch.Tensor,
        sample_sfa: torch.Tensor,
        sample_sfb: torch.Tensor,
        sample_c: torch.Tensor,
        sample_amax: torch.Tensor,
        acc_dtype: torch.dtype = torch.float32,
        mma_tiler_mn: Tuple[int, int] = (128, 128),
        cluster_shape_mn: Tuple[int, int] = (1, 1),
        sf_vec_size: int = 32,
    ):
        super().__init__()

        self._warn_experimental_api()
        self._logger.debug("Entering __init__")

        self.a_desc = self._make_tensor_desc(sample_a, name="sample_a")
        self.b_desc = self._make_tensor_desc(sample_b, name="sample_b")
        self.sfa_desc = self._make_tensor_desc(sample_sfa, name="sample_sfa")
        self.sfb_desc = self._make_tensor_desc(sample_sfb, name="sample_sfb")
        self.c_desc = self._make_tensor_desc(sample_c, name="sample_c")
        self.amax_desc = self._make_tensor_desc(sample_amax, name="sample_amax")

        self.acc_dtype = acc_dtype
        self.mma_tiler_mn = mma_tiler_mn
        self.cluster_shape_mn = cluster_shape_mn
        self.sf_vec_size = sf_vec_size

        self.num_cluster_overlap_margin = int(os.getenv("CUDNNFE_CLUSTER_OVERLAP_MARGIN", "0"))
        self._logger.debug(f"setting num_cluster_overlap_margin: {self.num_cluster_overlap_margin}")

        # used to reshape sfa/sfb tensors to atom layout
        self.atom_m = (32, 4)
        self.atom_k = 4

        self._interpret_uint8_as_fp4x2 = True
        self._logger.debug(
            f"__init__ completed with args: sample_a {self.a_desc.shape}, sample_b {self.b_desc.shape}, sample_sfa {self.sfa_desc.shape}, sample_sfb {self.sfb_desc.shape}, sample_c {self.c_desc.shape}, sample_amax {self.amax_desc.shape}, acc_dtype {acc_dtype}, mma_tiler_mn {mma_tiler_mn}, cluster_shape_mn {cluster_shape_mn}, sf_vec_size {sf_vec_size}"
        )

    def check_support(self) -> bool:
        self._logger.debug("Entering check_support")

        self._logger.debug("Checking dtypes and sf_vec_size")
        ab_dtype = self._check_dtype(
            self.a_desc,
            dtype=[torch.float4_e2m1fn_x2, torch.uint8, torch.float8_e5m2, torch.float8_e4m3fn],
            name="A",
        )
        self._check_dtype(
            self.b_desc,
            dtype=ab_dtype,
            name="B",
            extra_error_msg="A and B tensor dtypes must match",
        )
        if ab_dtype == torch.uint8:
            self._logger.warning("Uint8 ab_dtype will be interpreted as packed fp4, not as native uint8")

        self._value_error_if(
            self.sf_vec_size not in {16, 32},
            f"Unsupported sf_vec_size: received {self.sf_vec_size}, expected {{16, 32}}",
        )

        sf_dtype = self._check_dtype(
            self.sfa_desc,
            dtype=[torch.float8_e8m0fnu, torch.float8_e4m3fn, torch.int8],
            name="sfa",
        )
        self._check_dtype(
            self.sfb_desc,
            dtype=sf_dtype,
            name="sfb",
            extra_error_msg="sfa and sfb tensor dtypes must match",
        )
        if sf_dtype == torch.int8:
            self._logger.warning("Int8 sf_dtype will be interpreted as float8_e8m0fnu, not as native int8")

        self._value_error_if(
            sf_dtype == torch.float8_e4m3fn and self.sf_vec_size == 32,
            "Unsupported sf_dtype and sf_vec_size combination: float8_e4m3fn and 32 is not supported",
        )
        self._value_error_if(
            ab_dtype in {torch.float8_e5m2, torch.float8_e4m3fn} and self.sf_vec_size == 16,
            f"Unsupported ab_dtype and sf_vec_size combination: {{float8_e5m2, float8_e4m3fn}} and 16 is not supported",
        )

        c_dtype = self._check_dtype(
            self.c_desc,
            dtype=[torch.float32, torch.float16, torch.bfloat16, torch.float8_e5m2, torch.float8_e4m3fn, torch.float4_e2m1fn_x2, torch.uint8],
            name="C",
        )
        self._value_error_if(
            self._is_fp4x2(c_dtype) and not self._is_fp4x2(ab_dtype),
            f"Unsupported c_dtype and ab_dtype combination: fp4 c_dtype requires fp4 ab_dtype, got {ab_dtype}",
        )
        self._not_implemented_error_if(
            self._is_fp8(c_dtype) and self._is_fp8(ab_dtype),
            "Unsupported c_dtype and ab_dtype combination: fp8 ab_dtype and fp8 c_dtype (fails to launch)",
        )
        self._check_dtype(
            self.acc_dtype,
            dtype=torch.float32,
            name="Accumulator",
            extra_error_msg="Accumulator must be float32",
        )

        self.ab_dtype = ab_dtype
        self.c_dtype = c_dtype

        self._logger.debug("Checking tensor layout")
        m, k, l = self.a_desc.shape
        n, _, _ = self.b_desc.shape
        _, _, m_div_atom_m0_m1, _, sf_k_div_atom_k, _ = self.sfa_desc.shape
        _, _, n_div_atom_m0_m1, _, sf_k_div_atom_k, _ = self.sfb_desc.shape

        self._check_tensor_shape(self.a_desc, (m, k, l), "A")
        self._check_tensor_shape(self.b_desc, (n, k, l), "B")
        self._check_tensor_shape(self.c_desc, (m, n, l), "C")
        self._check_tensor_shape(
            self.sfa_desc,
            (self.atom_m[0], self.atom_m[1], m_div_atom_m0_m1, self.atom_k, sf_k_div_atom_k, l),
            "sfa",
        )
        self._check_tensor_shape(
            self.sfb_desc,
            (self.atom_m[0], self.atom_m[1], n_div_atom_m0_m1, self.atom_k, sf_k_div_atom_k, l),
            "sfb",
        )
        self.amax_desc = self._pad_tensor_to_ndim(self.amax_desc, 3, "amax")
        self._check_tensor_shape(self.amax_desc, (1, 1, 1), "amax")

        expected_m_div_atom = ceil_div(m, self.atom_m[0] * self.atom_m[1])
        expected_n_div_atom = ceil_div(n, self.atom_m[0] * self.atom_m[1])
        self._value_error_if(
            m_div_atom_m0_m1 != expected_m_div_atom,
            f"Input/Output shape mismatch: expected m_div_atom_m0_m1 (sfa.shape[2]) = {expected_m_div_atom}, got {m_div_atom_m0_m1}",
        )
        self._value_error_if(
            n_div_atom_m0_m1 != expected_n_div_atom,
            f"Input/Output shape mismatch: expected n_div_atom_m0_m1 (sfb.shape[2]) = {expected_n_div_atom}, got {n_div_atom_m0_m1}",
        )

        # Check tensor strides
        _ = self._check_tensor_stride(
            self.a_desc,
            stride=[(1, m, m * k), (k, 1, m * k)],
            name="A",
        )
        _ = self._check_tensor_stride(
            self.b_desc,
            stride=[(1, n, n * k), (k, 1, n * k)],
            name="B",
        )
        _ = self._check_tensor_stride(
            self.c_desc,
            stride=[(1, m, m * n), (n, 1, m * n)],
            name="C",
        )

        # Derive major mode from stride order
        self.a_major = "m" if self.a_desc.stride_order == (0, 1, 2) else "k"
        self.b_major = "n" if self.b_desc.stride_order == (0, 1, 2) else "k"
        self.c_major = "m" if self.c_desc.stride_order == (0, 1, 2) else "n"

        self._value_error_if(
            self._is_fp4x2(ab_dtype) and not (self.a_major == "k" and self.b_major == "k"),
            f"Unsupported A or B tensor stride: Float4 tensors require k-major layout for hardware efficiency, got {self.a_major} and {self.b_major}",
        )
        self._value_error_if(
            self._is_fp4x2(c_dtype) and self.c_major == "m",
            f"Unsupported C tensor stride: Float4 tensors require n-major layout for hardware efficiency, got {self.c_major}",
        )

        self._logger.debug("Checking mma tiler and cluster shape")
        self._value_error_if(
            self.mma_tiler_mn[0] not in [128, 256],
            f"Unsupported mma_tiler_mn[0]: expected {{128, 256}}, got {self.mma_tiler_mn[0]}",
        )
        self._value_error_if(
            self.mma_tiler_mn[1] not in [128, 256],
            f"Unsupported mma_tiler_mn[1]: expected {{128, 256}}, got {self.mma_tiler_mn[1]}",
        )
        self._not_implemented_error_if(
            self.mma_tiler_mn[0] == 256,
            "mma_tiler_mn[0] == 256 currently hangs",
        )
        self._value_error_if(
            self._is_fp4x2(self.ab_dtype) and self.mma_tiler_mn[1] == 256 and k <= 128,
            f"mma_tiler_mn (X, 256) requires k > 128 (packed x2), got {k}",
        )
        self._value_error_if(
            not (self.cluster_shape_mn[0] % (2 if self.mma_tiler_mn[0] == 256 else 1) == 0),
            "Illegal cluster shape",
        )
        self._not_implemented_error_if(
            self.mma_tiler_mn == (128, 256) and self.sf_vec_size == 16 and c_dtype in {torch.float32, torch.float16, torch.bfloat16},
            "mma_tiler_mn (128, 256), sf_vec_size 16, c_dtype {torch.float32, torch.float16, torch.bfloat16} fails to launch",
        )

        # Special cluster shape check for scale factor multicasts.
        # Due to limited size of scale factors, we can't multicast among more than 4 CTAs.
        self._value_error_if(
            not (
                self.cluster_shape_mn[0] <= 4
                and self.cluster_shape_mn[1] <= 4
                and self.cluster_shape_mn[0] > 0
                and self.cluster_shape_mn[1] > 0
                and is_power_of_2(self.cluster_shape_mn[0])
                and is_power_of_2(self.cluster_shape_mn[1])
            ),
            f"Invalid cluster shape: expected cluster_shape_mn values in {{1, 2, 4}}, got {self.cluster_shape_mn}",
        )

        self._logger.debug("Checking tensor alignment")

        def check_contigous_16B_alignment(dtype, is_mode0_major, tensor_shape):
            major_mode_idx = 0 if is_mode0_major else 1
            num_major_elements = tensor_shape[major_mode_idx]
            num_contiguous_elements = 16 * 8 // (_convert_to_cutlass_data_type(dtype).width)
            return num_major_elements % num_contiguous_elements == 0

        self._value_error_if(
            not (
                check_contigous_16B_alignment(ab_dtype, self.a_major == "m", (m, k, l))
                and check_contigous_16B_alignment(ab_dtype, self.b_major == "n", (n, k, l))
                and check_contigous_16B_alignment(c_dtype, self.c_major == "m", (m, n, l))
            ),
            "Unsupported tensor alignment: tensors must be 16B aligned",
        )

        self._logger.debug("Checking environment")
        self._runtime_error_if(not torch.cuda.is_available(), "CUDA is not available")
        device = torch.cuda.current_device()
        major, minor = torch.cuda.get_device_capability(device)
        compute_capability = major * 10 + minor
        self._runtime_error_if(
            compute_capability < 100,
            f"GemmAmax requires SM100+ compute capability, but found SM{compute_capability} on device {device}",
        )

        self._kernel = Sm100BlockScaledPersistentDenseGemmKernel

        self._is_supported = True
        self._logger.debug("check_support completed successfully")
        return True

    def compile(self) -> None:
        self._logger.debug("Entering compile")
        self._ensure_support_checked()
        if self._compiled_kernel is not None:
            self._logger.debug("Kernel already compiled; skipping recompilation")
            return

        gemm_amax = self._kernel(
            sf_vec_size=self.sf_vec_size,
            mma_tiler_mn=self.mma_tiler_mn,
            cluster_shape_mn=self.cluster_shape_mn,
        )
        hardware_info = cutlass.utils.HardwareInfo()
        max_active_clusters = hardware_info.get_max_active_clusters(self.cluster_shape_mn[0] * self.cluster_shape_mn[1])
        max_active_clusters -= self.num_cluster_overlap_margin
        self._value_error_if(
            max_active_clusters <= 0,
            "max_active_clusters must be > 0 after applying overlap margin; reduce CUDNNFE_CLUSTER_OVERLAP_MARGIN",
        )

        self._logger.debug("Compiling gemm_amax")
        a_cute = self._make_fake_cute_tensor_from_desc(self.a_desc, assumed_align=16)
        b_cute = self._make_fake_cute_tensor_from_desc(self.b_desc, assumed_align=16)
        sfa_cute = self._make_fake_cute_tensor_from_desc(self.sfa_desc, assumed_align=16)
        sfb_cute = self._make_fake_cute_tensor_from_desc(self.sfb_desc, assumed_align=16)
        c_cute = self._make_fake_cute_tensor_from_desc(self.c_desc, assumed_align=16)
        amax_cute = self._make_fake_cute_tensor_from_desc(self.amax_desc, assumed_align=16)
        fake_stream = make_fake_stream(use_tvm_ffi_env_stream=False)

        _compiled_kernel = cute.compile(
            gemm_amax,
            a_tensor=a_cute,
            b_tensor=b_cute,
            sfa_tensor=sfa_cute,
            sfb_tensor=sfb_cute,
            c_tensor=c_cute,
            amax_tensor=amax_cute,
            max_active_clusters=max_active_clusters,
            stream=fake_stream,
            options="--enable-tvm-ffi",
        )

        def tensor_api(
            a_tensor: torch.Tensor,
            b_tensor: torch.Tensor,
            sfa_tensor: torch.Tensor,
            sfb_tensor: torch.Tensor,
            c_tensor: torch.Tensor,
            amax_tensor: torch.Tensor,
            stream: cuda.CUstream,
        ):
            amax_tensor = self._pad_tensor_to_ndim(amax_tensor, 3, "amax")

            return _compiled_kernel(
                a_tensor,
                b_tensor,
                sfa_tensor,
                sfb_tensor,
                c_tensor,
                amax_tensor,
                stream,
            )

        self._compiled_kernel = tensor_api

        for attr_name in tuple(vars(self)):
            if attr_name.startswith("sample_"):
                setattr(self, attr_name, None)

        self._logger.debug("Kernel compiled successfully")

    def execute(
        self,
        a_tensor: torch.Tensor,
        b_tensor: torch.Tensor,
        sfa_tensor: torch.Tensor,
        sfb_tensor: torch.Tensor,
        c_tensor: torch.Tensor,
        amax_tensor: torch.Tensor,
        current_stream: Optional[cuda.CUstream] = None,
    ) -> None:
        self._logger.debug("Entering execute")
        current_stream = self._get_default_stream(current_stream)

        self._runtime_error_if(
            self._compiled_kernel is None,
            "GemmAmaxSm100 kernel not compiled; call compile() first",
        )
        self._logger.debug("Executing with compiled kernel")

        self._compiled_kernel(
            a_tensor=a_tensor,
            b_tensor=b_tensor,
            sfa_tensor=sfa_tensor,
            sfb_tensor=sfb_tensor,
            c_tensor=c_tensor,
            amax_tensor=amax_tensor,
            stream=current_stream,
        )
        self._logger.debug("Executed with compiled kernel successfully")


import logging

_logger = logging.getLogger(__name__)
_cache_of_GemmAmaxSm100Objects = {}


def gemm_amax_wrapper_sm100(
    a_tensor: torch.Tensor,
    b_tensor: torch.Tensor,
    sfa_tensor: torch.Tensor,
    sfb_tensor: torch.Tensor,
    c_major: str = "n",
    c_dtype: torch.dtype = torch.float32,
    acc_dtype: torch.dtype = torch.float32,
    mma_tiler_mn: Tuple[int, int] = (128, 128),
    cluster_shape_mn: Tuple[int, int] = (1, 1),
    sf_vec_size: int = 32,
    stream: Optional[cuda.CUstream] = None,
) -> TupleDict:

    _logger.debug("gemm_amax_wrapper_sm100: Creating empty output tensors c and amax")

    m, _, l = a_tensor.shape
    n, _, l = b_tensor.shape
    c_tensor = None
    if c_major == "m":
        c_tensor = torch.empty_strided((m, n, l), (1, m, m * n), dtype=c_dtype, device=a_tensor.device)
    elif c_major == "n":
        c_tensor = torch.empty_strided((m, n, l), (n, 1, m * n), dtype=c_dtype, device=a_tensor.device)
    else:
        raise ValueError(f"c_major must be either 'm' or 'n', got {c_major}")
    amax_tensor = torch.full((1, 1, 1), -float("inf"), device=a_tensor.device, dtype=torch.float32)

    cache_key = (
        a_tensor.shape,
        b_tensor.shape,
        sfa_tensor.shape,
        sfb_tensor.shape,
        a_tensor.dtype,
        b_tensor.dtype,
        sfa_tensor.dtype,
        sfb_tensor.dtype,
        a_tensor.stride(),
        b_tensor.stride(),
        sfa_tensor.stride(),
        sfb_tensor.stride(),
        c_major,
        c_dtype,
        acc_dtype,
        mma_tiler_mn,
        cluster_shape_mn,
        sf_vec_size,
    )
    if cache_key in _cache_of_GemmAmaxSm100Objects:
        _logger.debug("gemm_amax_wrapper_sm100: Using previously cached GemmAmaxSm100 object")
        gemm_amax = _cache_of_GemmAmaxSm100Objects[cache_key]
        gemm_amax.execute(
            a_tensor=a_tensor,
            b_tensor=b_tensor,
            sfa_tensor=sfa_tensor,
            sfb_tensor=sfb_tensor,
            c_tensor=c_tensor,
            amax_tensor=amax_tensor,
            current_stream=stream,
        )
    else:
        _logger.debug("gemm_amax_wrapper_sm100: No previously cached GemmAmaxSm100 object found, creating new GemmAmaxSm100 object")
        gemm_amax = GemmAmaxSm100(
            sample_a=a_tensor,
            sample_b=b_tensor,
            sample_sfa=sfa_tensor,
            sample_sfb=sfb_tensor,
            sample_c=c_tensor,
            sample_amax=amax_tensor,
            acc_dtype=acc_dtype,
            mma_tiler_mn=mma_tiler_mn,
            cluster_shape_mn=cluster_shape_mn,
            sf_vec_size=sf_vec_size,
        )
        assert gemm_amax.check_support()
        gemm_amax.compile()
        gemm_amax.execute(
            a_tensor=a_tensor,
            b_tensor=b_tensor,
            sfa_tensor=sfa_tensor,
            sfb_tensor=sfb_tensor,
            c_tensor=c_tensor,
            amax_tensor=amax_tensor,
            current_stream=stream,
        )
        _cache_of_GemmAmaxSm100Objects[cache_key] = gemm_amax

    return TupleDict(
        c_tensor=c_tensor,
        amax_tensor=amax_tensor,
    )
