from functools import partial
from typing import Any, Literal, Optional

import torch
import torch.nn as nn
import torch.nn.utils.parametrize as P

from .. import functional as F


class Bnb4bitParametrization(nn.Module):
    """
    A parametrization module that handles dequantization of a 4-bit quantized parameter.

    The parameter data is expected to be already quantized when this parametrization is applied.
    This module will dequantize the parameter data to its original floating-point representation
    when the forward method is called (i.e. when the parameter is accessed).

    Args:
        quant_state (`F.QuantState`):
            The quantization state containing the necessary information for dequantization.
    """

    def __init__(self, quant_state: F.QuantState):
        super().__init__()
        self.quant_state = quant_state

    @torch.no_grad()
    def forward(self, quantized_param: torch.Tensor) -> torch.Tensor:
        """
        Forward pass to dequantize the parameter.

        Args:
            quantized_param (`torch.Tensor`): The quantized parameter tensor (from .original)

        Returns:
            `torch.Tensor`: The dequantized parameter tensor in the original shape and dtype.
        """
        return F.dequantize_4bit(quantized_param, self.quant_state)


def replace_parameter_4bit_prequantized(
    module: nn.Module, param_name: str, qs_dict: dict[str, Any], device: torch.device
):
    if not hasattr(module, param_name):
        raise AttributeError(f"Module does not have parameter '{param_name}'")

    original_param = getattr(module, param_name)

    if not isinstance(original_param, nn.Parameter):
        raise TypeError(f"Parameter '{param_name}' is not an instance of nn.Parameter")

    quant_state = F.QuantState.from_dict(qs_dict, device=device)

    # Apply a parametrization to the module to handle dequantization.
    P.register_parametrization(module, param_name, Bnb4bitParametrization(quant_state), unsafe=True)

    # Next, register hooks.
    _register_parametrization_hooks(module, param_name)


def replace_parameter_4bit(
    module: nn.Module,
    param_name: str,
    compress_statistics: bool = False,
    quant_type: Literal["nf4", "fp4"] = "nf4",
    blocksize: Optional[int] = None,
):
    """
    Replace a module parameter with a 4-bit quantized version using parametrization.

    This function quantizes an existing parameter in a PyTorch module to 4-bit precision
    and sets up parametrization to handle automatic dequantization during forward passes.
    The original parameter is replaced with quantized data, and a parametrization layer
    is registered to manage the quantization state and dequantization process.

    Additional, it registers a state dict post-hook to ensure that the quantization state
    is saved correctly when the model's state dict is saved.

    It is useful for MoE models or other scenarios where you want to quantize parameters
    outside of nn.Linear layers without changing the model's architecture.

    <Tip warning={true}>This feature is experimental and may change in future releases.</Tip>

    Args:
        module (`nn.Module`):
            The PyTorch module containing the parameter to be quantized.
        param_name (`str`):
            The name of the parameter within the module to quantize.
        compress_statistics (`bool`, *optional*, defaults to `False`):
            Whether to compress quantization statistics to reduce memory usage.
        quant_type (`Literal["nf4", "fp4"]`, *optional*, defaults to `"nf4"`):
            The quantization format to use.
        blocksize (`int`, *optional*, defaults to `None`):
            The block size for quantization. If None, uses the default block size.

    Raises:
        AttributeError: If the module does not have the specified parameter.
        TypeError: If the specified attribute is not an instance of nn.Parameter.
    """

    if not hasattr(module, param_name):
        raise AttributeError(f"Module does not have parameter '{param_name}'")

    original_param = getattr(module, param_name)

    if not isinstance(original_param, nn.Parameter):
        raise TypeError(f"Parameter '{param_name}' is not an instance of nn.Parameter")

    # Quantize the original parameter.
    quantized_data, quant_state = F.quantize_4bit(
        original_param.data,
        blocksize=blocksize,
        compress_statistics=compress_statistics,
        quant_type=quant_type,
    )

    # Replace the parameter with the quantized data.
    setattr(module, param_name, nn.Parameter(quantized_data, requires_grad=False))
    del original_param

    # Apply a parametrization to the module to handle dequantization.
    P.register_parametrization(module, param_name, Bnb4bitParametrization(quant_state), unsafe=True)

    # Next, register hooks.
    _register_parametrization_hooks(module, param_name)


def _disable_parametrization_cache(module: nn.Module, inputs: tuple[Any, ...], output: Any):
    # Clamp instead of a bare decrement: with ``always_call=True`` this hook also runs
    # when the forward raised before the pre-hook incremented (e.g. an earlier pre-hook
    # failed), and the counter must never go negative — a negative value is truthy, so
    # ``if not P._cache_enabled`` would stop clearing the cache forever.
    P._cache_enabled = max(0, P._cache_enabled - 1)
    if not P._cache_enabled:
        P._cache = {}


def _enable_parametrization_cache(module: nn.Module, inputs: tuple[Any, ...]):
    P._cache_enabled += 1


def _register_parametrization_hooks(module: nn.Module, param_name: str):
    # Register a state dict hook for saving. Note that this requires torch >= 2.5.0.
    if torch.__version__ >= (2, 5):
        module.register_state_dict_post_hook(
            partial(
                _parametrized_state_dict_post_hook,
                param_name=param_name,
            )
        )

    # Register hooks to enable caching for the dequantization parametrization.
    # This helps preserve time and memory when the same quantized parameter
    # is accessed multiple times in the forward computation.
    #
    # ``always_call=True`` is load-bearing: activation checkpointing with
    # ``use_reentrant=False`` aborts its backward recompute mid-forward by design
    # (early stop, via an internal exception) once the last needed activation has
    # been rematerialized. A plain forward hook is skipped in that case, so the
    # global ``parametrize._cache_enabled`` counter leaks upward once per
    # checkpointed region per step, after which the cache is enabled (and never
    # cleared) for the remainder of training — every dequantized parameter this
    # module produces stays resident, i.e. a memory leak of the full dequantized
    # model size (4x the packed 4-bit bytes).
    module.register_forward_pre_hook(_enable_parametrization_cache)
    module.register_forward_hook(_disable_parametrization_cache, always_call=True)


def _parametrized_state_dict_post_hook(
    module: nn.Module,
    state_dict: dict[str, Any],
    prefix: str,
    local_metadata: Any,
    *,
    param_name: str = "weight",
    **kwargs: dict[str, Any],
) -> None:
    """
    Hook to modify the state dict to include the quantization state.
    """

    original_key = f"{prefix}parametrizations.{param_name}.original"

    if original_key in state_dict:
        # Create a clean entry.
        # The `parametrizations.{param_name}.original` key will have the quantized data,
        # but we would like it to keep it in the state_dict as `{param_name}`.
        clean_key = f"{prefix}{param_name}"
        state_dict[clean_key] = state_dict.pop(original_key)

        assert P.is_parametrized(module, param_name)

        # Find the parametrization, which should have the quantization state.
        parametrization: Bnb4bitParametrization = next(
            filter(lambda x: isinstance(x, Bnb4bitParametrization), module.parametrizations[param_name]), None
        )

        assert parametrization is not None, "Parametrization not found for the parameter."

        quant_state = parametrization.quant_state

        # Next, we need to store the quantization state.
        if quant_state is not None:
            for k, v in quant_state.as_dict(packed=True).items():
                state_dict[f"{prefix}{param_name}.{k}"] = v