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

"""Mixin class holding data specific for NVFP4Tensor"""

from __future__ import annotations
from collections.abc import Iterable
import functools
import math
from typing import Annotated, Any, Dict, Optional, Tuple, Union
import warnings

import torch

import transformer_engine_torch as tex

from ...quantized_tensor import InnerTensor, QuantizedTensorStorage, Quantizer
from .._quantization_helpers import safe_quantized_repr

from ...constants import TE_DType as torch_to_transformer_engine_dtype, DType
from ...utils import _empty_tensor


@functools.lru_cache(maxsize=None)
def _fp4_e2m1_vals(device: torch.device, dtype: torch.dtype) -> torch.Tensor:
    """Values representable in FP4 E2M1 format"""
    return torch.tensor(
        [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, -0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0],
        device=device,
        dtype=dtype,
    )


class _FromNVFP4Func(torch.autograd.Function):
    """Cast from NVFP4 to other dtype"""

    @staticmethod
    def forward(
        _ctx: Optional[torch.autograd.function.FunctionCtx],  # unused
        tensor: NVFP4TensorStorage,
        dtype: torch.dtype,
    ) -> torch.Tensor:
        # pylint: disable=missing-function-docstring
        if tensor._rowwise_data is not None and tensor._rowwise_data.numel() == 0:
            return torch.empty(tensor.size(), dtype=dtype, device=tensor.device)
        if tensor._columnwise_data is not None and tensor._columnwise_data.numel() == 0:
            return torch.empty(tensor.size(), dtype=dtype, device=tensor.device)

        if tensor._rowwise_data is None and tensor._columnwise_data is None:
            raise ValueError("Attempted to dequantize NVFP4 tensor with no data")
        if tensor._rowwise_data is None and tensor._columnwise_data is not None:
            raise NotImplementedError("Dequantizing column-wise NVFP4 data is not implemented yet!")

        # ``tex.dequantize`` requires CUDA-resident buffers. If the tensor has
        src_device = tensor.device
        if src_device.type != "cuda":
            cuda_tensor = tensor.to(device=torch.device("cuda"))
            result = tex.dequantize(cuda_tensor, torch_to_transformer_engine_dtype[dtype])
            return result.to(device=src_device)
        return tex.dequantize(tensor, torch_to_transformer_engine_dtype[dtype])

    @staticmethod
    def backward(
        _ctx: torch.autograd.function.FunctionCtx,  # unused
        grad: torch.Tensor,
    ) -> Tuple[Optional[torch.Tensor], ...]:
        # pylint: disable=missing-function-docstring
        # Assume that we want gradients in full precision
        return grad, None


class NVFP4TensorStorage(QuantizedTensorStorage):
    """Mixin class that holds data attributes of NVFP4Tensor.

    NVFP4Tensor inherits from the PyTorch tensor class and this mixin
    class. If this class is instantiated directly, it has the same
    data, lower CPU overhead, and less functionality. It should only
    be instantiated directly for performance-critical internal usage.

    """

    # Row-scaled FP4 data and its block scaling factors
    _rowwise_data: Annotated[Optional[torch.Tensor], InnerTensor("rowwise_data")]
    _rowwise_scale_inv: Annotated[torch.Tensor, InnerTensor("rowwise_scale_inv")]
    # Column-scaled FP4 data and its block scaling factors
    _columnwise_data: Annotated[Optional[torch.Tensor], InnerTensor("columnwise_data")]
    _columnwise_scale_inv: Annotated[torch.Tensor, InnerTensor("columnwise_scale_inv")]
    # Input absolute maximum values, used to compute the tensor scale
    _amax_rowwise: Annotated[torch.Tensor, InnerTensor("amax_rowwise")]
    _amax_columnwise: Annotated[torch.Tensor, InnerTensor("amax_columnwise")]

    # Builder class for casting to MXFP8
    _quantizer: Optional[Quantizer]
    # FP4 data type
    _fp4_dtype: DType
    # Whether scaling factors are in the swizzled format expected by
    # GEMM
    _with_gemm_swizzled_scales: bool
    # Whether this NVFP4 tensor uses row-scaled amax metadata
    _row_scaled_nvfp4: bool
    # Whether this NVFP4 tensor uses 4over6 map-to-4/map-to-6 block selection
    _nvfp4_use_4over6: bool
    # Global E4M3 scale bound used by this NVFP4 tensor
    _nvfp4_e4m3_max: int

    def __new__(
        cls,
        rowwise_data: Optional[torch.Tensor],
        rowwise_scale_inv: torch.Tensor,
        columnwise_data: Optional[torch.Tensor],
        columnwise_scale_inv: torch.Tensor,
        amax_rowwise: torch.Tensor,
        amax_columnwise: torch.Tensor,
        fp4_dtype: Union[DType, tex.DType],
        quantizer: Optional[Quantizer],
        with_gemm_swizzled_scales: bool,
        *args,
        fake_dtype: Optional[torch.dtype] = None,
        row_scaled_nvfp4: bool = False,
        nvfp4_use_4over6: bool = False,
        nvfp4_e4m3_max: int = 448,
        **kwargs,
    ):
        if cls is NVFP4TensorStorage:
            instance = object.__new__(cls)
            instance._dtype = fake_dtype if fake_dtype is not None else torch.float32
        else:
            instance = super().__new__(cls, *args, fake_dtype=fake_dtype, **kwargs)

        instance._rowwise_data = rowwise_data
        instance._columnwise_data = columnwise_data
        instance._fp4_dtype = DType.cast(fp4_dtype)
        instance._quantizer = quantizer.copy() if quantizer is not None else None
        instance._rowwise_scale_inv = rowwise_scale_inv
        instance._columnwise_scale_inv = columnwise_scale_inv
        instance._amax_rowwise = amax_rowwise
        instance._amax_columnwise = amax_columnwise
        instance._with_gemm_swizzled_scales = with_gemm_swizzled_scales
        instance._row_scaled_nvfp4 = row_scaled_nvfp4
        instance._nvfp4_use_4over6 = nvfp4_use_4over6
        instance._nvfp4_e4m3_max = nvfp4_e4m3_max if nvfp4_use_4over6 else 448

        return instance

    def clear(self):
        """Deallocate this tensor's memory. Typically not needed and must be used carefully."""
        for t in (
            self._rowwise_data,
            self._columnwise_data,
            self._rowwise_scale_inv,
            self._columnwise_scale_inv,
            self._amax_rowwise,
            self._amax_columnwise,
        ):
            if t is not None:
                t.data = _empty_tensor()

    def copy_from_storage(self, src: QuantizedTensorStorage) -> None:
        """Copy data buffers from another NVFP4TensorStorage."""
        if not isinstance(src, NVFP4TensorStorage):
            raise TypeError("copy_from_storage expects NVFP4TensorStorage")
        if self._fp4_dtype != src._fp4_dtype:
            raise RuntimeError("FP4 dtype mismatch in copy_from_storage")
        if self._with_gemm_swizzled_scales != src._with_gemm_swizzled_scales:
            raise RuntimeError("Scale layout mismatch in copy_from_storage")
        if self._row_scaled_nvfp4 != src._row_scaled_nvfp4:
            raise RuntimeError("Rowwise amax scaling mode mismatch in copy_from_storage")
        if self._nvfp4_use_4over6 != src._nvfp4_use_4over6:
            raise RuntimeError("NVFP4 4over6 mode mismatch in copy_from_storage")
        if self._nvfp4_e4m3_max != src._nvfp4_e4m3_max:
            raise RuntimeError("NVFP4 4over6 E4M3 scale bound mismatch in copy_from_storage")

        def _copy_optional(dst: Optional[torch.Tensor], src_tensor: Optional[torch.Tensor]):
            if dst is not None and src_tensor is not None:
                dst.copy_(src_tensor)

        _copy_optional(self._rowwise_data, src._rowwise_data)
        _copy_optional(self._columnwise_data, src._columnwise_data)
        _copy_optional(self._rowwise_scale_inv, src._rowwise_scale_inv)
        _copy_optional(self._columnwise_scale_inv, src._columnwise_scale_inv)
        _copy_optional(self._amax_rowwise, src._amax_rowwise)
        _copy_optional(self._amax_columnwise, src._amax_columnwise)

    def get_metadata(self) -> Dict[str, Any]:
        """Get this tensor's metadata."""
        return {
            "rowwise_data": self._rowwise_data,
            "rowwise_scale_inv": self._rowwise_scale_inv,
            "columnwise_data": self._columnwise_data,
            "columnwise_scale_inv": self._columnwise_scale_inv,
            "amax_rowwise": self._amax_rowwise,
            "amax_columnwise": self._amax_columnwise,
            "fp4_dtype": self._fp4_dtype,
            "quantizer": self._quantizer,
            "with_gemm_swizzled_scales": self._with_gemm_swizzled_scales,
            "row_scaled_nvfp4": self._row_scaled_nvfp4,
            "nvfp4_use_4over6": self._nvfp4_use_4over6,
            "nvfp4_e4m3_max": self._nvfp4_e4m3_max,
            "fake_dtype": self._dtype,
        }

    def prepare_for_saving(self) -> Tuple[list[Optional[torch.Tensor]], NVFP4TensorStorage]:
        """Prepare the tensor base for saving for backward"""
        tensors = [
            self._rowwise_data,
            self._columnwise_data,
            self._rowwise_scale_inv,
            self._columnwise_scale_inv,
            self._amax_rowwise,
            self._amax_columnwise,
        ]
        self._rowwise_data = None
        self._columnwise_data = None
        self._rowwise_scale_inv = None
        self._columnwise_scale_inv = None
        self._amax_rowwise = None
        self._amax_columnwise = None
        return tensors, self

    def restore_from_saved(
        self, tensors: list[Optional[torch.Tensor]]
    ) -> list[Optional[torch.Tensor]]:
        """Restore the tensor base data from the saved tensors list."""
        self._rowwise_data = tensors[0]
        self._columnwise_data = tensors[1]
        self._rowwise_scale_inv = tensors[2]
        self._columnwise_scale_inv = tensors[3]
        self._amax_rowwise = tensors[4]
        self._amax_columnwise = tensors[5]
        return tensors[6:]

    def get_data_tensors(self):
        """Get this Tensor's data."""
        return self._rowwise_data, self._columnwise_data

    def dequantize(self, *, dtype: Optional[torch.dtype] = None) -> torch.Tensor:
        """Dequantize to a higher precision."""
        if dtype is None:
            dtype = self._dtype
        if self._rowwise_data is not None and self._rowwise_data.numel() == 0:
            return torch.empty(self.size(), dtype=dtype, device=self.device)
        return _FromNVFP4Func.forward(None, self, dtype)

    def size(self, dim: Optional[int] = None) -> Union[torch.Size, int]:
        # pylint: disable=missing-function-docstring

        # Infer tensor shape
        shape = None
        if self._rowwise_data is not None:
            byte_shape = list(self._rowwise_data.size())
            shape = byte_shape[:-1] + [byte_shape[-1] * 2]
        elif self._columnwise_data is not None:
            warnings.warn("Attempting to get shape of NVFP4 tensor with only column-wise data.")
            byte_shape = list(self._columnwise_data.size())
            shape = byte_shape[1:-1] + [byte_shape[-1] * 2, byte_shape[0]]
        if shape is None:
            raise RuntimeError("Attempted to get shape of NVFP4 tensor with no data")

        # Return shape or dim
        if dim is None:
            return torch.Size(shape)
        return shape[dim]

    @property
    def device(self):
        """Return the device of the tensor. Define this to avoid expensive PyObject lookups."""
        if self._rowwise_data is not None:
            return self._rowwise_data.device
        if self._columnwise_data is not None:
            return self._columnwise_data.device
        raise RuntimeError("NVFP4TensorStorage has no data!")

    def view(self, shape: torch.Size):
        # pylint: disable=missing-function-docstring

        # Return input tensor if view not needed
        cur_shape = self.size()
        if shape is None or shape == cur_shape:
            return self

        # Canonicalize shape
        if not isinstance(shape, Iterable):
            shape = [shape]
        elif len(shape) == 1 and isinstance(shape[0], Iterable):
            shape = shape[0]
        if -1 in shape:
            shape = list(shape)
            d_inferred = -math.prod(cur_shape) // math.prod(shape)
            for i, d in enumerate(shape):
                if d == -1:
                    shape[i] = d_inferred
                    break
        if shape[-1] != cur_shape[-1]:
            raise RuntimeError(
                "NVFP4Tensor does not support reshaping inner dimension "
                f"(attempted to reshape dims={tuple(cur_shape)} to {tuple(shape)})"
            )

        # Reshape data
        new_rowwise_data = None
        new_columnwise_data = None
        if self._rowwise_data is not None:
            if shape[-1] % 2 != 0:
                raise ValueError(
                    "Cannot represent row-wise data for NVFP4 tensor "
                    f"with shape={shape} as byte array."
                )
            byte_shape = list(shape[:-1]) + [shape[-1] // 2]
            new_rowwise_data = self._rowwise_data.view(byte_shape)
        if self._columnwise_data is not None:
            columnwise_shape = (shape[-1], math.prod(shape[:-1]))
            if columnwise_shape[-1] % 2 != 0:
                raise ValueError(
                    "Cannot represent column-wise data for NVFP4 tensor "
                    f"with shape={shape} as byte array."
                )
            byte_shape = (columnwise_shape[0], columnwise_shape[1] // 2)
            new_columnwise_data = self._columnwise_data.view(byte_shape)

        # Construct tensor
        return NVFP4TensorStorage(
            rowwise_data=new_rowwise_data,
            rowwise_scale_inv=self._rowwise_scale_inv,
            columnwise_data=new_columnwise_data,
            columnwise_scale_inv=self._columnwise_scale_inv,
            amax_rowwise=self._amax_rowwise,
            amax_columnwise=self._amax_columnwise,
            quantizer=self._quantizer,
            fp4_dtype=self._fp4_dtype,
            with_gemm_swizzled_scales=self._with_gemm_swizzled_scales,
            row_scaled_nvfp4=self._row_scaled_nvfp4,
            nvfp4_use_4over6=self._nvfp4_use_4over6,
            nvfp4_e4m3_max=self._nvfp4_e4m3_max,
            fake_dtype=self._dtype,
        )

    def __repr__(self):
        try:
            data_rowwise = self.dequantize()

            return (
                "NVFP4TensorStorage("
                f"rowwise_scaled_data={data_rowwise},"
                f"rowwise_scale_inv={self._rowwise_scale_inv},"
                f"amax_rowwise={self._amax_rowwise},"
                f"amax_columnwise={self._amax_columnwise},"
                ")"
            )
        except Exception as exc:  # pylint: disable=broad-except
            return safe_quantized_repr(self, "NVFP4TensorStorage", error=exc)

    def update_usage(
        self,
        rowwise_usage: Optional[bool] = None,
        columnwise_usage: Optional[bool] = None,
    ):
        """
        For the NVFP4 format, columnwise scaled output is only produced by x2
        scaling kernels, so this function only disables usages.
        """

        # Default usage is based on available data
        if rowwise_usage is None:
            rowwise_usage = self._rowwise_data is not None
        if columnwise_usage is None:
            columnwise_usage = self._columnwise_data is not None

        # If both rowwise and columnwise are requested, create columnwise from rowwise if needed
        if rowwise_usage and columnwise_usage:
            if (
                self._rowwise_data is None
                or self._rowwise_scale_inv is None
                or self._amax_rowwise is None
            ):
                raise RuntimeError(
                    "Cannot update to rowwise and columnwise usage because rowwise data is None."
                )
            if self._columnwise_data is None or self._columnwise_scale_inv is None:
                self._create_columnwise()
                return

        # Update row-scaled data
        if rowwise_usage:
            if self._rowwise_data is None:
                raise RuntimeError(
                    "Requested row-wise usage, but NVFP4Tensor is missing row-scaled NVFP4 data"
                )
            if self._rowwise_scale_inv is None:
                raise RuntimeError(
                    "Requested row-wise usage, but NVFP4Tensor is missing row-scaled scale-inverses"
                )
            if self._amax_rowwise is None:
                raise RuntimeError(
                    "Requested row-wise usage, but NVFP4Tensor is missing per tensor"
                    " row-scaled scale-inverse"
                )
        else:
            self._rowwise_data = None
            self._rowwise_scale_inv = None
            self._amax_rowwise = None

        # Update column-scaled data
        if columnwise_usage:
            if self._columnwise_data is None:
                raise RuntimeError(
                    "Requested column-wise usage, but NVFP4Tensor is missing column-scaled FP8 data"
                )
            if self._columnwise_scale_inv is None:
                raise RuntimeError(
                    "Requested column-wise usage, "
                    "but NVFP4Tensor is missing column-scaled scale-inverses"
                )
            if self._amax_columnwise is None:
                raise RuntimeError(
                    "Requested column-wise usage, "
                    "but NVFP4Tensor is missing per tensor column-scaled scale-inverse"
                )
        else:
            self._columnwise_data = None
            self._columnwise_scale_inv = None
            self._amax_columnwise = None

    def _create_columnwise(self):
        """
        Update columnwise data and columnwise scale inv. Can only be used when using 2D scaling.
        """
        if self._quantizer is None or not self._quantizer.with_2d_quantization:
            raise RuntimeError("Cannot create columnwise data without 2D quantization enabled.")
        rowwise_data = self._rowwise_data
        if not rowwise_data.is_contiguous():
            rowwise_data = rowwise_data.contiguous()
        # NVFP4 requires a specialized transpose that handles nibble repacking
        self._columnwise_data = tex.nvfp4_data_transpose(rowwise_data, out=self._columnwise_data)
        if self._columnwise_scale_inv is None:
            if self._quantizer is None:
                raise RuntimeError("Cannot create columnwise scale inverse: quantizer is None.")
            # Use logical shape (self.size()), not packed byte shape (rowwise_data.shape)
            # NVFP4 packs 2 elements per byte, so rowwise_data.shape[-1] is K/2
            logical_shape = self.size()
            columnwise_scale_inv_shape = self._quantizer.get_scale_shape(logical_shape, True)
            self._columnwise_scale_inv = torch.empty(
                columnwise_scale_inv_shape,
                dtype=self._rowwise_scale_inv.dtype,
                device=self._rowwise_scale_inv.device,
            )
        if len(self._rowwise_scale_inv.shape) != 2:
            raise ValueError(
                "Expected rowwise_scale_inv to be 2D, but got"
                f" {len(self._rowwise_scale_inv.shape)}D with shape"
                f" {self._rowwise_scale_inv.shape}."
            )
        if len(self._columnwise_scale_inv.shape) != 2:
            raise ValueError(
                "Expected columnwise_scale_inv to be 2D, but got"
                f" {len(self._columnwise_scale_inv.shape)}D with shape"
                f" {self._columnwise_scale_inv.shape}."
            )

        # rowwise_scale_inv has shape [M_padded, K_tiles] where each tile's scale
        # is repeated 16 times (once per row in the 16x16 tile).
        # columnwise_scale_inv has shape [K_padded, M_tiles] where scales are
        # repeated 16 times per tile row.
        TILE_SIZE = 16
        logical_shape = self.size()
        M, K = logical_shape[0], logical_shape[-1]
        M_tiles = (M + TILE_SIZE - 1) // TILE_SIZE
        K_tiles = (K + TILE_SIZE - 1) // TILE_SIZE

        tex.nvfp4_2d_scale_transpose(
            self._rowwise_scale_inv,
            self._columnwise_scale_inv,
            M_tiles,
            K_tiles,
        )

        # Also set columnwise amax (same as rowwise since it's just transposed data)
        if self._amax_columnwise is None:
            self._amax_columnwise = torch.empty_like(self._amax_rowwise)
        self._amax_columnwise.copy_(self._amax_rowwise)
