"""Quantization utilities for model compression.

Supports INT8 dynamic quantization (PyTorch native) and INT4 quantization
(via bitsandbytes when available).
"""

import torch
import torch.nn as nn
from typing import Optional, Set, Type, Union
import warnings


def quantize_dynamic(
    model: nn.Module,
    dtype: torch.dtype = torch.qint8,
    layers_to_quantize: Optional[Set[Type[nn.Module]]] = None,
    inplace: bool = False,
) -> nn.Module:
    """Apply dynamic quantization to a model for faster CPU inference.
    
    Dynamic quantization quantizes weights at load time and activations
    dynamically during inference. Best for models with large linear layers.
    
    .. warning::
        Dynamic quantization in PyTorch only supports **CPU**. This function
        automatically moves the model to CPU. Using the quantized model in a
        GPU pipeline will cause severe performance degradation due to device
        transfers. For GPU inference, use :func:`quantize_4bit` with bitsandbytes
        or ``torch.compile()`` instead.
    
    Args:
        model: The model to quantize.
        dtype: Quantization dtype (torch.qint8 or torch.float16). Defaults to qint8.
        layers_to_quantize: Set of layer types to quantize. Defaults to {nn.Linear}.
        inplace: If True, modifies model in place. Defaults to False.
        
    Returns:
        Quantized model (on CPU).
        
    Example:
        >>> model = GPT.load("model.pt")
        >>> model_q = quantize_dynamic(model)
        >>> # ~2-4x smaller, faster on CPU
    """
    if layers_to_quantize is None:
        layers_to_quantize = {nn.Linear}
    
    # Warn if model is on GPU since dynamic quantization only supports CPU
    if any(p.is_cuda for p in model.parameters()):
        warnings.warn(
            "Model is on GPU but dynamic quantization only supports CPU. "
            "Model will be moved to CPU. For GPU inference, use quantize_4bit() "
            "with bitsandbytes or torch.compile() instead.",
            UserWarning
        )
    
    if not inplace:
        import copy
        model = copy.deepcopy(model)
    
    # Move to CPU for quantization (required by PyTorch quantization backend)
    model = model.cpu()
    
    quantized = torch.quantization.quantize_dynamic(
        model,
        layers_to_quantize,
        dtype=dtype
    )
    
    return quantized


def quantize_4bit(
    model: nn.Module,
    compute_dtype: torch.dtype = torch.float16,
    quant_type: str = "nf4",
) -> nn.Module:
    """Apply 4-bit quantization using bitsandbytes (if available).
    
    4-bit quantization provides ~4x memory reduction with minimal quality loss.
    Requires bitsandbytes library.
    
    Args:
        model: The model to quantize.
        compute_dtype: Dtype for computation (float16 or bfloat16).
        quant_type: Quantization type ("nf4" or "fp4"). NF4 is recommended.
        
    Returns:
        Quantized model.
        
    Example:
        >>> model = GPT.load("model.pt")
        >>> model_4bit = quantize_4bit(model)
        >>> # ~4x smaller, runs on GPU
    """
    try:
        import bitsandbytes as bnb
    except ImportError:
        raise ImportError(
            "4-bit quantization requires bitsandbytes. "
            "Install with: pip install bitsandbytes"
        )
    
    def _replace_linear_4bit(module: nn.Module):
        """Recursively replace Linear layers with 4-bit versions."""
        for name, child in module.named_children():
            if isinstance(child, nn.Linear):
                # Create 4-bit linear
                new_layer = bnb.nn.Linear4bit(
                    child.in_features,
                    child.out_features,
                    bias=child.bias is not None,
                    compute_dtype=compute_dtype,
                    quant_type=quant_type,
                )
                # Copy weights (will be quantized)
                new_layer.weight.data = child.weight.data
                if child.bias is not None:
                    new_layer.bias.data = child.bias.data
                
                setattr(module, name, new_layer)
            else:
                _replace_linear_4bit(child)
    
    import copy
    model = copy.deepcopy(model)
    _replace_linear_4bit(model)
    
    return model


def estimate_memory_savings(
    model: nn.Module,
    quantization: str = "int8"
) -> dict:
    """Estimate memory savings from quantization.
    
    Args:
        model: The model to analyze.
        quantization: Type of quantization ("int8", "int4", "fp16").
        
    Returns:
        Dict with memory estimates.
        
    Example:
        >>> savings = estimate_memory_savings(model, "int8")
        >>> print(f"Original: {savings['original_mb']:.1f}MB")
        >>> print(f"Quantized: {savings['quantized_mb']:.1f}MB")
        >>> print(f"Reduction: {savings['reduction_ratio']:.1f}x")
    """
    # Calculate original size
    param_bytes = sum(p.numel() * p.element_size() for p in model.parameters())
    buffer_bytes = sum(b.numel() * b.element_size() for b in model.buffers())
    original_bytes = param_bytes + buffer_bytes
    
    # Estimate quantized size based on type
    reduction_factors = {
        "int8": 4,   # float32 -> int8 = 4x reduction
        "int4": 8,   # float32 -> int4 = 8x reduction
        "fp16": 2,   # float32 -> fp16 = 2x reduction
    }
    
    factor = reduction_factors.get(quantization, 1)
    quantized_bytes = original_bytes / factor
    
    return {
        "original_mb": original_bytes / (1024 * 1024),
        "quantized_mb": quantized_bytes / (1024 * 1024),
        "reduction_ratio": factor,
        "original_bytes": original_bytes,
        "quantized_bytes": quantized_bytes,
    }


def _detect_quantization_method(model: nn.Module) -> dict:
    """Detect the quantization method used on a model.
    
    Args:
        model: The model to analyze.
        
    Returns:
        Dict with quantization info (method, compute_dtype, quant_type).
    """
    quant_info = {
        "method": None,
        "compute_dtype": None,
        "quant_type": None,
    }
    
    # Check for bitsandbytes 4-bit quantization
    try:
        import bitsandbytes as bnb
        for module in model.modules():
            if isinstance(module, bnb.nn.Linear4bit):
                quant_info["method"] = "4bit"
                quant_info["compute_dtype"] = str(module.compute_dtype)
                quant_info["quant_type"] = getattr(module, "quant_type", "nf4")
                return quant_info
    except ImportError:
        pass
    
    # Check for PyTorch dynamic quantization
    for module in model.modules():
        if hasattr(module, "_packed_params"):  # Quantized Linear layers
            quant_info["method"] = "dynamic"
            return quant_info
    
    return quant_info


def save_quantized(
    model: nn.Module,
    path: str,
    include_config: bool = True,
    quantization_method: Optional[str] = None,
    compute_dtype: Optional[str] = None,
    quant_type: Optional[str] = None,
):
    """Save a quantized model with metadata for proper reloading.
    
    Args:
        model: Quantized model to save.
        path: Save path.
        include_config: Include model config if available.
        quantization_method: Override auto-detection ("dynamic", "4bit", or None).
        compute_dtype: Compute dtype for 4-bit quantization.
        quant_type: Quantization type for 4-bit ("nf4" or "fp4").
    """
    # Auto-detect quantization method if not provided
    if quantization_method is None:
        quant_info = _detect_quantization_method(model)
    else:
        quant_info = {
            "method": quantization_method,
            "compute_dtype": compute_dtype,
            "quant_type": quant_type,
        }
    
    save_dict = {
        "model_state_dict": model.state_dict(),
        "quantization_info": quant_info,
        "quantized": True,  # For legacy compatibility and tests
    }
    
    if include_config and hasattr(model, "get_config"):
        save_dict["config"] = model.get_config()
    
    torch.save(save_dict, path)
    print(f"Quantized model saved to {path} (method: {quant_info.get('method', 'none')})")


def load_quantized(
    model_class: Type[nn.Module],
    path: str,
    device: str = "cpu",
    **model_kwargs
) -> nn.Module:
    """Load a quantized model with proper structure reconstruction.
    
    This function applies the correct quantization structure to the model
    BEFORE loading the state dict, ensuring compatibility with 4-bit and
    dynamic quantization.
    
    Args:
        model_class: The model class to instantiate.
        path: Path to saved quantized model.
        device: Device to load on.
        **model_kwargs: Additional model arguments.
        
    Returns:
        Loaded quantized model.
        
    Raises:
        ValueError: If quantization method is unsupported.
    """
    checkpoint = torch.load(path, map_location=device, weights_only=False)
    
    # Get config from checkpoint or use provided kwargs
    if "config" in checkpoint:
        config = checkpoint["config"].copy()
        config.update(model_kwargs)
    else:
        config = model_kwargs
    
    # Create base model
    model = model_class(**config)
    
    # Get quantization info
    quant_info = checkpoint.get("quantization_info", {})
    method = quant_info.get("method")
    
    # Apply quantization structure BEFORE loading state dict
    if method == "4bit":
        compute_dtype_str = quant_info.get("compute_dtype", "torch.float16")
        # Parse dtype string to actual dtype
        if "float16" in compute_dtype_str:
            compute_dtype = torch.float16
        elif "bfloat16" in compute_dtype_str:
            compute_dtype = torch.bfloat16
        else:
            compute_dtype = torch.float32
        
        quant_type = quant_info.get("quant_type", "nf4")
        model = quantize_4bit(model, compute_dtype=compute_dtype, quant_type=quant_type)
        
    elif method == "dynamic":
        # Apply dynamic quantization structure to match saved state dict
        # We need to convert nn.Linear -> nn.quantized.dynamic.Linear
        # This mirrors the structure created by quantize_dynamic
        model = quantize_dynamic(model, inplace=True)
    
    # Handle legacy format (just "quantized": True flag)
    elif checkpoint.get("quantized") and method is None:
        warnings.warn(
            "Loading model saved with legacy quantization format. "
            "Attempting to load directly - this may fail for 4-bit models."
        )
    
    # Load state dict
    model.load_state_dict(checkpoint["model_state_dict"])
    model.to(device)
    model.eval()
    
    return model
