"""
Xavante - smoothquant_compressor.py
Responsabilidade: Compressao W8A8 via SmoothQuant (Lema 4.1).
Integra a versao validada do FlexNet.
"""
from __future__ import annotations

import logging

import torch
import torch.nn as nn

logger = logging.getLogger(__name__)


class SmoothQuantCompressor(nn.Module):
    """
    SmoothQuant (Xiao et al., 2022) migra a escala da ativacao para o peso,
    permitindo quantizacao W8A8 sem perda significativa.

    W_smooth = W * diag(s)
    X_smooth = X * diag(s)^-1
    s_j = max|X_j|^alpha / max|W_j|^(1-alpha),  alpha=0.5 tipicamente.
    """

    def __init__(self, alpha: float = 0.5, n_bits: int = 8, calibration_samples: int = 128):
        super().__init__()
        self.alpha = alpha
        self.n_bits = n_bits
        self.calibration_samples = calibration_samples
        self.register_buffer("smooth_scale", torch.ones(1))
        self._calibrated = False

    def calibrate(self, weight: torch.Tensor, activation_samples: torch.Tensor) -> None:
        """
        weight: [out, in]
        activation_samples: [N, in] (diversas amostras reais)
        """
        with torch.no_grad():
            # max abs por canal de entrada
            w_max = weight.abs().amax(dim=0).clamp(min=1e-5)  # [in]
            x_max = activation_samples.abs().amax(dim=0).clamp(min=1e-5)  # [in]
            s = (x_max.pow(self.alpha) / w_max.pow(1 - self.alpha)).clamp(min=1e-5)
            self.smooth_scale = s.to(weight.device, weight.dtype)
            self._calibrated = True
            logger.info("SmoothQuant calibrado: scale mean=%.4f, std=%.4f", s.mean().item(), s.std().item())

    def smooth_weight(self, weight: torch.Tensor) -> torch.Tensor:
        """Aplica W_smooth = W * diag(s) (multiplicacao por coluna)."""
        if not self._calibrated:
            return weight
        return weight * self.smooth_scale.unsqueeze(0)

    def smooth_activation(self, x: torch.Tensor) -> torch.Tensor:
        """Aplica X_smooth = X / diag(s)."""
        if not self._calibrated:
            return x
        return x / self.smooth_scale

    def quantize_per_tensor_symmetric(self, t: torch.Tensor) -> torch.Tensor:
        # Quantizacao simetrica per-tensor em n_bits
        qmax = 2 ** (self.n_bits - 1) - 1
        scale = t.abs().amax().clamp(min=1e-8) / qmax
        q = torch.clamp((t / scale).round(), -qmax, qmax)
        return q.to(torch.int8) if self.n_bits == 8 else q.to(torch.int16)

    def fake_quantize(self, t: torch.Tensor) -> torch.Tensor:
        """Quantizacao fake (diferenciavel via STE) para treino QAT."""
        qmax = 2 ** (self.n_bits - 1) - 1
        scale = t.abs().amax().clamp(min=1e-8) / qmax
        q = torch.clamp((t / scale).round(), -qmax, qmax)
        return q * scale  # retorna em float para backward

    def forward(self, weight: torch.Tensor, x: torch.Tensor, fake: bool = True) -> torch.Tensor:
        w_s = self.smooth_weight(weight)
        x_s = self.smooth_activation(x)
        if fake:
            w_q = self.fake_quantize(w_s)
            x_q = self.fake_quantize(x_s)
        else:
            w_q = self.dequantize(self.quantize_per_tensor_symmetric(w_s), w_s)
            x_q = self.dequantize(self.quantize_per_tensor_symmetric(x_s), x_s)
        return torch.nn.functional.linear(x_q, w_q)

    def dequantize(self, qt: torch.Tensor, ref: torch.Tensor) -> torch.Tensor:
        qmax = 2 ** (self.n_bits - 1) - 1
        scale = ref.abs().amax().clamp(min=1e-8) / qmax
        return qt.to(ref.dtype) * scale


__all__ = ["SmoothQuantCompressor"]