# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# See LICENSE for license information.

"""GEMM API that enables custom GEMM logic for custom quantization recipes."""

import dataclasses
import enum
from typing import Iterable, Optional

import torch

from transformer_engine.pytorch.quantized_tensor import QuantizedTensorStorage, Quantizer
from transformer_engine.pytorch.tensor.utils import is_custom


@enum.unique
class GEMMType(enum.Enum):
    """Type of GEMM operation being performed."""

    FPROP = "fprop"
    DGRAD = "dgrad"
    WGRAD = "wgrad"


@dataclasses.dataclass(frozen=True)
class MMParams:
    """Matrix multiplication parameters."""

    out_dtype: torch.dtype | None = None
    # Use split accumulator for more accurate FP8 GEMM
    use_split_accumulator: bool = True


def custom_gemm(
    A: QuantizedTensorStorage,
    B: QuantizedTensorStorage,
    workspace: torch.Tensor,  # pylint: disable=unused-argument
    out_dtype: Optional[torch.dtype] = None,
    quantization_params: Optional[Quantizer] = None,  # pylint: disable=unused-argument
    gelu: bool = False,  # pylint: disable=unused-argument
    gelu_in: torch.Tensor = None,  # pylint: disable=unused-argument
    accumulate: bool = False,  # pylint: disable=unused-argument
    layout: str = "TN",
    out: Optional[torch.Tensor] = None,  # pylint: disable=unused-argument
    bias: Optional[torch.Tensor] = None,
    use_split_accumulator: bool = False,
    grad: bool = False,
) -> Iterable[Optional[torch.Tensor]]:
    """Dispatch GEMM to quantizer's qgemm method."""
    if not (is_custom(A) and is_custom(B)):
        raise TypeError("A and B must be custom tensors")

    A, B = B, A

    # Determine GEMM type based on grad flag and layout
    if not grad:
        gemm_type = GEMMType.FPROP
    else:
        if layout == "NN":
            gemm_type = GEMMType.DGRAD
        elif layout == "NT":
            gemm_type = GEMMType.WGRAD
        else:
            # Default to FPROP for other layouts
            gemm_type = GEMMType.FPROP

    # Extract quantizer from QuantizedTensor to get qgemm logic
    # TODO(etsykunov): make it more flexible, what if we might want to use gemm logic from B._quantizer?
    quantizer = None
    if hasattr(A, "_quantizer") and A._quantizer is not None:
        quantizer = A._quantizer
    elif hasattr(B, "_quantizer") and B._quantizer is not None:
        quantizer = B._quantizer
    else:
        raise ValueError("No quantizer found in QuantizedTensor objects")

    # Create MMParams
    m_params = MMParams(
        out_dtype=out_dtype,
        use_split_accumulator=use_split_accumulator,
    )
    out_dtype = A.dtype if m_params.out_dtype is None else m_params.out_dtype

    if gemm_type == GEMMType.FPROP:
        qx, sx = A.data, A.scale
        qw, sw = B.data, B.scale
        if qx is None:
            raise ValueError("FPROP GEMM: quantized activation data (A.data) is None")
        if sx is None:
            raise ValueError("FPROP GEMM: activation scale (A.scale) is None")
        if qw is None:
            raise ValueError("FPROP GEMM: quantized weight data (B.data) is None")
        if sw is None:
            raise ValueError("FPROP GEMM: weight scale (B.scale) is None")
        if A.original_shape is None:
            raise ValueError("FPROP GEMM: A.original_shape is None, cannot determine output shape")

        # Call quantizer's qgemm method
        result = quantizer.qgemm(
            qx,
            qw,
            m_params,
            out_dtype,
            sx,
            sw,
            bias,
            gemm_type=GEMMType.FPROP,
            qresult_x=A,
            qresult_w=B,
        )
        if len(A.original_shape) > 2:
            # Original input was 3D, so we need to reshape result back to 3D
            batch_size = A.original_shape[0]
            seq_len = A.original_shape[1]
            result = result.view(batch_size, seq_len, result.shape[-1])
    elif gemm_type == GEMMType.DGRAD:
        qdy, sdy = A.data, A.scale
        qw_t, sw_t = B.data_t, B.scale_t
        if qdy is None:
            raise ValueError("DGRAD GEMM: quantized gradient data (A.data) is None")
        if sdy is None:
            raise ValueError("DGRAD GEMM: gradient scale (A.scale) is None")
        if qw_t is None:
            raise ValueError("DGRAD GEMM: transposed quantized weight data (B.data_t) is None")
        if sw_t is None:
            raise ValueError("DGRAD GEMM: transposed weight scale (B.scale_t) is None")

        result = quantizer.qgemm(
            qdy,
            qw_t,
            m_params,
            out_dtype,
            sdy,
            sw_t,
            None,
            gemm_type=GEMMType.DGRAD,
            qresult_x=A,
            qresult_w=B,
        )
    elif gemm_type == GEMMType.WGRAD:
        qdy_t, sdy_t = A.data_t, A.scale_t
        qx_t, sx_t = B.data_t, B.scale_t
        if qdy_t is None:
            raise ValueError("WGRAD GEMM: transposed quantized gradient data (A.data_t) is None")
        if sdy_t is None:
            raise ValueError("WGRAD GEMM: transposed gradient scale (A.scale_t) is None")
        if qx_t is None:
            raise ValueError("WGRAD GEMM: transposed quantized activation data (B.data_t) is None")
        if sx_t is None:
            raise ValueError("WGRAD GEMM: transposed activation scale (B.scale_t) is None")

        result = quantizer.qgemm(
            qdy_t,
            qx_t,
            m_params,
            out_dtype,
            sdy_t,
            sx_t,
            None,
            gemm_type=GEMMType.WGRAD,
            qresult_x=A,
            qresult_w=B,
        )

    # Return in the same format as general_gemm
    return result, None, None, None
