"""
INT8 Absmax Weight Quantization.

Maps float weights to 8-bit integers using symmetric scaling:
    scale = max(|W|) / 127
    W_int8 = round(W / scale)
    W_reconstructed = W_int8 * scale

Per-tensor mode uses one scale for the entire weight matrix.
Per-channel mode uses one scale per output channel (row), which
preserves more precision since each row is scaled independently —
a single outlier row no longer degrades the entire tensor.

INT8 halves memory vs float16/bfloat16. The tradeoff is a small
accuracy loss from rounding, measurable as a perplexity increase.
"""

import torch
import torch.nn as nn


def absmax_quantize(tensor, per_channel=False):
    """
    Quantize a float tensor to INT8 using absmax scaling.

    Args:
        tensor: Float tensor (typically a weight matrix, shape: out_features x in_features)
        per_channel: If True, compute one scale per row (output channel)

    Returns:
        (quantized_int8_tensor, scale)
    """
    if per_channel:
        scale = tensor.abs().amax(dim=1, keepdim=True) / 127
    else:
        scale = tensor.abs().max() / 127

    scale = torch.clamp(scale, min=1e-8)

    quantized = torch.clamp(torch.round(tensor / scale), -127, 127).to(torch.int8)
    return quantized, scale


def absmax_dequantize(quantized, scale):
    """Reconstruct float tensor from INT8 + scale."""
    return quantized.float() * scale


class QuantizedLinear(nn.Module):
    """Drop-in replacement for nn.Linear that stores INT8 weights and dequantizes on forward."""

    def __init__(self, in_features, out_features, weight_int8, scale, bias=None):
        super().__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.register_buffer("weight_int8", weight_int8)
        self.register_buffer("scale", scale)
        if bias is not None:
            self.register_buffer("bias", bias)
        else:
            self.bias = None

    @classmethod
    def from_linear(cls, linear, per_channel=False):
        """Create a QuantizedLinear from an existing nn.Linear."""
        weight_int8, scale = absmax_quantize(linear.weight.data.float(), per_channel=per_channel)
        bias = linear.bias.data if linear.bias is not None else None
        return cls(linear.in_features, linear.out_features, weight_int8, scale, bias)

    def forward(self, x):
        w = absmax_dequantize(self.weight_int8, self.scale).to(x.dtype)
        return nn.functional.linear(x, w, self.bias)


def quantize_model(model, per_channel=False):
    """Walk the model and replace all nn.Linear layers with QuantizedLinear.

    Skips Linear layers whose weights are tied to an Embedding (e.g. a
    shared output head), since quantizing them would break the tie and
    produce a mismatch between embedding lookups and output projections.
    """
    tied_ptrs = {
        m.weight.data_ptr()
        for m in model.modules()
        if isinstance(m, nn.Embedding)
    }
    _quantize_recursive(model, per_channel, tied_ptrs)


def _quantize_recursive(module, per_channel, tied_ptrs=frozenset()):
    for name, child in module.named_children():
        if isinstance(child, nn.Linear):
            if child.weight.data_ptr() in tied_ptrs:
                continue  # skip weight-tied layers
            setattr(module, name, QuantizedLinear.from_linear(child, per_channel=per_channel))
        else:
            _quantize_recursive(child, per_channel, tied_ptrs)