"""
Production FP8 Training Engine

High-performance FP8 training implementation with dynamic format selection,
automatic scaling, and production-grade reliability features.

Key Features:
- E4M3 for forward pass (higher precision)
- E5M2 for backward pass (broader dynamic range)
- Dynamic scaling strategies for numerical stability
- Transformer Engine integration
- 2x speedup on H100/Blackwell with maintained accuracy

References:
    - FP8 Formats Paper: https://arxiv.org/abs/2209.05433
    - Transformer Engine: https://docs.nvidia.com/deeplearning/transformer-engine/
    - NVIDIA FP8 Training: https://developer.nvidia.com/blog/nvidia-h100-transformer-engine/
"""

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

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

try:
    import transformer_engine.pytorch as te
    from transformer_engine.common.recipe import DelayedScaling, Format
    TRANSFORMER_ENGINE_AVAILABLE = True
except ImportError:
    TRANSFORMER_ENGINE_AVAILABLE = False
    # Note: FP8 training will use fallback implementations when Transformer Engine
    # is not installed. This is normal for most users - only warn when FP8 is used.

try:
    from ..hardware.abstraction.hal_core import HardwareAbstractionLayer
    HAL_AVAILABLE = True
except ImportError:
    HAL_AVAILABLE = False

logger = logging.getLogger(__name__)


class FP8Format(Enum):
    """Supported FP8 formats"""
    E4M3 = "e4m3"  # 1 sign, 4 exponent, 3 mantissa bits - higher precision
    E5M2 = "e5m2"  # 1 sign, 5 exponent, 2 mantissa bits - wider range


class FP8Config:
    """Configuration for FP8 training"""

    def __init__(
        self,
        forward_format: FP8Format = FP8Format.E4M3,
        backward_format: FP8Format = FP8Format.E5M2,
        scaling_strategy: str = "dynamic",
        scaling_interval: int = 1000,
        initial_scale: float = 2**15,
        growth_factor: float = 2.0,
        backoff_factor: float = 0.5,
        overflow_check_interval: int = 100,
        enable_amax_history: bool = True,
        amax_history_len: int = 1024,
        margin: int = 2,
        fp8_dpa: bool = True,  # Delayed Param Allocation
        fp8_mha: bool = True,  # Multi-Head Attention in FP8
        use_te_linear: bool = True  # Use Transformer Engine Linear layers
    ):
        self.forward_format = forward_format
        self.backward_format = backward_format
        self.scaling_strategy = scaling_strategy
        self.scaling_interval = scaling_interval
        self.initial_scale = initial_scale
        self.growth_factor = growth_factor
        self.backoff_factor = backoff_factor
        self.overflow_check_interval = overflow_check_interval
        self.enable_amax_history = enable_amax_history
        self.amax_history_len = amax_history_len
        self.margin = margin
        self.fp8_dpa = fp8_dpa
        self.fp8_mha = fp8_mha
        self.use_te_linear = use_te_linear

        # Validate configuration
        self._validate()

    def _validate(self):
        """Validate configuration parameters"""
        assert self.scaling_interval > 0, "Scaling interval must be positive"
        assert self.initial_scale > 0, "Initial scale must be positive"
        assert self.growth_factor > 1.0, "Growth factor must be > 1.0"
        assert 0.0 < self.backoff_factor < 1.0, "Backoff factor must be between 0 and 1"
        assert self.overflow_check_interval > 0, "Overflow check interval must be positive"
        assert self.amax_history_len > 0, "AMAX history length must be positive"
        assert self.margin >= 0, "Margin must be non-negative"

    def to_te_recipe(self):
        """Convert to Transformer Engine recipe format"""
        if not TRANSFORMER_ENGINE_AVAILABLE:
            return None

        # Map our formats to TE formats
        format_mapping = {
            FP8Format.E4M3: Format.E4M3,
            FP8Format.E5M2: Format.E5M2
        }

        fp8_format = format_mapping.get(self.forward_format, Format.E4M3)

        return DelayedScaling(
            fp8_format=fp8_format,
            amax_history_len=self.amax_history_len,
            amax_compute_algo="most_recent",
            scaling_factor_compute_algo="max",
            override_linear_precision=(False, False, not self.use_te_linear)
        )


class FP8ScaleManager:
    """Manages FP8 scaling factors for numerical stability"""

    def __init__(self, config: FP8Config, device: torch.device):
        self.config = config
        self.device = device

        # Initialize scaling factors
        self.scale = torch.tensor(config.initial_scale, device=device, dtype=torch.float32)
        self.inv_scale = torch.tensor(1.0 / config.initial_scale, device=device, dtype=torch.float32)

        # Overflow tracking
        self.overflow_count = 0
        self.step_count = 0
        self.amax_history = [] if config.enable_amax_history else None

        # Dynamic scaling state
        self.last_overflow_step = -1
        self.growth_tracker = 0

    def update_scale(self, amax: torch.Tensor, force_update: bool = False) -> bool:
        """
        Update scaling factor based on AMAX values

        Args:
            amax: Maximum absolute value from forward/backward pass
            force_update: Force scale update regardless of interval

        Returns:
            True if scaling was updated
        """
        self.step_count += 1

        # Update AMAX history
        if self.amax_history is not None:
            self.amax_history.append(float(amax.max()))
            if len(self.amax_history) > self.config.amax_history_len:
                self.amax_history.pop(0)

        # Check for overflow
        overflow_detected = self._check_overflow(amax)

        if overflow_detected:
            self._handle_overflow()
            return True

        # Dynamic scaling update
        if force_update or (self.step_count % self.config.scaling_interval == 0):
            return self._update_dynamic_scale(amax)

        return False

    def _check_overflow(self, amax: torch.Tensor) -> bool:
        """Check if overflow occurred based on AMAX values"""
        # For E4M3 format, max representable value is 448
        # For E5M2 format, max representable value is 57344
        max_values = {
            FP8Format.E4M3: 448.0,
            FP8Format.E5M2: 57344.0
        }

        max_fp8_value = max_values.get(self.config.forward_format, 448.0)
        scaled_amax = amax * self.scale

        # Check if any values exceed FP8 range
        overflow = scaled_amax > max_fp8_value
        return bool(overflow.any().item())

    def _handle_overflow(self):
        """Handle overflow by reducing scale"""
        self.overflow_count += 1
        self.last_overflow_step = self.step_count

        # Reduce scale
        self.scale = self.scale * self.config.backoff_factor
        self.inv_scale = 1.0 / self.scale

        # Reset growth tracker
        self.growth_tracker = 0

        logger.warning("FP8 overflow detected at step %d, reducing scale to %.4f", self.step_count, self.scale.item())

    def _update_dynamic_scale(self, amax: torch.Tensor) -> bool:
        """Update scale dynamically based on AMAX history"""
        if self.amax_history is None or len(self.amax_history) < 100:
            return False

        # Calculate recent AMAX statistics
        recent_amax = self.amax_history[-100:]
        max_amax = max(recent_amax)

        # Determine if we can safely increase scale
        max_values = {
            FP8Format.E4M3: 448.0,
            FP8Format.E5M2: 57344.0
        }

        max_fp8_value = max_values.get(self.config.forward_format, 448.0)
        safety_margin = 2 ** self.config.margin

        # Check if we have headroom to increase scale
        scaled_max_amax = max_amax * self.scale.item()
        headroom_ratio = max_fp8_value / (scaled_max_amax * safety_margin)

        # Only grow if we haven't had recent overflow and have sufficient headroom
        steps_since_overflow = self.step_count - self.last_overflow_step
        if steps_since_overflow > self.config.scaling_interval * 2 and headroom_ratio > 2.0:
            self.growth_tracker += 1

            # Grow scale if consistently safe
            if self.growth_tracker >= 5:
                self.scale = self.scale * self.config.growth_factor
                self.inv_scale = 1.0 / self.scale
                self.growth_tracker = 0
                logger.info("Growing FP8 scale to %.4f at step %d", self.scale.item(), self.step_count)
                return True

        return False

    def get_scale_info(self) -> dict[str, Any]:
        """Get current scaling information"""
        return {
            'scale': float(self.scale.item()),
            'inv_scale': float(self.inv_scale.item()),
            'overflow_count': self.overflow_count,
            'step_count': self.step_count,
            'recent_amax': self.amax_history[-10:] if self.amax_history else [],
            'growth_tracker': self.growth_tracker
        }


class FP8TrainingEngine:
    """
    Production-grade FP8 training engine with dynamic format selection

    Args:
        model: PyTorch model to train with FP8
        config: FP8Config for training configuration
        device: Target device for training

    Example:
        >>> config = FP8Config(forward_format=FP8Format.E4M3, backward_format=FP8Format.E5M2)  # doctest: +SKIP
        >>> engine = FP8TrainingEngine(model, config, device)  # doctest: +SKIP
        >>> engine.setup_fp8_training()  # doctest: +SKIP
    """

    def __init__(
        self,
        model: nn.Module,
        config: FP8Config,
        device: torch.device | None = None
    ):
        self.model = model
        self.config = config
        self.device = device or torch.device('cuda' if torch.cuda.is_available() else 'cpu')

        # Initialize components
        self.scale_manager = FP8ScaleManager(config, self.device)
        self.fp8_recipe = config.to_te_recipe() if TRANSFORMER_ENGINE_AVAILABLE else None

        # Training state
        self.is_setup = False
        self.fp8_enabled = False
        self.original_layers = {}  # Store original layers for restoration

        # Statistics tracking
        self.training_stats: dict[str, Any] = {
            'steps': 0,
            'overflows': 0,
            'scale_updates': 0,
            'avg_loss': 0.0
        }

        # Initialize HAL if available
        self.hal: HardwareAbstractionLayer | None
        if HAL_AVAILABLE:
            self.hal = HardwareAbstractionLayer()
        else:
            self.hal = None

    def setup_fp8_training(self) -> bool:
        """
        Initialize FP8 training setup

        Returns:
            True if setup was successful
        """
        if self.is_setup:
            warnings.warn("FP8 training already setup", stacklevel=2)
            return True

        logger.info("Setting up FP8 training...")

        # Check hardware compatibility
        if not self._check_hardware_compatibility():
            warnings.warn("Hardware may not support FP8 training optimally", stacklevel=2)

        # Replace linear layers with FP8 versions if Transformer Engine is available
        if TRANSFORMER_ENGINE_AVAILABLE and self.config.use_te_linear:
            self._replace_linear_layers()

        # Move model to device
        self.model = self.model.to(self.device)

        # Enable FP8 training mode
        self.fp8_enabled = True
        self.is_setup = True

        logger.info(
            "FP8 training setup complete — forward=%s, backward=%s, TE=%s, HAL=%s",
            self.config.forward_format.value,
            self.config.backward_format.value,
            TRANSFORMER_ENGINE_AVAILABLE,
            self.hal is not None,
        )

        return True

    def _check_hardware_compatibility(self) -> bool:
        """Check if current hardware supports FP8 training"""
        if not torch.cuda.is_available():
            return False

        # Check GPU generation - H100 (compute capability 9.0) and above support FP8
        device_capability = torch.cuda.get_device_capability(self.device)
        major, minor = device_capability

        if major >= 9:  # Hopper (H100) and newer
            return True
        elif major == 8 and minor >= 9:  # Some Ada Lovelace GPUs
            return True
        else:
            return False

    def _replace_linear_layers(self):
        """Replace standard Linear layers with Transformer Engine FP8 layers"""
        if not TRANSFORMER_ENGINE_AVAILABLE:
            return

        def replace_layer(module, name, layer):
            """Replace a single linear layer"""
            if isinstance(layer, nn.Linear):
                # Store original for potential restoration
                self.original_layers[name] = layer

                # Create FP8 replacement
                fp8_layer = te.Linear(
                    layer.in_features,
                    layer.out_features,
                    bias=layer.bias is not None,
                    device=layer.weight.device,
                    dtype=layer.weight.dtype
                )

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

                # Replace in model
                setattr(module, name, fp8_layer)
                return True
            return False

        # Recursively replace linear layers
        replaced_count = 0
        for _name, module in self.model.named_modules():
            for child_name, child in module.named_children():
                if replace_layer(module, child_name, child):
                    replaced_count += 1

        logger.info("Replaced %d Linear layers with FP8 versions", replaced_count)

    def training_step(
        self,
        inputs: torch.Tensor,
        targets: torch.Tensor,
        return_loss: bool = True
    ) -> torch.Tensor | tuple[torch.Tensor, dict[str, Any]]:
        """
        Single FP8 training step with automatic scaling

        Args:
            inputs: Input tensor
            targets: Target tensor
            return_loss: Whether to return loss or (loss, metadata)

        Returns:
            Loss tensor or tuple of (loss, metadata)
        """
        if not self.is_setup:
            raise RuntimeError("FP8 training not setup. Call setup_fp8_training() first.")

        self.training_stats['steps'] += 1

        # FP8 context manager for Transformer Engine
        if TRANSFORMER_ENGINE_AVAILABLE and self.fp8_recipe is not None:
            with te.fp8_autocast(enabled=self.fp8_enabled, fp8_recipe=self.fp8_recipe):
                outputs = self.model(inputs)
                # Handle sequence outputs: reshape [B, S, C] to [B*S, C] for cross_entropy
                if outputs.dim() == 3 and targets.dim() == 1:
                    outputs = outputs[:, -1, :]  # Take last token
                elif outputs.dim() == 3 and targets.dim() == 2:
                    outputs = outputs.reshape(-1, outputs.size(-1))
                    targets = targets.reshape(-1)
                loss = F.cross_entropy(outputs, targets)
        else:
            # Fallback to standard training with manual scaling
            with torch.autocast(device_type=str(self.device).split(':')[0], dtype=torch.bfloat16):
                outputs = self.model(inputs)
                # Handle sequence outputs: reshape [B, S, C] to [B*S, C] for cross_entropy
                if outputs.dim() == 3 and targets.dim() == 1:
                    outputs = outputs[:, -1, :]  # Take last token
                elif outputs.dim() == 3 and targets.dim() == 2:
                    outputs = outputs.reshape(-1, outputs.size(-1))
                    targets = targets.reshape(-1)
                loss = F.cross_entropy(outputs, targets)

                # Apply FP8 scaling manually
                if self.fp8_enabled:
                    loss = loss * self.scale_manager.scale

        # Update training statistics
        self.training_stats['avg_loss'] = (
            self.training_stats['avg_loss'] * 0.99 + float(loss.item()) * 0.01
        )

        if return_loss:
            return loss
        else:
            metadata = {
                'scale': self.scale_manager.get_scale_info(),
                'training_stats': self.training_stats.copy()
            }
            return loss, metadata

    def optimizer_step(self, optimizer: torch.optim.Optimizer, grad_scaler: Any | None = None) -> bool:
        """
        Optimizer step with FP8 scaling and overflow handling

        Args:
            optimizer: PyTorch optimizer
            grad_scaler: Optional gradient scaler for mixed precision

        Returns:
            True if optimizer step was successful (no overflow)
        """
        if not self.is_setup:
            raise RuntimeError("FP8 training not setup. Call setup_fp8_training() first.")

        # Check for gradient overflow
        amax = self._compute_grad_amax()
        overflow_detected = self.scale_manager.update_scale(amax)

        if overflow_detected:
            self.training_stats['overflows'] += 1
            # Skip optimizer step on overflow
            optimizer.zero_grad()
            return False

        # Unscale gradients if using FP8 scaling
        if self.fp8_enabled:
            self._unscale_gradients(optimizer)

        # Standard optimizer step
        if grad_scaler is not None:
            grad_scaler.step(optimizer)
            grad_scaler.update()
        else:
            optimizer.step()

        optimizer.zero_grad()
        return True

    def _compute_grad_amax(self) -> torch.Tensor:
        """Compute maximum absolute value of gradients"""
        amax_values = []

        for param in self.model.parameters():
            if param.grad is not None:
                amax_values.append(param.grad.abs().max())

        if amax_values:
            return torch.stack(amax_values).max()
        else:
            return torch.tensor(0.0, device=self.device)

    def _unscale_gradients(self, optimizer: torch.optim.Optimizer):
        """Unscale gradients for FP8 training"""
        inv_scale = self.scale_manager.inv_scale

        for group in optimizer.param_groups:
            for param in group['params']:
                if param.grad is not None:
                    param.grad.data.mul_(inv_scale)

    def get_training_statistics(self) -> dict[str, Any]:
        """Get comprehensive training statistics"""
        stats = self.training_stats.copy()
        stats.update({
            'scale_info': self.scale_manager.get_scale_info(),
            'fp8_enabled': self.fp8_enabled,
            'transformer_engine_available': TRANSFORMER_ENGINE_AVAILABLE,
            'hardware_compatible': self._check_hardware_compatibility(),
            'replaced_layers': len(self.original_layers)
        })
        return stats

    def save_checkpoint(self, filepath: str):
        """Save FP8 training checkpoint"""
        checkpoint = {
            'model_state_dict': self.model.state_dict(),
            'config': self.config.__dict__,
            'scale_manager_state': self.scale_manager.__dict__,
            'training_stats': self.training_stats,
            'fp8_enabled': self.fp8_enabled
        }
        torch.save(checkpoint, filepath)

    def load_checkpoint(self, filepath: str):
        """Load FP8 training checkpoint"""
        checkpoint = torch.load(filepath, map_location=self.device, weights_only=True)

        self.model.load_state_dict(checkpoint['model_state_dict'])
        self.training_stats = checkpoint.get('training_stats', self.training_stats)
        self.fp8_enabled = checkpoint.get('fp8_enabled', True)

        # Restore scale manager state
        if 'scale_manager_state' in checkpoint:
            for key, value in checkpoint['scale_manager_state'].items():
                if hasattr(self.scale_manager, key):
                    setattr(self.scale_manager, key, value)

    def restore_original_model(self):
        """Restore original model layers (for inference or non-FP8 training)"""
        if not self.original_layers:
            return

        for layer_path, original_layer in self.original_layers.items():
            # Navigate to parent module and restore layer
            parts = layer_path.split('.')
            module = self.model
            for part in parts[:-1]:
                module = getattr(module, part)
            setattr(module, parts[-1], original_layer)

        self.fp8_enabled = False
        logger.info("Restored %d original layers", len(self.original_layers))

    def __enter__(self):
        """Context manager entry"""
        if not self.is_setup:
            self.setup_fp8_training()
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        """Context manager exit.

        Note: FP8 context cleanup is minimal as scaling factors are
        managed internally. Future versions may add cache clearing
        or precision restoration for nested contexts.
        """
        # FP8 context exit - no cleanup currently required
        pass


# Factory function
def create_fp8_trainer(
    model: nn.Module,
    device: torch.device | None = None,
    forward_format: FP8Format = FP8Format.E4M3,
    backward_format: FP8Format = FP8Format.E5M2,
    **kwargs
) -> FP8TrainingEngine:
    """
    Factory function to create FP8 training engine with sensible defaults

    Args:
        model: PyTorch model to train
        device: Target device
        forward_format: FP8 format for forward pass
        backward_format: FP8 format for backward pass
        **kwargs: Additional config parameters

    Returns:
        Configured FP8TrainingEngine

    Example:
        >>> model = torch.nn.Linear(128, 64)  # doctest: +SKIP
        >>> trainer = create_fp8_trainer(model, device='cpu')  # doctest: +SKIP
    """
    config = FP8Config(
        forward_format=forward_format,
        backward_format=backward_format,
        **kwargs
    )

    return FP8TrainingEngine(model, config, device)


# Validation function
def validate_fp8_setup(
    model: nn.Module,
    device: torch.device | None = None
) -> dict[str, Any]:
    """
    Validate FP8 training setup and provide recommendations

    Args:
        model: PyTorch model to validate
        device: Target device

    Returns:
        Dictionary with validation results and recommendations
    """
    device = device or torch.device('cuda' if torch.cuda.is_available() else 'cpu')

    validation_results = {
        'valid': True,
        'warnings': [],
        'errors': [],
        'recommendations': [],
        'capabilities': {}
    }

    # Check hardware
    if device.type != 'cuda':
        validation_results['warnings'].append("FP8 training requires CUDA device")
        validation_results['valid'] = False

    if torch.cuda.is_available():
        capability = torch.cuda.get_device_capability(device)
        validation_results['capabilities']['compute_capability'] = capability

        if capability[0] < 9:
            validation_results['warnings'].append(
                f"Compute capability {capability} may not fully support FP8. "
                f"H100 (9.0) or newer recommended."
            )

    # Check Transformer Engine availability
    validation_results['capabilities']['transformer_engine'] = TRANSFORMER_ENGINE_AVAILABLE
    if not TRANSFORMER_ENGINE_AVAILABLE:
        validation_results['recommendations'].append(
            "Install Transformer Engine for optimal FP8 performance: pip install transformer-engine"
        )

    # Count linear layers
    linear_count = sum(1 for module in model.modules() if isinstance(module, nn.Linear))
    validation_results['capabilities']['linear_layers'] = linear_count

    if linear_count == 0:
        validation_results['warnings'].append("No Linear layers found - FP8 benefits may be limited")

    # Check model size
    param_count = sum(p.numel() for p in model.parameters())
    validation_results['capabilities']['parameter_count'] = param_count

    if param_count < 100_000_000:  # 100M parameters
        validation_results['recommendations'].append(
            "FP8 training is most beneficial for large models (100M+ parameters)"
        )

    return validation_results
