import torch
import torch.nn as nn

from torch.ao.quantization.fake_quantize import FakeQuantize
from torch.ao.quantization.observer import MinMaxObserver
from torch.ao.quantization.qconfig import QConfig
from torch.ao.quantization import convert

from .model import BitTransformerLM


def quantize_dynamic(model: BitTransformerLM, dtype: torch.dtype = torch.qint8) -> BitTransformerLM:
    """Return a dynamically quantized copy of the model for inference."""
    quantized = torch.quantization.quantize_dynamic(
        model, {nn.Linear}, dtype=dtype
    )
    return quantized


class FourBitObserver(MinMaxObserver):
    """Min-max observer configured for 4-bit quantization."""

    def __init__(self, **kwargs):
        super().__init__(
            quant_min=0,
            quant_max=15,
            dtype=torch.quint8,
            qscheme=torch.per_tensor_affine,
            **kwargs,
        )


FourBitFakeQuantize = FakeQuantize.with_args(observer=FourBitObserver)

four_bit_qconfig = QConfig(activation=FourBitFakeQuantize, weight=FourBitFakeQuantize)


class QATLinear(nn.Linear):
    """Linear layer with fake quantization for QAT."""

    def __init__(self, in_features: int, out_features: int, bias: bool = True) -> None:
        super().__init__(in_features, out_features, bias)
        self.weight_fake_quant = FourBitFakeQuantize()
        self.activation_post_process = FourBitFakeQuantize()

    @classmethod
    def from_float(cls, mod: nn.Linear) -> "QATLinear":
        qat = cls(mod.in_features, mod.out_features, mod.bias is not None)
        qat.weight = mod.weight
        qat.bias = mod.bias
        return qat

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.activation_post_process(x)
        w = self.weight_fake_quant(self.weight)
        return nn.functional.linear(x, w, self.bias)


def prepare_qat_fx(model: BitTransformerLM) -> BitTransformerLM:
    """Prepare BitTransformerLM for quantization-aware training."""

    for name, module in model.named_children():
        if isinstance(module, nn.Linear):
            setattr(model, name, QATLinear.from_float(module))
        else:
            prepare_qat_fx(module)
    return model


def convert_qat_fx(model: BitTransformerLM) -> BitTransformerLM:
    """Convert a QAT-prepared model to a quantized version."""

    for name, module in model.named_children():
        if isinstance(module, QATLinear):
            w = module.weight.data
            qmin, qmax = 0, 15
            min_w = w.min()
            max_w = w.max()
            scale = (max_w - min_w) / (qmax - qmin + 1e-8)
            zero_point = qmin - torch.round(min_w / scale)
            q_w = torch.clamp(torch.round(w / scale + zero_point), qmin, qmax)
            new_mod = nn.Linear(module.in_features, module.out_features, module.bias is not None)
            new_mod.weight = nn.Parameter((q_w - zero_point) * scale)
            if module.bias is not None:
                new_mod.bias = nn.Parameter(module.bias.data)
            setattr(model, name, new_mod)
        else:
            convert_qat_fx(module)
    return model