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

from typing import Optional
import torch
import warnings
import triton
from aiter.ops.triton._triton_kernels.gemm.feed_forward.ff_a16w16_fused_ungated import (
    _ff_a16w16_fused_ungated,
    _get_config,
)
from aiter.ops.triton._triton_kernels.activation import _get_activation_from_str

from aiter.ops.triton.utils.logger import AiterTritonLogger

_LOGGER = AiterTritonLogger()


def ff_a16w16_fused_ungated(
    x,
    w_up,
    w_down,
    dtype: Optional[float] = torch.bfloat16,
    y: Optional[torch.Tensor] = None,
    config: Optional[dict] = None,
    activation: Optional[str] = None,
):
    """
    Computes a full feed-forward operation with activation (e.g FF with Relu)

    Key parameters:
    - X: Matrix X with shape (M, K).
    - w_up: Up-projection W with shape (N, K).
    - w_down: Down-projection W with shape (N, K).
    - dtype: Optional parameter to specify bf16 or fp16 datatype. Default is bf16
    - Y: Output Matrix Y with shape (M, K).
    If this is none, then it's created by this API and returned as output.
    - activation: Optional activation function to apply.
    One of ("gelu", "gelu_tanh", "silu", "silu_exp2", "relu", None)

    Returns:
    - Y: The output matrix with shape (M, K).
    """

    _LOGGER.info(
        f"FF_A16W16_FUSED_UNGATED: x={tuple(x.shape)} w_up={tuple(w_up.shape)} w_down={tuple(w_down.shape) }"
    )

    # Shape checks
    assert (
        x.shape[1] == w_up.shape[1] == w_down.shape[1]
    ), f"Incompatible matrix shapes: x:{x.shape}, w_up:{w_up.shape}, w_down:{w_down.shape}"
    assert (
        w_up.shape[0] == w_down.shape[0]
    ), f"Incompatible matrix shapes: w_up:{w_up.shape}, w_down:{w_down.shape}"

    N, K = w_up.shape
    M = x.shape[0]
    if M > 64:
        warnings.warn(
            "The fused FF kernel is slower than the unfused equivalent for large batch sizes (>64)."
        )

    w_up = w_up.T

    if y is None:
        y = torch.zeros(
            (M, K), dtype=dtype, device=x.device
        )  # zeros, as this does atomic adds on top

    if config is None:
        config, _ = _get_config(M, N, K)

    grid = lambda META: (  # noqa: E731
        triton.cdiv(M, META["BLOCK_SIZE_M"]) * triton.cdiv(N, META["BLOCK_SIZE_N"]),
    )
    _ff_a16w16_fused_ungated[grid](
        x,
        w_up,
        w_down,
        y,
        M,
        N,
        K,
        x.stride(0),
        x.stride(1),
        w_up.stride(0),
        w_up.stride(1),
        w_down.stride(0),
        w_down.stride(1),
        y.stride(0),
        y.stride(1),
        activation=_get_activation_from_str(activation) if activation else "",
        use_activation=activation is not None,
        **config,
    )

    return y
