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

"""Mixin class holding data specific for MXFP8Tensor"""

from __future__ import annotations
from typing import Annotated, Optional, Dict, Any, Tuple, Union
from collections.abc import Iterable
import math
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,
    MXFP8_BLOCK_SCALING_SIZE,
    DType,
)

from ...utils import _empty_tensor


class _FromMXFP8Func(torch.autograd.Function):
    """Cast from MXFP8 to other dtype"""

    @staticmethod
    def forward(
        _ctx: Optional[torch.autograd.function.FunctionCtx],  # unused
        tensor: MXFP8TensorStorage,
        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("Cannot dequantize MXFP8 tensor with no data")
        te_dtype = torch_to_transformer_engine_dtype[dtype]
        # ``tex.dequantize`` requires CUDA-resident buffers.
        src_device = tensor.device
        if src_device.type != "cuda":
            cuda_tensor = tensor.to(device=torch.device("cuda"))
            result = tex.dequantize(cuda_tensor, te_dtype)
            return result.to(device=src_device)
        return tex.dequantize(tensor, te_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 MXFP8TensorStorage(QuantizedTensorStorage):
    """Mixin class that holds data attributes of MXFP8Tensor.

    MXFP8Tensor 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 FP8 data and its scaling factors
    _rowwise_data: Annotated[Optional[torch.Tensor], InnerTensor("rowwise_data")]
    _rowwise_scale_inv: Annotated[torch.Tensor, InnerTensor("rowwise_scale_inv")]
    # Column-scaled FP8 data and its scaling factors
    _columnwise_data: Annotated[Optional[torch.Tensor], InnerTensor("columnwise_data")]
    _columnwise_scale_inv: Annotated[torch.Tensor, InnerTensor("columnwise_scale_inv")]

    # Builder class for casting to MXFP8
    _quantizer: Optional[Quantizer]
    # FP8 data type
    _fp8_dtype: DType
    # Whether scaling factors are in the swizzled format expected by
    # GEMM
    _with_gemm_swizzled_scales: bool

    def __new__(
        cls,
        rowwise_data: Optional[torch.Tensor],
        rowwise_scale_inv: Optional[torch.Tensor],
        columnwise_data: Optional[torch.Tensor],
        columnwise_scale_inv: Optional[torch.Tensor],
        fp8_dtype: Union[DType, tex.DType],
        quantizer: Optional[Quantizer],
        with_gemm_swizzled_scales: bool,
        *args,
        fake_dtype: Optional[torch.dtype] = None,
        **kwargs,
    ):
        if cls is MXFP8TensorStorage:
            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._rowwise_scale_inv = rowwise_scale_inv
        instance._columnwise_scale_inv = columnwise_scale_inv
        instance._quantizer = quantizer.copy() if quantizer is not None else None
        instance._fp8_dtype = DType.cast(fp8_dtype)
        instance._with_gemm_swizzled_scales = with_gemm_swizzled_scales

        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,
        ):
            if t is not None:
                t.data = _empty_tensor()

    def copy_from_storage(self, src: QuantizedTensorStorage) -> None:
        """Copy data buffers from another MXFP8TensorStorage."""
        if not isinstance(src, MXFP8TensorStorage):
            raise TypeError("copy_from_storage expects MXFP8TensorStorage")
        if self._fp8_dtype != src._fp8_dtype:
            raise RuntimeError("FP8 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")

        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)

    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,
            "fp8_dtype": self._fp8_dtype,
            "quantizer": self._quantizer,
            "with_gemm_swizzled_scales": self._with_gemm_swizzled_scales,
            "fake_dtype": self._dtype,
        }

    def prepare_for_saving(self) -> Tuple[list[Optional[torch.Tensor]], MXFP8TensorStorage]:
        """Prepare the tensor base for saving for backward"""
        tensors = [
            self._rowwise_data,
            self._columnwise_data,
            self._rowwise_scale_inv,
            self._columnwise_scale_inv,
        ]
        self._rowwise_data = None
        self._columnwise_data = None
        self._rowwise_scale_inv = None
        self._columnwise_scale_inv = 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]
        return tensors[4:]

    def get_data_tensors(self, rowwise_data: bool = True, columnwise_data: bool = True):
        """Get this Tensor's data."""
        if rowwise_data and columnwise_data:
            return self._rowwise_data, self._columnwise_data
        if rowwise_data:
            return self._rowwise_data
        if columnwise_data:
            return self._columnwise_data
        raise ValueError("No data to get, both rowwise_data and columnwise_data are False")

    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 _FromMXFP8Func.forward(None, self, dtype)

    def size(self, *args, **kwargs):
        # pylint: disable=missing-function-docstring
        if self._rowwise_data is not None:
            return self._rowwise_data.size(*args, **kwargs)
        return self._columnwise_data.size(*args, **kwargs)

    @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("MXFP8TensorStorage 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(
                "MXFP8Tensor does not support reshaping inner dimension "
                f"(attempted to reshape dims={tuple(cur_shape)} to {tuple(shape)})"
            )

        # Construct new tensor
        cur_rowwise_data = self._rowwise_data
        cur_columnwise_data = self._columnwise_data
        new_rowwise_data = None
        new_columnwise_data = None
        if cur_rowwise_data is not None:
            new_rowwise_data = cur_rowwise_data.view(*shape)
        if cur_columnwise_data is not None:
            new_columnwise_data = cur_columnwise_data.view(*shape)

        return MXFP8TensorStorage(
            rowwise_data=new_rowwise_data,
            rowwise_scale_inv=self._rowwise_scale_inv,
            columnwise_data=new_columnwise_data,
            columnwise_scale_inv=self._columnwise_scale_inv,
            fp8_dtype=self._fp8_dtype,
            quantizer=self._quantizer,
            with_gemm_swizzled_scales=self._with_gemm_swizzled_scales,
            fake_dtype=self._dtype,
        )

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

            return (
                "MXFP8TensorStorage("
                f"fp8_dtype={self._fp8_dtype}, "
                f"rowwise_scaled_data={data_rowwise}"
                f"rowwise_scale_inv={self._rowwise_scale_inv}, "
                ")"
            )
        except Exception as exc:  # pylint: disable=broad-except
            return safe_quantized_repr(self, "MXFP8TensorStorage", error=exc)

    def update_usage(
        self,
        rowwise_usage: Optional[bool] = None,
        columnwise_usage: Optional[bool] = None,
    ):
        """
        For MXFP8, 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

        # Update row-scaled data
        if rowwise_usage:
            if self._rowwise_data is None:
                raise RuntimeError(
                    "Requested row-wise usage, but MXFP8Tensor is missing row-scaled FP8 data"
                )
            if self._rowwise_scale_inv is None:
                raise RuntimeError(
                    "Requested row-wise usage, but MXFP8Tensor is missing row-scaled scale-inverses"
                )
        else:
            self._rowwise_data = None
            self._rowwise_scale_inv = None

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

    def get_usages(self) -> Dict[str, bool]:
        """Get the usage of the tensor"""
        return {
            "rowwise": self._rowwise_data is not None,
            "columnwise": self._columnwise_data is not None,
        }

    def fsdp_buffer_fields(self) -> Tuple[str, ...]:
        """Fields gathered by FSDP2 for MXFP8.

        Block scales are per-block and direction-specific — each direction
        gathers both its data buffer and its scale-inv buffer. ``None``-valued
        directions (e.g. a columnwise-only sub-storage in hybrid quantization)
        are excluded so the gather tuple only contains real tensors.
        """
        fields = []
        if self._rowwise_data is not None:
            fields.extend(("_rowwise_data", "_rowwise_scale_inv"))
        if self._columnwise_data is not None:
            fields.extend(("_columnwise_data", "_columnwise_scale_inv"))
        return tuple(fields)

    def fsdp_extract_buffers(
        self,
    ) -> Tuple[Tuple[Optional[torch.Tensor], ...], Dict[str, Any]]:
        """Extract MXFP8 buffers, unpadding block-scale alignment before gather.

        MXFP8 kernels require scale-inv tensors aligned to ``[128, 4]``
        (rowwise) and ``[4, 128]`` (columnwise). That padding is attached to
        the local shard but would produce misaligned concatenation under
        FSDP2's dim-0 all-gather. Strip it here and re-apply in
        :meth:`fsdp_assign_gathered`.
        """
        if self._with_gemm_swizzled_scales:
            raise NotImplementedError(
                "FSDP2 is only supported for MXFP8Tensors with compact scales"
            )
        names = self.fsdp_buffer_fields()
        buffers = []
        shape = self.size()
        flattened_in_shape0 = math.prod(shape[:-1])
        for name in names:
            t = getattr(self, name)
            if name == "_rowwise_scale_inv" and t is not None:
                if t.size(0) != flattened_in_shape0:
                    t = t[:flattened_in_shape0]
            elif name == "_columnwise_scale_inv" and t is not None:
                expected = math.ceil(flattened_in_shape0 / MXFP8_BLOCK_SCALING_SIZE)
                if t.size(0) != expected:
                    t = t[:expected]
            buffers.append(t)
        return tuple(buffers), {"field_names": names}

    def fsdp_assign_gathered(
        self,
        gathered: Tuple[Optional[torch.Tensor], ...],
        meta: Dict[str, Any],
    ) -> None:
        """Write gathered MXFP8 buffers back, re-padding block scales.

        Inverse of :meth:`fsdp_extract_buffers`: the gathered scale-inv tensors
        are padded back up to ``[128, 4]`` / ``[4, 128]`` alignment before
        being assigned to the storage.
        """
        names = meta["field_names"]
        if len(names) != len(gathered):
            raise RuntimeError(
                "MXFP8TensorStorage.fsdp_assign_gathered got "
                f"{len(gathered)} buffers for {len(names)} fields"
            )
        for name, buf in zip(names, gathered):
            if buf is not None and name == "_rowwise_scale_inv":
                pad = (128 - buf.size(0) % 128) % 128
                if pad > 0:
                    buf = torch.nn.functional.pad(buf, (0, 0, 0, pad))
            elif buf is not None and name == "_columnwise_scale_inv":
                pad = (4 - buf.size(0) % 4) % 4
                if pad > 0:
                    buf = torch.nn.functional.pad(buf, (0, 0, 0, pad))
            setattr(self, name, buf)
