from typing import Any

import torch
import torch.nn as nn
from fouroversix.matmul import fp4_matmul
from fouroversix.model.config import ModuleQuantizationConfig
from fouroversix.model.quantize import QuantizedModule
from fouroversix.quantize import (
    QuantizationConfig,
    QuantizedTensor,
    quantize_to_fp4,
)
from fouroversix.utils import RoundStyle


class FourOverSixLinearFunction(torch.autograd.Function):
    """Differentiable FP4 linear layer."""

    @staticmethod
    def forward(
        ctx: torch.autograd.function.FunctionCtx,
        config: ModuleQuantizationConfig,
        input: torch.Tensor,
        weight: torch.Tensor | QuantizedTensor,
        bias: torch.Tensor = None,
    ) -> tuple[torch.Tensor,]:
        """
        Perform an FP4 matrix multiplication. The input is provided in high precision
        and quantized to FP4 prior to the matrix multiplication, while the weight is
        provided in low precision.
        """

        fprop_activation_config = config.get_activation_config()
        fprop_weight_config = config.get_weight_config(
            block_scale_2d=config.weight_scale_2d,
        )

        if isinstance(weight, torch.Tensor):
            ctx.save_for_backward(input, weight, bias)
            weight = quantize_to_fp4(weight, fprop_weight_config)

        ctx.config = config

        out = fp4_matmul(
            input.reshape(-1, input.shape[-1]),
            weight,
            backend=config.matmul_backend,
            input_config=fprop_activation_config,
            out_dtype=config.output_dtype,
        ).reshape(*input.shape[:-1], weight.original_shape[0])

        if bias is not None:
            out = out + bias

        return out

    @staticmethod
    def backward(
        ctx: torch.autograd.function.FunctionCtx,
        grad_output: torch.Tensor,
    ) -> tuple[torch.Tensor, ...]:
        """Backward pass for the FP4 linear layer."""

        input, weight, bias = ctx.saved_tensors

        assert grad_output.shape[0] == 1  # noqa: S101

        dgrad_grad_config = ctx.config.get_gradient_config(
            round_style=RoundStyle.stochastic,
        )

        dgrad_weight_config = ctx.config.get_weight_config(
            block_scale_2d=ctx.config.weight_scale_2d,
            transpose=True,
        )

        grad_input = fp4_matmul(
            grad_output[0],
            weight,
            backend=ctx.matmul_backend,
            input_config=dgrad_grad_config,
            other_config=dgrad_weight_config,
            out_dtype=ctx.config.output_dtype,
        ).unsqueeze(0)

        wgrad_grad_config = ctx.config.get_gradient_config(
            rht=True,
            round_style=RoundStyle.stochastic,
            transpose=True,
        )

        wgrad_activation_config = ctx.config.get_activation_config(
            rht=True,
            transpose=True,
        )

        grad_weight = fp4_matmul(
            grad_output[0],
            input[0],
            backend=ctx.matmul_backend,
            input_config=wgrad_grad_config,
            other_config=wgrad_activation_config,
            out_dtype=ctx.config.output_dtype,
        ).unsqueeze(0)

        grad_bias = (
            grad_output.sum(0) if bias is not None and ctx.needs_input_grad[3] else None
        )

        return (
            None,
            grad_input,
            grad_weight,
            grad_bias,
        )


@QuantizedModule.register(nn.Linear)
class FourOverSixLinear(nn.Linear):
    """
    Drop-in replacement for `nn.Linear` that quantizes weights, activations, and
    gradients.
    """

    def __init__(
        self,
        module: nn.Linear,
        config: ModuleQuantizationConfig,
    ) -> None:
        """
        Initialize the FourOverSixLinear layer.

        Args:
            module (nn.Linear): The high-precision module that this quantized layer will
                replace.
            config (ModuleQuantizationConfig): The quantization configuration to use for
                the layer.

        """

        super().__init__(
            module.in_features,
            module.out_features,
            module.bias is not None,
            module.weight.device,
            module.weight.dtype,
        )

        self.weight = module.weight
        self.bias = module.bias
        self.config = config

        if not self.config.keep_master_weights:
            self.register_buffer(
                "quantized_weight_values",
                nn.Parameter(
                    torch.zeros(
                        self.out_features,
                        self.in_features // 2,
                        dtype=torch.uint8,
                    ),
                    requires_grad=False,
                ),
            )
            self.register_buffer(
                "quantized_weight_scale_factors",
                nn.Parameter(
                    torch.zeros(
                        self.out_features
                        * self.in_features
                        // self.config.dtype.block_size(),
                        dtype=self.config.dtype.scale_dtype(),
                    ),
                    requires_grad=False,
                ),
            )
            self.register_buffer(
                "quantized_weight_amax",
                nn.Parameter(torch.zeros(1, dtype=torch.float32), requires_grad=False),
            )
            self.register_buffer(
                "quantized_weight_metadata",
                nn.Parameter(
                    torch.zeros(2 + 2, dtype=torch.int32),
                    requires_grad=False,
                ),
            )

    @property
    def parameters_to_quantize(self) -> tuple[str, ...]:
        """Return high precision parameters to be quantized and deleted."""
        return ("weight",)

    def get_packing_factor(self, parameter_name: str) -> float:
        """Get the packing factor for a parameter."""
        return 2 if parameter_name == "quantized_weight_values" else 1

    def get_quantized_parameters(
        self,
        parameter_name: str,
        parameter: torch.Tensor,
    ) -> dict[str, Any]:
        """Get the quantized parameters for the layer."""

        if parameter_name == "weight":
            config = QuantizationConfig(
                backend=self.config.quantize_backend,
                block_scale_2d=self.config.weight_scale_2d,
                dtype=self.config.dtype,
                scale_rule=self.config.get_weight_scale_rule(),
            )

            quantized_weight = quantize_to_fp4(parameter, config)

            return {
                "quantized_weight_values": quantized_weight.values,
                "quantized_weight_scale_factors": quantized_weight.scale_factors,
                "quantized_weight_amax": quantized_weight.amax,
                "quantized_weight_metadata": torch.tensor(
                    [
                        quantized_weight.original_shape[0],
                        quantized_weight.original_shape[1],
                        quantized_weight.padded_shape[0],
                        quantized_weight.padded_shape[1],
                    ],
                ),
            }

        msg = f"Unsupported high-preciison parameter: {parameter_name}"
        raise ValueError(msg)

    def quantized_weight(self) -> QuantizedTensor:
        """
        Prepare this layer for post-training quantization by quantizing the weight,
        storing the quantized weight, and deleting the original weight. This should not
        be done if the layer is used for training, as training requires storage of the
        high-precision weight.
        """

        if not hasattr(self, "_quantized_weight"):
            if self.config.keep_master_weights:
                return quantize_to_fp4(self.weight, self.config.get_weight_config())

            original_shape = tuple(self.quantized_weight_metadata.data[:2].tolist())
            padded_shape = tuple(self.quantized_weight_metadata.data[2:].tolist())

            self._quantized_weight = QuantizedTensor(
                self.quantized_weight_values.data,
                self.quantized_weight_scale_factors.data,
                self.quantized_weight_amax.data,
                self.config.dtype,
                original_shape,
                self.config.get_weight_scale_rule(),
                padded_shape,
            )

        return self._quantized_weight

    def forward(self, input: torch.Tensor) -> torch.Tensor:
        """Forward pass for the FP4 linear layer."""

        return FourOverSixLinearFunction.apply(
            self.config,
            input,
            self.quantized_weight(),
            self.bias,
        )
