# Unsloth
# Copyright 2023-present Daniel Han-Chen, Michael Han-Chen & the Unsloth team. All rights reserved.
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as published
# by the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with this program.  If not, see <https://www.gnu.org/licenses/>.

"""
Autotuning utils
"""

import logging
from itertools import product
from typing import List

import torch
import triton

logger = logging.getLogger(__name__)

DEFAULT_M_BLOCK_SIZES = [64, 128]
DEFAULT_N_BLOCK_SIZES = [64, 128, 256]
DEFAULT_K_BLOCK_SIZES = [64, 128, 256]
DEFAULT_NUM_CTAS = 1
DEFAULT_NUM_WARPS = [4, 8]
DEFAULT_NUM_STAGES = [3, 4, 5]
BOOLS = [True, False]


def val_to_list(val):
    if val is None:
        return None
    elif isinstance(val, list):
        return val
    else:
        return [val]


def convert_args_to_list(args):
    return [val_to_list(arg) for arg in args]


def _triton_supports_tma():
    """Check if current Triton version supports TMA API."""
    import triton.language as tl

    # Check for both old experimental and new stable API names
    return hasattr(tl, "make_tensor_descriptor") or hasattr(
        tl, "_experimental_make_tensor_descriptor"
    )


# TMA disabled for now: incompatible with permute_x/permute_y in the MoE
# grouped GEMM passes. Re-enable once resolved.
_TRITON_HAS_TMA = False  # _triton_supports_tma()


def get_forward_configs(
    BLOCK_M = DEFAULT_M_BLOCK_SIZES,
    BLOCK_N = DEFAULT_N_BLOCK_SIZES,
    BLOCK_K = DEFAULT_K_BLOCK_SIZES,
    TMA_LOAD_X = None,  # Auto-detect if not specified
    TMA_LOAD_W = None,  # Auto-detect if not specified
    TMA_STORE = False,  # NOTE: TMA_STORE is disabled for now
    num_warps = DEFAULT_NUM_WARPS,
    num_stages = DEFAULT_NUM_STAGES,
    num_ctas = DEFAULT_NUM_CTAS,
):
    # Auto-detect TMA support
    if TMA_LOAD_X is None:
        TMA_LOAD_X = _TRITON_HAS_TMA
    if TMA_LOAD_W is None:
        TMA_LOAD_W = _TRITON_HAS_TMA

    (
        BLOCK_M,
        BLOCK_N,
        BLOCK_K,
        TMA_LOAD_X,
        TMA_LOAD_W,
        TMA_STORE,
        num_warps,
        num_stages,
        num_ctas,
    ) = convert_args_to_list(
        [
            BLOCK_M,
            BLOCK_N,
            BLOCK_K,
            TMA_LOAD_X,
            TMA_LOAD_W,
            TMA_STORE,
            num_warps,
            num_stages,
            num_ctas,
        ]
    )
    kernel_configs = []
    for (
        block_m,
        block_n,
        block_k,
        w,
        s,
        tma_load_x,
        tma_load_w,
        tma_store,
        num_ctas,
    ) in product(
        BLOCK_M,
        BLOCK_N,
        BLOCK_K,
        num_warps,
        num_stages,
        TMA_LOAD_X,
        TMA_LOAD_W,
        TMA_STORE,
        num_ctas,
    ):
        kernel_configs.append(
            triton.Config(
                dict(
                    BLOCK_SIZE_M = block_m,
                    BLOCK_SIZE_N = block_n,
                    BLOCK_SIZE_K = block_k,
                    USE_TMA_LOAD_X = tma_load_x,
                    USE_TMA_LOAD_W = tma_load_w,
                    USE_TMA_STORE = tma_store,
                ),
                num_warps = w,
                num_stages = s,
                num_ctas = num_ctas,
            )
        )

    return kernel_configs


def get_dX_kernel_configs(
    BLOCK_M = DEFAULT_M_BLOCK_SIZES,
    BLOCK_N = DEFAULT_N_BLOCK_SIZES,
    BLOCK_K = DEFAULT_K_BLOCK_SIZES,
    TMA_LOAD_dY = None,  # Auto-detect if not specified
    TMA_LOAD_W = None,  # Auto-detect if not specified
    TMA_STORE = False,  # NOTE: TMA_STORE is disabled for now
    num_warps = DEFAULT_NUM_WARPS,
    num_stages = DEFAULT_NUM_STAGES,
    num_ctas = DEFAULT_NUM_CTAS,
):
    # Auto-detect TMA support
    if TMA_LOAD_dY is None:
        TMA_LOAD_dY = _TRITON_HAS_TMA
    if TMA_LOAD_W is None:
        TMA_LOAD_W = _TRITON_HAS_TMA
    (
        BLOCK_M,
        BLOCK_N,
        BLOCK_K,
        TMA_LOAD_dY,
        TMA_LOAD_W,
        TMA_STORE,
        num_warps,
        num_stages,
        num_ctas,
    ) = convert_args_to_list(
        [
            BLOCK_M,
            BLOCK_N,
            BLOCK_K,
            TMA_LOAD_dY,
            TMA_LOAD_W,
            TMA_STORE,
            num_warps,
            num_stages,
            num_ctas,
        ]
    )
    kernel_configs = []
    for (
        block_m,
        block_n,
        block_k,
        w,
        s,
        tma_load_dy,
        tma_load_w,
        tma_store,
        num_ctas,
    ) in product(
        BLOCK_M,
        BLOCK_N,
        BLOCK_K,
        num_warps,
        num_stages,
        TMA_LOAD_dY,
        TMA_LOAD_W,
        TMA_STORE,
        num_ctas,
    ):
        kernel_configs.append(
            triton.Config(
                dict(
                    BLOCK_SIZE_M = block_m,
                    BLOCK_SIZE_N = block_n,
                    BLOCK_SIZE_K = block_k,
                    USE_TMA_LOAD_dY = tma_load_dy,
                    USE_TMA_LOAD_W = tma_load_w,
                    USE_TMA_STORE = tma_store,
                ),
                num_warps = w,
                num_stages = s,
                num_ctas = num_ctas,
            )
        )

    return kernel_configs


def get_dW_kernel_configs(
    BLOCK_M = DEFAULT_M_BLOCK_SIZES,
    BLOCK_N = DEFAULT_N_BLOCK_SIZES,
    BLOCK_K = DEFAULT_K_BLOCK_SIZES,
    num_warps = DEFAULT_NUM_WARPS,
    num_stages = DEFAULT_NUM_STAGES,
    num_ctas = DEFAULT_NUM_CTAS,
    TMA_LOAD_dY = None,  # Auto-detect if not specified
    TMA_LOAD_X = None,  # Auto-detect if not specified
    TMA_STORE = False,
):
    # Auto-detect TMA support
    if TMA_LOAD_dY is None:
        TMA_LOAD_dY = _TRITON_HAS_TMA
    if TMA_LOAD_X is None:
        TMA_LOAD_X = _TRITON_HAS_TMA
    (
        BLOCK_M,
        BLOCK_N,
        BLOCK_K,
        num_warps,
        num_stages,
        num_ctas,
        TMA_LOAD_dY,
        TMA_LOAD_X,
        TMA_STORE,
    ) = convert_args_to_list(
        [
            BLOCK_M,
            BLOCK_N,
            BLOCK_K,
            num_warps,
            num_stages,
            num_ctas,
            TMA_LOAD_dY,
            TMA_LOAD_X,
            TMA_STORE,
        ]
    )
    kernel_configs = []
    for (
        block_m,
        block_n,
        block_k,
        w,
        s,
        tma_load_dy,
        tma_load_x,
        tma_store,
        num_ctas,
    ) in product(
        BLOCK_M,
        BLOCK_N,
        BLOCK_K,
        num_warps,
        num_stages,
        TMA_LOAD_dY,
        TMA_LOAD_X,
        TMA_STORE,
        num_ctas,
    ):
        kernel_configs.append(
            triton.Config(
                dict(
                    BLOCK_SIZE_M = block_m,
                    BLOCK_SIZE_N = block_n,
                    BLOCK_SIZE_K = block_k,
                    USE_TMA_LOAD_dY = tma_load_dy,
                    USE_TMA_LOAD_X = tma_load_x,
                    USE_TMA_STORE = tma_store,
                ),
                num_warps = w,
                num_stages = s,
                num_ctas = num_ctas,
            )
        )

    return kernel_configs


def estimate_smem_reqs(
    num_stages: int, BLOCK_SIZE_M: int, BLOCK_SIZE_N: int, BLOCK_SIZE_K: int, dtype: torch.dtype
):
    num_bytes = dtype.itemsize
    return (
        num_stages * BLOCK_SIZE_K * (BLOCK_SIZE_M + BLOCK_SIZE_N) + BLOCK_SIZE_M * BLOCK_SIZE_N
    ) * num_bytes


def exceeds_smem_capacity(
    num_stages: int,
    BLOCK_SIZE_M: int,
    BLOCK_SIZE_N: int,
    BLOCK_SIZE_K: int,
    dtype: torch.dtype,
    smem_size: int,
    slack: float = 50000,
):
    smem_reqs = estimate_smem_reqs(num_stages, BLOCK_SIZE_M, BLOCK_SIZE_N, BLOCK_SIZE_K, dtype)
    return smem_reqs > smem_size + slack


def common_prune_criteria(config: triton.Config, kwargs: dict, dtype):
    from ..interface import supports_tma
    from .tuning import get_device_properties

    smem_size = get_device_properties().SIZE_SMEM

    num_stages = config.num_stages
    BLOCK_SIZE_M = config.kwargs["BLOCK_SIZE_M"]
    BLOCK_SIZE_N = config.kwargs["BLOCK_SIZE_N"]
    BLOCK_SIZE_K = config.kwargs["BLOCK_SIZE_K"]

    num_tokens = kwargs["NUM_TOKENS"]
    num_experts = kwargs["NUM_EXPERTS"]
    permute_x = kwargs["PERMUTE_X"]
    permute_y = kwargs["PERMUTE_Y"]
    tokens_per_expert = num_tokens // num_experts

    # use_tma = [k for k in config.kwargs.keys() if k.startswith("USE_TMA_")]
    MIN_BLOCK_SIZE_M = DEFAULT_M_BLOCK_SIZES[0]
    if exceeds_smem_capacity(
        num_stages, BLOCK_SIZE_M, BLOCK_SIZE_N, BLOCK_SIZE_K, dtype, smem_size
    ):
        return True
    if BLOCK_SIZE_M > tokens_per_expert * 2 and tokens_per_expert > MIN_BLOCK_SIZE_M:
        return True
    if permute_x and permute_y:
        return True
    # if not supports_tma() and any(use_tma):
    #     return True
    return False


def maybe_disable_tma(config: triton.Config):
    from ..interface import supports_tma
    tma_keys = [k for k in config.kwargs.keys() if k.startswith("USE_TMA_")]
    if not supports_tma():
        logger.info("Disabling TMA")
        for k in tma_keys:
            config.kwargs[k] = False


def prune_kernel_configs_fwd(configs: list[triton.Config], args, **kwargs):
    x = kwargs["x_ptr"]
    dtype = x.dtype

    logger.debug(f"Pruning configs: {len(configs)}")

    pruned_configs = []
    for config in configs:
        # disable TMA if gpu does not support it
        maybe_disable_tma(config)

        if common_prune_criteria(config, kwargs, dtype):
            continue
        if config.kwargs["USE_TMA_LOAD_X"] and kwargs["PERMUTE_X"]:
            # Dynamically disable TMA_LOAD_X for permuted X
            config.kwargs["USE_TMA_LOAD_X"] = False
        if config.kwargs["USE_TMA_STORE"] and kwargs["PERMUTE_Y"]:
            continue

        pruned_configs.append(config)

    logger.debug(f"Pruned configs: {len(pruned_configs)}")
    return pruned_configs


def prune_dX_configs(configs: List[triton.Config], args, **kwargs):
    dtype = kwargs["w_ptr"].dtype

    logger.debug(f"Pruning configs: {len(configs)}")
    pruned_configs = []

    for config in configs:
        if common_prune_criteria(config, kwargs, dtype):
            continue
        if config.kwargs["USE_TMA_LOAD_dY"] and kwargs["PERMUTE_Y"]:
            # dynamically disable TMA_LOAD_dY for permuted Y
            config.kwargs["USE_TMA_LOAD_dY"] = False
        if config.kwargs["USE_TMA_STORE"] and kwargs["PERMUTE_X"]:
            continue
        pruned_configs.append(config)

    logger.debug(f"Pruned configs: {len(pruned_configs)}")
    return pruned_configs


def prune_kernel_configs_backward_dW(configs: list[triton.Config], args, **kwargs):
    dtype = kwargs["x_ptr"].dtype

    pruned_configs = []
    logger.debug(f"Pruning configs: {len(configs)}")

    for config in configs:
        if common_prune_criteria(config, kwargs, dtype):
            continue
        if config.kwargs["USE_TMA_LOAD_dY"] and kwargs["PERMUTE_Y"]:
            config.kwargs["USE_TMA_LOAD_dY"] = False
        if config.kwargs["USE_TMA_LOAD_X"] and kwargs["PERMUTE_X"]:
            config.kwargs["USE_TMA_LOAD_X"] = False
        pruned_configs.append(config)

    logger.debug(f"Pruned configs: {len(pruned_configs)}")
    return pruned_configs
