"""
Native FP8 Operations for PyTorch 2.1+

This module provides actual FP8 computation using PyTorch's native FP8 types:
- torch.float8_e4m3fn: 8-bit float with 4-bit exponent, 3-bit mantissa (higher precision)
- torch.float8_e5m2: 8-bit float with 5-bit exponent, 2-bit mantissa (wider range)

Key Features:
- Real FP8 quantization and dequantization
- FP8 linear layers with actual FP8 GEMM operations
- Dynamic scaling for numerical stability
- Inference pipeline for FP8 model serving
- Hardware-aware optimizations for H100/Blackwell

References:
    - PyTorch FP8: https://pytorch.org/docs/stable/torch.html#fp8-dtypes
    - FP8 Formats Paper: https://arxiv.org/abs/2209.05433
    - NVIDIA FP8 Training: https://developer.nvidia.com/blog/nvidia-h100-transformer-engine/
"""

import logging
import warnings
from dataclasses import dataclass
from enum import Enum
from typing import Any

import torch
import torch.nn as nn
import torch.nn.functional as F

logger = logging.getLogger(__name__)

# Check PyTorch version for FP8 support
PYTORCH_VERSION = tuple(int(x) for x in torch.__version__.split('.')[:2])
FP8_NATIVE_AVAILABLE = PYTORCH_VERSION >= (2, 1)

# FP8 dtype availability check
try:
    _E4M3_DTYPE: torch.dtype | None = torch.float8_e4m3fn
    _E5M2_DTYPE: torch.dtype | None = torch.float8_e5m2
    FP8_DTYPES_AVAILABLE = True
except AttributeError:
    _E4M3_DTYPE = None
    _E5M2_DTYPE = None
    FP8_DTYPES_AVAILABLE = False
    warnings.warn(
        f"PyTorch {torch.__version__} does not support native FP8 types. "
        "FP8 operations will use simulated quantization. "
        "Upgrade to PyTorch 2.1+ for native FP8 support.",
    stacklevel=2,
    )

# Check for scaled_mm availability (PyTorch 2.4+)
FP8_SCALED_MM_AVAILABLE = hasattr(torch, '_scaled_mm') or hasattr(torch.ops.aten, '_scaled_mm')


class FP8Dtype(Enum):
    """FP8 data types"""
    E4M3 = "e4m3fn"  # Higher precision, narrower range
    E5M2 = "e5m2"    # Lower precision, wider range


@dataclass
class FP8TensorSpec:
    """Specification for FP8 tensor with scale"""
    dtype: FP8Dtype
    scale: torch.Tensor
    amax: torch.Tensor

    @property
    def max_value(self) -> float:
        """Get maximum representable value for this FP8 format"""
        if self.dtype == FP8Dtype.E4M3:
            return 448.0  # Max for E4M3FN
        else:  # E5M2
            return 57344.0  # Max for E5M2


def get_fp8_dtype(format: FP8Dtype) -> torch.dtype | None:
    """Get PyTorch FP8 dtype for format"""
    if not FP8_DTYPES_AVAILABLE:
        return None

    if format == FP8Dtype.E4M3:
        return _E4M3_DTYPE
    else:
        return _E5M2_DTYPE


def compute_fp8_scale(
    tensor: torch.Tensor,
    format: FP8Dtype,
    margin: int = 0
) -> torch.Tensor:
    """
    Compute optimal scale factor for FP8 quantization.

    Args:
        tensor: Input tensor to compute scale for
        format: FP8 format (E4M3 or E5M2)
        margin: Safety margin in bits (0-3 recommended)

    Returns:
        Scale factor tensor
    """
    # Compute AMAX (maximum absolute value)
    amax = tensor.abs().max()

    # Get max representable value for format
    if format == FP8Dtype.E4M3:
        max_val = 448.0
    else:
        max_val = 57344.0

    # Compute scale with safety margin
    safety_factor = 2 ** margin
    scale = max_val / (amax * safety_factor + 1e-12)

    return scale.clamp(min=1e-12, max=1e12)


def quantize_to_fp8(
    tensor: torch.Tensor,
    scale: torch.Tensor,
    format: FP8Dtype
) -> tuple[torch.Tensor, torch.Tensor]:
    """
    Quantize tensor to FP8 format.

    Args:
        tensor: Input tensor (float32/float16/bfloat16)
        scale: Scaling factor
        format: FP8 format

    Returns:
        Tuple of (quantized tensor, scale used)
    """
    # Scale tensor
    scaled_tensor = tensor * scale

    if FP8_DTYPES_AVAILABLE:
        # Use native FP8 type
        fp8_dtype = get_fp8_dtype(format)
        quantized = scaled_tensor.to(fp8_dtype)
    else:
        # Simulate FP8 quantization via clamping and rounding
        if format == FP8Dtype.E4M3:
            max_val = 448.0
        else:
            max_val = 57344.0

        # Clamp to FP8 range
        quantized = scaled_tensor.clamp(-max_val, max_val)

        # Simulate reduced precision via rounding
        # This is a simplified simulation - real FP8 has different precision
        if format == FP8Dtype.E4M3:
            # E4M3 has ~3 bits of mantissa precision
            quantized = torch.round(quantized * 8) / 8
        else:
            # E5M2 has ~2 bits of mantissa precision
            quantized = torch.round(quantized * 4) / 4

    return quantized, scale


def dequantize_from_fp8(
    tensor: torch.Tensor,
    scale: torch.Tensor,
    output_dtype: torch.dtype = torch.float32
) -> torch.Tensor:
    """
    Dequantize FP8 tensor back to higher precision.

    Args:
        tensor: FP8 tensor
        scale: Scale factor used during quantization
        output_dtype: Output data type

    Returns:
        Dequantized tensor
    """
    return tensor.to(output_dtype) / scale


class FP8QuantizedTensor:
    """
    Wrapper for FP8 quantized tensor with scale tracking.

    This class encapsulates an FP8 tensor along with its scale factor,
    enabling proper arithmetic operations and dequantization.
    """

    def __init__(
        self,
        data: torch.Tensor,
        scale: torch.Tensor,
        format: FP8Dtype
    ):
        self.data = data
        self.scale = scale
        self.format = format

    def dequantize(self, dtype: torch.dtype = torch.float32) -> torch.Tensor:
        """Dequantize to specified dtype"""
        return dequantize_from_fp8(self.data, self.scale, dtype)

    @classmethod
    def from_tensor(
        cls,
        tensor: torch.Tensor,
        format: FP8Dtype = FP8Dtype.E4M3,
        scale: torch.Tensor | None = None
    ) -> 'FP8QuantizedTensor':
        """Create FP8 quantized tensor from float tensor"""
        if scale is None:
            scale = compute_fp8_scale(tensor, format)

        quantized, scale = quantize_to_fp8(tensor, scale, format)
        return cls(quantized, scale, format)

    @property
    def shape(self) -> torch.Size:
        return self.data.shape

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

    def to(self, device: torch.device) -> 'FP8QuantizedTensor':
        """Move to device"""
        return FP8QuantizedTensor(
            self.data.to(device),
            self.scale.to(device),
            self.format
        )


class NativeFP8Linear(nn.Module):
    """
    Linear layer with native FP8 computation.

    This layer performs actual FP8 matrix multiplication when supported:
    - Weights are stored in FP8 format
    - Activations are quantized to FP8 on-the-fly
    - Uses scaled_mm for FP8 GEMM when available (PyTorch 2.4+)
    - Falls back to dequantize-compute-quantize otherwise

    Args:
        in_features: Size of input features
        out_features: Size of output features
        bias: Whether to include bias
        weight_format: FP8 format for weights (default E4M3)
        activation_format: FP8 format for activations (default E4M3)
        device: Target device

    Example:
        >>> layer = NativeFP8Linear(512, 256, device=torch.device("cpu"))
        >>> x = torch.randn(32, 512)
        >>> output = layer(x)
        >>> output.shape
        torch.Size([32, 256])
    """

    def __init__(
        self,
        in_features: int,
        out_features: int,
        bias: bool = True,
        weight_format: FP8Dtype = FP8Dtype.E4M3,
        activation_format: FP8Dtype = FP8Dtype.E4M3,
        device: torch.device | None = None
    ):
        super().__init__()

        self.in_features = in_features
        self.out_features = out_features
        self.weight_format = weight_format
        self.activation_format = activation_format

        device = device or torch.device('cuda' if torch.cuda.is_available() else 'cpu')

        # Store master weights in high precision for updates
        self.register_buffer(
            'weight_master',
            torch.randn(out_features, in_features, device=device)
        )

        # FP8 weight buffer (will be quantized from master)
        self.register_buffer(
            'weight_fp8',
            torch.zeros(out_features, in_features, device=device)
        )

        # Weight scale
        self.register_buffer(
            'weight_scale',
            torch.ones(1, device=device)
        )

        # Activation scale (updated during forward)
        self.register_buffer(
            'activation_scale',
            torch.ones(1, device=device)
        )

        # AMAX tracking for dynamic scaling
        self.register_buffer(
            'weight_amax',
            torch.zeros(1, device=device)
        )
        self.register_buffer(
            'activation_amax',
            torch.zeros(1, device=device)
        )

        # Bias (kept in high precision)
        if bias:
            self.bias = nn.Parameter(torch.zeros(out_features, device=device))
        else:
            self.register_parameter('bias', None)

        # Initialize
        self._init_weights()
        self._quantize_weights()

        # Track if FP8 is actually being used
        self._fp8_active = FP8_DTYPES_AVAILABLE

    def _init_weights(self):
        """Initialize weights using Kaiming initialization"""
        nn.init.kaiming_uniform_(self.weight_master, a=5**0.5)
        if self.bias is not None:
            nn.init.zeros_(self.bias)

    def _quantize_weights(self):
        """Quantize master weights to FP8"""
        # Compute scale
        self.weight_scale.data = compute_fp8_scale(
            self.weight_master,
            self.weight_format,
            margin=1  # Small safety margin
        )

        # Update AMAX
        self.weight_amax.data = self.weight_master.abs().max()

        # Quantize
        quantized, _ = quantize_to_fp8(
            self.weight_master,
            self.weight_scale,
            self.weight_format
        )
        self.weight_fp8.data = quantized

    def _update_activation_scale(self, input: torch.Tensor):
        """Update activation scale based on input statistics"""
        current_amax = input.abs().max()

        # EMA for activation AMAX
        if self.training:
            self.activation_amax.data = (
                0.99 * self.activation_amax + 0.01 * current_amax
            )

        # Compute scale
        self.activation_scale.data = compute_fp8_scale(
            input,
            self.activation_format,
            margin=1
        )

    def forward(self, input: torch.Tensor) -> torch.Tensor:
        """Forward pass with FP8 computation"""

        # Update activation scale
        self._update_activation_scale(input)

        if FP8_SCALED_MM_AVAILABLE and FP8_DTYPES_AVAILABLE and self._fp8_active:
            # Use native FP8 scaled matmul (PyTorch 2.4+)
            output = self._fp8_forward_native(input)
        elif FP8_DTYPES_AVAILABLE and self._fp8_active:
            # FP8 types available but no scaled_mm - use dequantize approach
            output = self._fp8_forward_dequant(input)
        else:
            # Fallback to simulated FP8
            output = self._fp8_forward_simulated(input)

        if self.bias is not None:
            output = output + self.bias

        return output

    def _fp8_forward_native(self, input: torch.Tensor) -> torch.Tensor:
        """Forward with native FP8.

        Note: For training, we use the dequantize approach to preserve gradients,
        since _scaled_mm doesn't have autograd support. For inference, we could
        use _scaled_mm directly for better performance.
        """
        # For training mode, always use dequantize approach for gradient support
        if self.training:
            return self._fp8_forward_dequant(input)

        # For inference, try to use scaled_mm for better performance
        try:
            # Quantize input to FP8
            input_fp8, _ = quantize_to_fp8(
                input,
                self.activation_scale,
                self.activation_format
            )

            # Get FP8 dtypes
            act_dtype = get_fp8_dtype(self.activation_format)
            weight_dtype = get_fp8_dtype(self.weight_format)

            # Convert to FP8
            input_fp8 = input_fp8.to(act_dtype)
            weight_fp8 = self.weight_fp8.to(weight_dtype)

            # Compute output scale
            output_scale = 1.0 / (self.activation_scale * self.weight_scale)

            if hasattr(torch, '_scaled_mm'):
                output = torch._scaled_mm(
                    input_fp8,
                    weight_fp8.t(),
                    scale_a=self.activation_scale,
                    scale_b=self.weight_scale,
                    out_dtype=input.dtype
                )
            else:
                # Fallback to manual scaling
                output = torch.mm(
                    input_fp8.to(input.dtype),
                    weight_fp8.t().to(input.dtype)
                ) * output_scale

            return output
        except Exception:
            logger.debug("Native FP8 forward pass failed, falling back to dequantize approach", exc_info=True)
            # If native fails, use dequantize approach
            return self._fp8_forward_dequant(input)

    def _fp8_forward_dequant(self, input: torch.Tensor) -> torch.Tensor:
        """Forward with FP8 quantization and dequantization"""
        # Quantize input
        input_fp8, _ = quantize_to_fp8(
            input,
            self.activation_scale,
            self.activation_format
        )

        # Dequantize for computation
        input_deq = dequantize_from_fp8(input_fp8, self.activation_scale, input.dtype)
        weight_deq = dequantize_from_fp8(self.weight_fp8, self.weight_scale, input.dtype)

        # Standard matmul
        output = F.linear(input_deq, weight_deq)

        return output

    def _fp8_forward_simulated(self, input: torch.Tensor) -> torch.Tensor:
        """Forward with simulated FP8 (no native support)"""
        # Simulate FP8 quantization via clamping and rounding
        input_quant, _ = quantize_to_fp8(
            input,
            self.activation_scale,
            self.activation_format
        )

        weight_quant, _ = quantize_to_fp8(
            self.weight_master,
            self.weight_scale,
            self.weight_format
        )

        # Dequantize for computation
        input_deq = input_quant / self.activation_scale
        weight_deq = weight_quant / self.weight_scale

        output = F.linear(input_deq, weight_deq)

        return output

    def sync_weights(self):
        """Sync FP8 weights from master weights (call after optimizer step)"""
        self._quantize_weights()

    def get_fp8_info(self) -> dict[str, Any]:
        """Get FP8 layer information"""
        return {
            'in_features': self.in_features,
            'out_features': self.out_features,
            'weight_format': self.weight_format.value,
            'activation_format': self.activation_format.value,
            'weight_scale': float(self.weight_scale.item()),
            'activation_scale': float(self.activation_scale.item()),
            'weight_amax': float(self.weight_amax.item()),
            'activation_amax': float(self.activation_amax.item()),
            'fp8_native': FP8_DTYPES_AVAILABLE,
            'fp8_scaled_mm': FP8_SCALED_MM_AVAILABLE,
            'fp8_active': self._fp8_active
        }

    def extra_repr(self) -> str:
        return (
            f'in_features={self.in_features}, out_features={self.out_features}, '
            f'bias={self.bias is not None}, weight_format={self.weight_format.value}, '
            f'fp8_native={FP8_DTYPES_AVAILABLE}'
        )


class FP8InferenceEngine:
    """
    Engine for FP8 model inference.

    This engine handles:
    - Model quantization to FP8
    - Efficient FP8 inference
    - Dynamic batching with FP8
    - Memory-efficient serving

    Args:
        model: PyTorch model to serve
        weight_format: FP8 format for weights
        activation_format: FP8 format for activations
        calibration_data: Optional calibration data for scale computation

    Example:
        >>> model = torch.nn.Sequential(torch.nn.Linear(64, 32))  # doctest: +SKIP
        >>> engine = FP8InferenceEngine(model)  # doctest: +SKIP
        >>> engine.prepare()  # doctest: +SKIP
    """

    def __init__(
        self,
        model: nn.Module,
        weight_format: FP8Dtype = FP8Dtype.E4M3,
        activation_format: FP8Dtype = FP8Dtype.E4M3,
        calibration_data: torch.Tensor | None = None
    ):
        self.model = model
        self.weight_format = weight_format
        self.activation_format = activation_format
        self.calibration_data = calibration_data

        self._prepared = False
        self._fp8_layers = {}
        self._layer_scales = {}

    def prepare(self, device: torch.device | None = None) -> 'FP8InferenceEngine':
        """
        Prepare model for FP8 inference.

        This converts linear layers to FP8 and computes optimal scales.
        """
        device = device or next(self.model.parameters()).device

        # Convert linear layers to FP8
        self._convert_linear_layers(device)

        # Calibrate scales if calibration data provided
        if self.calibration_data is not None:
            self._calibrate_scales()

        self.model.eval()
        self._prepared = True

        return self

    def _convert_linear_layers(self, device: torch.device):
        """Convert nn.Linear to NativeFP8Linear"""

        def convert_module(module: nn.Module, name: str = ""):
            for child_name, child in list(module.named_children()):
                full_name = f"{name}.{child_name}" if name else child_name

                if isinstance(child, nn.Linear):
                    # Create FP8 layer
                    fp8_layer = NativeFP8Linear(
                        child.in_features,
                        child.out_features,
                        bias=child.bias is not None,
                        weight_format=self.weight_format,
                        activation_format=self.activation_format,
                        device=device
                    )

                    # Copy weights
                    fp8_layer.weight_master.data.copy_(child.weight.data)
                    if child.bias is not None:
                        fp8_layer.bias.data.copy_(child.bias.data)

                    # Quantize weights
                    fp8_layer._quantize_weights()

                    # Replace in model
                    setattr(module, child_name, fp8_layer)
                    self._fp8_layers[full_name] = fp8_layer

                else:
                    convert_module(child, full_name)

        convert_module(self.model)

    def _calibrate_scales(self):
        """Calibrate activation scales using calibration data"""
        if self.calibration_data is None:
            return

        self.model.eval()
        with torch.no_grad():
            # Run calibration data through model
            _ = self.model(self.calibration_data)

        # Scales are updated via forward hooks in FP8 layers

    @torch.no_grad()
    def infer(self, inputs: torch.Tensor) -> torch.Tensor:
        """Run FP8 inference"""
        if not self._prepared:
            raise RuntimeError("Engine not prepared. Call prepare() first.")

        return self.model(inputs)

    def get_memory_savings(self) -> dict[str, Any]:
        """Estimate memory savings from FP8"""
        fp8_bytes = 0
        fp32_bytes = 0

        for _name, layer in self._fp8_layers.items():
            # FP8: 1 byte per element
            fp8_bytes += layer.in_features * layer.out_features * 1
            # FP32: 4 bytes per element
            fp32_bytes += layer.in_features * layer.out_features * 4

        savings = (fp32_bytes - fp8_bytes) / fp32_bytes if fp32_bytes > 0 else 0

        return {
            'fp8_memory_mb': fp8_bytes / (1024 * 1024),
            'fp32_memory_mb': fp32_bytes / (1024 * 1024),
            'savings_ratio': savings,
            'savings_percent': savings * 100,
            'fp8_layers_count': len(self._fp8_layers)
        }

    def get_layer_info(self) -> dict[str, dict[str, Any]]:
        """Get info for all FP8 layers"""
        return {name: layer.get_fp8_info() for name, layer in self._fp8_layers.items()}


# Utility functions

def is_fp8_available() -> bool:
    """Check if native FP8 is available"""
    return FP8_DTYPES_AVAILABLE


def get_fp8_info() -> dict[str, Any]:
    """Get FP8 support information"""
    return {
        'pytorch_version': torch.__version__,
        'fp8_native_available': FP8_DTYPES_AVAILABLE,
        'fp8_scaled_mm_available': FP8_SCALED_MM_AVAILABLE,
        'supported_formats': ['e4m3fn', 'e5m2'] if FP8_DTYPES_AVAILABLE else [],
        'e4m3_max_value': 448.0,
        'e5m2_max_value': 57344.0,
        'recommended_use': {
            'e4m3fn': 'Forward pass (weights, activations) - higher precision',
            'e5m2': 'Backward pass (gradients) - wider dynamic range'
        }
    }


def convert_model_to_native_fp8(
    model: nn.Module,
    weight_format: FP8Dtype = FP8Dtype.E4M3,
    activation_format: FP8Dtype = FP8Dtype.E4M3,
    inplace: bool = False
) -> nn.Module:
    """
    Convert model to use native FP8 layers.

    Args:
        model: Model to convert
        weight_format: FP8 format for weights
        activation_format: FP8 format for activations
        inplace: Whether to modify model in place

    Returns:
        Model with FP8 layers
    """
    import copy

    if not inplace:
        model = copy.deepcopy(model)

    device = next(model.parameters()).device
    converted_count = 0

    def convert_module(module: nn.Module, name: str = ""):
        nonlocal converted_count

        for child_name, child in list(module.named_children()):
            full_name = f"{name}.{child_name}" if name else child_name

            if isinstance(child, nn.Linear):
                fp8_layer = NativeFP8Linear(
                    child.in_features,
                    child.out_features,
                    bias=child.bias is not None,
                    weight_format=weight_format,
                    activation_format=activation_format,
                    device=device
                )

                fp8_layer.weight_master.data.copy_(child.weight.data)
                if child.bias is not None:
                    fp8_layer.bias.data.copy_(child.bias.data)

                fp8_layer._quantize_weights()

                setattr(module, child_name, fp8_layer)
                converted_count += 1
            else:
                convert_module(child, full_name)

    convert_module(model)

    if converted_count > 0:
        print(f"Converted {converted_count} Linear layers to NativeFP8Linear")

    return model


def benchmark_fp8_layer(
    in_features: int = 1024,
    out_features: int = 1024,
    batch_size: int = 32,
    num_iterations: int = 100,
    device: torch.device | None = None
) -> dict[str, Any]:
    """
    Benchmark FP8 vs standard linear layer performance.

    Args:
        in_features: Input feature size
        out_features: Output feature size
        batch_size: Batch size
        num_iterations: Number of iterations
        device: Target device

    Returns:
        Dictionary with timing results
    """
    import time

    device = device or torch.device('cuda' if torch.cuda.is_available() else 'cpu')

    # Create layers
    fp8_layer = NativeFP8Linear(in_features, out_features, device=device)
    std_layer = nn.Linear(in_features, out_features, device=device)

    # Test input
    x = torch.randn(batch_size, in_features, device=device)

    # Warmup
    for _ in range(10):
        _ = fp8_layer(x)
        _ = std_layer(x)

    if device.type == 'cuda':
        torch.cuda.synchronize()

    # Benchmark FP8
    start = time.perf_counter()
    for _ in range(num_iterations):
        _ = fp8_layer(x)
    if device.type == 'cuda':
        torch.cuda.synchronize()
    fp8_time = (time.perf_counter() - start) / num_iterations * 1000

    # Benchmark standard
    start = time.perf_counter()
    for _ in range(num_iterations):
        _ = std_layer(x)
    if device.type == 'cuda':
        torch.cuda.synchronize()
    std_time = (time.perf_counter() - start) / num_iterations * 1000

    speedup = std_time / fp8_time if fp8_time > 0 else 1.0

    return {
        'fp8_time_ms': fp8_time,
        'standard_time_ms': std_time,
        'speedup': speedup,
        'device': str(device),
        'batch_size': batch_size,
        'in_features': in_features,
        'out_features': out_features,
        'fp8_native': FP8_DTYPES_AVAILABLE
    }
