# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.
#
# This source code is licensed under the BSD license found in the
# LICENSE file in the root directory of this source tree.

import math
from dataclasses import dataclass
from functools import partial
from typing import (
    Any,
    Callable,
    Iterable,
    List,
    Mapping,
    Optional,
    Set,
    Tuple,
    Type,
    Union,
)

import torch

from ..._cpp_lib import _built_with_cuda
from ..common import BaseOperator
from .attn_bias import (
    AttentionBias,
    AttentionBiasSubTensor,
    BlockDiagonalGappyKeysMask,
    BlockDiagonalMask,
    BlockDiagonalPaddedKeysMask,
    LowerTriangularMask,
    PagedBlockDiagonalGappyKeysMask,
    PagedBlockDiagonalPaddedKeysMask,
)


def _is_bias_type_supported_in_BMK(attn_bias_type: Any) -> bool:
    # NoneType
    if isinstance(None, attn_bias_type):
        return True
    if attn_bias_type in [LowerTriangularMask, torch.Tensor]:
        return True
    return False


def _attn_bias_apply(
    attn_bias: Optional[Union[torch.Tensor, AttentionBias]],
    op: Callable[[torch.Tensor], torch.Tensor],
) -> Optional[Union[torch.Tensor, AttentionBias]]:
    if isinstance(attn_bias, torch.Tensor) and attn_bias.ndim != 0:
        return op(attn_bias)
    return attn_bias


class ScaledTensor(torch.Tensor):
    __slots__ = ["scale", "dequant_func", "original_dtype"]

    # Disabling custom torch function handling for this class
    __torch_function__ = torch._C._disabled_torch_function_impl

    def __new__(
        cls,
        data: torch.Tensor,
        scale: torch.Tensor,
        dequant_func: Callable[[torch.Tensor, torch.Tensor], torch.Tensor],
        original_dtype: torch.dtype,
        require_grad: bool = False,
    ) -> "ScaledTensor":
        """
        Creates a new ScaledTensor subclass instance.

        Parameters:
        - data: The underlying quantized tensor (e.g., int8, int4).
        - scale: The scale tensor or scalar to be used for dequantization.
        - dequant_func: A callable that applies dequantization, which takes both the data and scale as input.
        - original_dtype: The data type before quantization (e.g., float32, float16).
        - require_grad: Whether or not to track gradients (default: False for inference use).
        """
        # Use _make_subclass to create a new ScaledTensor instance, which is a subclass of torch.Tensor.
        instance = torch.Tensor._make_subclass(cls, data, require_grad=require_grad)

        # Store the dequantization scale and function as attributes.
        instance.scale = scale  # type: ignore
        instance.dequant_func = dequant_func  # type: ignore

        # Store the original data type of the tensor, so we can cast it back after dequantization.
        instance.original_dtype = original_dtype  # type: ignore

        # Return the new instance of ScaledTensor.
        return instance

    def dequantize(self) -> torch.Tensor:
        """
        Applies the custom dequantization function provided at the tensor's creation.
        After dequantization, the data is cast back to its original data type.
        """
        # Explicitly create a new torch.Tensor to ensure the return type is torch.Tensor, not ScaledTensor.
        data = torch.Tensor(self.float())

        # Call the dequantization function, passing in the data and the scale.
        dequantized_data = self.dequant_func(data, self.scale)  # type: ignore

        # Cast the dequantized data back to the original data type.
        return dequantized_data.to(self.original_dtype)  # type: ignore

    def unpack(self) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        Unpacks the ScaledTensor by returning its data and scale as a tuple.
        Returns:
        - A tuple of (data, scale), both of which are torch.Tensor objects.
        """
        return self.data, self.scale  # type: ignore

    def __repr__(self):
        """
        Custom string representation for ScaledTensor.
        """
        return f"ScaledTensor(data={self.data}, scale={self.scale}, original_dtype={self.original_dtype})"


def pack_fp8_tensorwise_per_head(
    x: torch.Tensor, scale: Union[torch.Tensor, float], original_dtype
) -> ScaledTensor:
    """
    Pack a tensor into a tensorwise fp8 ScaledTensor.
    """
    if isinstance(scale, float):
        scale = torch.tensor([scale], device=x.device)

    def dequant_func(x, scale):
        return x * scale[:, None, :, None]

    return ScaledTensor(
        data=x,
        scale=scale,
        dequant_func=dequant_func,
        original_dtype=original_dtype,
    )


@dataclass
class Inputs:
    """
    Stores inputs to the `memory_efficient_attention` operators
    """

    query: torch.Tensor
    key: torch.Tensor
    value: torch.Tensor
    attn_bias: Optional[Union[torch.Tensor, AttentionBias]] = None
    p: float = 0.0
    scale: Optional[float] = None
    output_dtype: Optional[torch.dtype] = None
    is_partial: bool = False

    @property
    def device(self) -> torch.device:
        return self.query.device

    @property
    def scale_float(self) -> float:
        return self.query.shape[-1] ** (-0.5) if self.scale is None else self.scale

    def get_qkv_in_bmghk(self) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        if self.query.ndim == 5:
            return self.query, self.key, self.value
        if self.query.ndim == 4:
            return (
                self.query.unsqueeze(2),
                self.key.unsqueeze(2),
                self.value.unsqueeze(2),
            )
        if self.value.ndim == 3:
            return (
                self.query[:, :, None, None],
                self.key[:, :, None, None],
                self.value[:, :, None, None],
            )
        assert False

    def normalize_bmhk(self) -> Tuple[int, ...]:
        if self.query.ndim not in [3, 4, 5]:
            raise ValueError(
                f"Invalid shape for query: {self.query.shape}. "
                "Expected shape [batch, seqlen, head_groups, num_heads_per_group, K]"
                ", [batch, seqlen, num_heads, K], or [batch, seqlen, K]."
            )
        if self.value.dtype == torch.int32:
            # Quantized K/V case, in which the last dims of Q and K are different.
            # NB we currently don't have any implementations for quantized KV with
            # SUPPORTS_DIFFERENT_VALUE_EMBED.
            output_shape = tuple(self.query.shape)
        else:
            output_shape = (self.query.shape[:-1]) + (self.value.shape[-1],)
        # Convert from legacy format
        if self.query.ndim == 3:
            self.query = self.query.unsqueeze(2)
            self.key = self.key.unsqueeze(2)
            self.value = self.value.unsqueeze(2)
            self.attn_bias = _attn_bias_apply(
                self.attn_bias, partial(torch.unsqueeze, dim=1)
            )
        return output_shape

    def validate_inputs(self) -> None:
        qkv = (self.query, self.key, self.value)
        if self.query.ndim not in (3, 4, 5) or any(
            x.ndim != self.query.ndim for x in qkv
        ):
            raise ValueError(
                f"Query/Key/Value should all have BMGHK, BMHK or BMK shape.\n"
                f"  query.shape: {self.query.shape}\n"
                f"  key.shape  : {self.key.shape}\n"
                f"  value.shape: {self.value.shape}"
            )
        if any(x.device != self.query.device for x in qkv):
            raise ValueError("Query/Key/Value should all be on the same device")
        if isinstance(
            self.attn_bias,
            (
                BlockDiagonalMask,
                BlockDiagonalPaddedKeysMask,
                PagedBlockDiagonalPaddedKeysMask,
                BlockDiagonalGappyKeysMask,
                PagedBlockDiagonalGappyKeysMask,
            ),
        ):
            bias_device = self.attn_bias.q_seqinfo.seqstart.device
            if bias_device != self.query.device:
                raise ValueError(
                    f"Attention bias and Query/Key/Value should be on the same device\n"
                    f"  query.device: {self.query.device}\n"
                    f"  attn_bias   : {bias_device}\n"
                )

        quantized_dtypes = self.key.dtype == self.value.dtype == torch.int32
        non_quantized_dtypes = all(x.dtype == self.query.dtype for x in qkv)
        if not (quantized_dtypes or non_quantized_dtypes):
            raise ValueError(
                "Query/Key/Value should either all have the same dtype, or "
                "(in the quantized case) Key/Value should have dtype torch.int32\n"
                f"  query.dtype: {self.query.dtype}\n"
                f"  key.dtype  : {self.key.dtype}\n"
                f"  value.dtype: {self.value.dtype}"
            )
        # Biases with tensors attached are meant to be in BMHK format
        # This would require to permute biases/gradients which can be expensive,
        # so let's just forbid it - BMK is a legacy format anyway
        if self.query.ndim == 3 and not _is_bias_type_supported_in_BMK(
            type(self.attn_bias)
        ):
            raise ValueError(
                f"Please provide inputs in BMHK format rather "
                f"than BMK when using bias type `{type(self.attn_bias).__name__}`"
            )
        attn_bias_t: Optional[torch.Tensor] = None
        if isinstance(self.attn_bias, AttentionBiasSubTensor):
            if self.attn_bias.HOLDS_DENSE_TENSOR:
                attn_bias_t = self.attn_bias._subtensor
        elif isinstance(self.attn_bias, torch.Tensor):
            attn_bias_t = self.attn_bias
        if self.query.ndim == 4 and attn_bias_t is not None:
            expected_shape = (
                self.query.shape[0],
                self.query.shape[2],
                self.query.shape[1],
                self.key.shape[1],
            )
            if attn_bias_t.shape != expected_shape:
                raise ValueError(
                    f"Invalid shape for attention bias: {attn_bias_t.shape} (expected {expected_shape})\n"
                    f"  query.shape: {self.query.shape}\n"
                    f"  key.shape  : {self.key.shape}\n"
                    f"  value.shape: {self.value.shape}"
                )
        if isinstance(self.attn_bias, BlockDiagonalMask):
            if any(x.shape[0] != 1 for x in qkv):
                raise ValueError(
                    f"Expected batch_size=1 when using block-diagonal bias\n"
                    f"  query.shape: {self.query.shape}\n"
                    f"  key.shape  : {self.key.shape}\n"
                    f"  value.shape: {self.value.shape}"
                )
        if self.p < 0.0 or self.p > 1.0:
            raise ValueError(f"Invalid dropout probability: p={self.p}")
        # Check that shapes match between inputs
        B, Mq = self.query.shape[:2]
        K = self.query.shape[-1]
        B, Mkv = self.key.shape[:2]
        Kv = self.value.shape[-1]
        quantized_kv_cache = self.value.dtype == torch.int32
        key_embed_dim = Kv if quantized_kv_cache else K

        valid_shapes = True
        if self.query.ndim == 3:  # BMK
            valid_shapes = (
                self.query.shape == (B, Mq, K)
                and self.key.shape == (B, Mkv, K)
                and self.value.shape == (B, Mkv, Kv)
            )
        H = self.query.shape[-2]
        if self.query.ndim == 4:  # BMHK
            valid_shapes = (
                self.query.shape == (B, Mq, H, K)
                and self.key.shape == (B, Mkv, H, key_embed_dim)
                and self.value.shape == (B, Mkv, H, Kv)
            )
        G = self.query.shape[2]
        if self.query.ndim == 5:  # BMNHK
            valid_shapes = (
                self.query.shape == (B, Mq, G, H, K)
                and self.key.shape == (B, Mkv, G, H, key_embed_dim)
                and self.value.shape == (B, Mkv, G, H, Kv)
            )
        if not valid_shapes:
            raise ValueError(
                f"Incompatible shapes for attention inputs:\n"
                f"  query.shape: {self.query.shape}\n"
                f"  key.shape  : {self.key.shape}\n"
                f"  value.shape: {self.value.shape}\n"
                "HINT: We don't support broadcasting, please use `expand` "
                "yourself before calling `memory_efficient_attention` if you need to"
            )

    def get_output_dtype(self) -> torch.dtype:
        if self.output_dtype is None:
            if self.is_partial and self.query.dtype is not torch.float64:
                return torch.float32
            return self.query.dtype
        return self.output_dtype

    @property
    def nbytes(self) -> int:
        """
        Number of bytes in the input, not counting the attention bias.
        """
        return sum(
            x.untyped_storage().nbytes() for x in [self.query, self.key, self.value]
        )


@dataclass
class Context:
    lse: torch.Tensor
    out: torch.Tensor
    # NOTE: If `rng_state` is set, `op_bw` should be set as well
    # as the randomness is backend-dependant
    op_bw: Optional[Type["AttentionBwOpBase"]] = None
    rng_state: Optional[Any] = None
    qkv_share_storage: bool = False

    def get_padded_lse(self, pad_to: int, force_pad_inf: bool = False) -> torch.Tensor:
        pad_amount = (pad_to - (self.lse.shape[2] % pad_to)) % pad_to
        lse = self.lse
        if pad_amount > 0:
            if force_pad_inf:
                lse = lse[:, :, : self.out.shape[1]]
                pad_amount = (pad_to - (lse.shape[2] % pad_to)) % pad_to
            lse = torch.nn.functional.pad(lse, [0, pad_amount], value=math.inf)
        elif force_pad_inf and self.out.shape[1] != lse.shape[2]:
            lse[:, :, self.out.shape[1] :].fill_(math.inf)
        return lse


@dataclass
class Gradients:
    dq: torch.Tensor
    dk: torch.Tensor
    dv: torch.Tensor
    # bias gradient. None if there is no tensor bias or if it doesn't require grad
    db: Optional[torch.Tensor] = None


class AttentionOpBase(BaseOperator):
    """Base class for any attention operator in xFormers

    See:

    - :attr:`xformers.ops.fmha.cutlass.FwOp`
    - :attr:`xformers.ops.fmha.cutlass.BwOp`
    - :attr:`xformers.ops.fmha.flash.FwOp`
    - :attr:`xformers.ops.fmha.flash.BwOp`
    - :attr:`xformers.ops.fmha.triton.FwOp`
    - :attr:`xformers.ops.fmha.triton.BwOp`
    """

    OPERATOR: Any
    SUPPORTED_DEVICES: Set[str]
    CUDA_MINIMUM_COMPUTE_CAPABILITY: Tuple[int, int] = (5, 0)
    SUPPORTED_DTYPES: Set[torch.dtype]
    SUPPORTED_MAX_K: float
    SUPPORTED_MIN_K: int = 0
    SUPPORTED_ATTN_BIAS_TYPES: Iterable[Any] = (type(None),)
    SUPPORTS_DROPOUT: bool
    SUPPORTS_CUSTOM_SCALE: bool = False
    SUPPORTS_DIFFERENT_VALUE_EMBED: bool = False
    SUPPORTS_OUTPUT_DTYPE: bool = False
    SUPPORTS_PARTIAL: bool = False
    IS_DETERMINISTIC: bool = True
    SUPPORTS_BMGHK: bool = False
    NAME: str
    OPERATOR_CATEGORY = "memory_efficient_attention"
    # Format for the LSE computed in the FW pass, and accepted in the BW pass,
    # for BlockDiagonalMask and children.
    # When using a varlen bias, both the FW and BW operators must have the
    # same value for `VARLEN_LSE_PACKED`
    VARLEN_LSE_PACKED: bool = True

    _TEST_BATCH_SIZES: List[int] = [1, 300]
    _TEST_K: List[int] = [32, 128]

    @classmethod
    def supports(cls, d: Inputs) -> bool:
        return not cls.not_supported_reasons(d)

    @classmethod
    def shape_not_supported_reasons(
        cls, Mq: int, Mkv: int, K: int, Kv: int
    ) -> List[str]:
        reasons = []
        if not cls.SUPPORTS_DIFFERENT_VALUE_EMBED and K != Kv:
            reasons.append("query.shape[-1] != value.shape[-1]")
        if max(K, Kv) > cls.SUPPORTED_MAX_K:
            reasons.append(
                f"max(query.shape[-1], value.shape[-1]) > {cls.SUPPORTED_MAX_K}"
            )
        if min(K, Kv) < cls.SUPPORTED_MIN_K:
            reasons.append(
                f"min(query.shape[-1], value.shape[-1]) < {cls.SUPPORTED_MIN_K}"
            )
        return reasons

    @classmethod
    def not_supported_reasons(cls, d: Inputs) -> List[str]:
        """
        Returns a list of reasons why this is not supported.
        The kernel can run these inputs only if the returned list is empty
        """
        query_shape = d.query.shape
        reasons = cls.shape_not_supported_reasons(
            Mq=query_shape[1],
            Mkv=d.key.shape[1],
            K=query_shape[-1],
            Kv=query_shape[-1] if d.value.dtype == torch.int32 else d.value.shape[-1],
        )
        device_type = d.query.device.type
        dtype = d.query.dtype
        if device_type not in cls.SUPPORTED_DEVICES:
            reasons.append(f"device={device_type} (supported: {cls.SUPPORTED_DEVICES})")
        if (
            device_type == "cuda"
            and not _built_with_cuda
            and (torch.version.hip is None)
        ):
            reasons.append("xFormers wasn't build with CUDA support")
        if device_type == "cuda" and (torch.version.hip is None):
            device_capability = torch.cuda.get_device_capability(d.device)
            if device_capability < cls.CUDA_MINIMUM_COMPUTE_CAPABILITY:
                reasons.append(
                    f"requires device with capability > {cls.CUDA_MINIMUM_COMPUTE_CAPABILITY} "
                    f"but your GPU has capability {device_capability} (too old)"
                )
        if dtype not in cls.SUPPORTED_DTYPES:
            reasons.append(f"dtype={dtype} (supported: {cls.SUPPORTED_DTYPES})")
        if type(d.attn_bias) not in cls.SUPPORTED_ATTN_BIAS_TYPES:
            reasons.append(f"attn_bias type is {type(d.attn_bias)}")
        if not cls.SUPPORTS_OUTPUT_DTYPE:
            if d.output_dtype is not None and d.output_dtype is not dtype:
                reasons.append("Custom output dtype not supported")
        if d.is_partial and not cls.SUPPORTS_PARTIAL:
            reasons.append("Partial attention not supported")
        if (d.p != 0.0) and not cls.SUPPORTS_DROPOUT:
            reasons.append("dropout > 0.0")
        if d.scale is not None and not cls.SUPPORTS_CUSTOM_SCALE:
            reasons.append("has custom scale")
        # bfloat16 is only supported on A100+
        # ... although the kernels can still run and give the
        # correct result
        if dtype is torch.bfloat16 and (
            not device_type.startswith("cuda")
            or torch.cuda.get_device_capability(d.query.device)[0] < 8
        ):
            reasons.append("bf16 is only supported on A100+ GPUs")
        if not cls.is_available():
            reasons.append(
                "operator wasn't built - see `python -m xformers.info` for more info"
            )
        if not cls.IS_DETERMINISTIC and torch.are_deterministic_algorithms_enabled():
            reasons.append(
                "operator is non-deterministic, but `torch.use_deterministic_algorithms` is set"
            )
        if not cls.SUPPORTS_BMGHK and d.query.ndim == 5:
            reasons.append("operator does not support BMGHK format")
        return reasons


class AttentionFwOpBase(AttentionOpBase):
    ERROR_ATOL: Mapping[torch.dtype, float] = {
        torch.float: 3e-4,
        torch.half: 4e-3,
        torch.bfloat16: 2e-2,
    }
    ERROR_RTOL: Mapping[torch.dtype, float] = {
        torch.float: 2e-5,
        torch.half: 4e-4,
        torch.bfloat16: 5e-3,
    }

    @classmethod
    def apply(
        cls, inp: Inputs, needs_gradient: bool
    ) -> Tuple[torch.Tensor, Optional[Context]]:
        raise NotImplementedError()


class AttentionBwOpBase(AttentionOpBase):
    # NOTE on tolerances: These are tested for `scales => (1/32)**0.5`
    # In the BW pass, imprecisions accumulate in the Q@K.T recalculation
    # These imprecisions are multiplied by the `scale` and then exponentiated
    # So if the scale is too high, we get a lot of errors

    ERROR_ATOL: Mapping[torch.dtype, float] = {
        torch.float: 9e-4,
        torch.half: 0.2,
        torch.bfloat16: 0.9,
    }
    ERROR_RTOL: Mapping[torch.dtype, float] = {
        torch.float: 1e-4,
        torch.half: 2e-2,
        torch.bfloat16: 0.1,
    }
    SUPPORTS_ATTN_BIAS_GRAD = False
    SUPPORTS_PARTIAL = True

    @classmethod
    def not_supported_reasons(cls, d: Inputs) -> List[str]:
        reasons = super(AttentionBwOpBase, cls).not_supported_reasons(d)
        if (
            isinstance(d.attn_bias, torch.Tensor)
            and d.attn_bias.requires_grad
            and not cls.SUPPORTS_ATTN_BIAS_GRAD
        ):
            reasons.append(
                "Computing the bias gradient is not supported (attn_bias.requires_grad = True)"
            )

        return reasons

    @classmethod
    def apply(cls, ctx: Context, inp: Inputs, grad: torch.Tensor) -> Gradients:
        raise NotImplementedError()


AttentionOp = Tuple[
    Optional[Type[AttentionFwOpBase]], Optional[Type[AttentionBwOpBase]]
]


def bmk2bmhk(tensor, num_heads: int) -> torch.Tensor:
    if tensor.ndim == 4:
        return tensor
    return tensor.reshape(
        [tensor.shape[0] // num_heads, num_heads, tensor.shape[1], tensor.shape[2]]
    ).permute((0, 2, 1, 3))


def check_lastdim_alignment_stride1(
    reasons: List[str], name: str, x: torch.Tensor, alignment: int
) -> None:
    if x.shape[-1] % alignment != 0:
        reasons.append(f"{name}.shape[-1] % {alignment} != 0")
    elif x.stride(-2) % alignment != 0:
        reasons.append(
            f"{name}.stride(-2) % {alignment} != 0 ({name}.stride() = {x.stride()})"
        )
    # We can have stride=0 sometimes if dimension=1
    if x.stride(-1) > 1:
        reasons.append(
            f"{name}.stride(-1) > 1 ({name}.stride() = {x.stride()}) - you should call `.contiguous()` on the input"
        )