"""
HyperSparse Adaptive Quantization — sparsity adaptativa via kthvalue

HSAQ cuantiza ACTIVACIONES a {0, valor} mediante máscara binaria dinámica.
NO es cuantización de pesos. NO es INT8/INT4. NO es bitsandbytes.

Objetivo: optimizar la cuantización mejor que TurboQuant (Google) haciendo
que TODO sea adaptativo, para no usar recursos innecesariamente. Esto permite
usar modelos más grandes que la capacidad del hardware, ejecutando solo los
tokens/neuronas necesarias para el trabajo operativo.
"""
import torch
import torch.nn as nn


class HSAQ(nn.Module):
    """HyperSparse Adaptive Quantization.

    Aplica una máscara de sparse adaptativa por batch.
    Sparsity escalonada: capas tempranas preservan más información,
    capas tardías comprimen más agresivamente.

    Args:
        sparsity: Fracción de neuronas a enmascarar (0.3 = 30% → 0).
    """
    def __init__(self, sparsity: float = 0.3):
        super().__init__()
        self.sparsity = sparsity

    def forward(self, x: torch.Tensor, sparsity_override: float | None = None) -> torch.Tensor:
        """Aplica máscara de sparsity adaptativa por batch.

        Args:
            x: Tensor de entrada (B, D) o (B, T, D).
            sparsity_override: Sparsity para esta llamada (opcional).
                              Si se pasa, usa este valor en lugar de self.sparsity.
        """
        s = sparsity_override if sparsity_override is not None else self.sparsity
        if s <= 0.0:
            return x

        flat = x.abs().view(x.size(0), -1)
        n = flat.size(1)
        k = max(1, min(n - 1, int(n * s)))
        # kthvalue: encuentra el valor en el percentil k (bottom-k%)
        # Los valores >= thresh son el top-(1-s)% que pasan
        thresh = torch.kthvalue(flat, k, dim=1).values
        thresh = thresh.view(-1, *([1] * (x.dim() - 1)))
        mask = x.abs() >= thresh
        self._last_sparsity = 1.0 - mask.float().mean().item()
        self._last_threshold = thresh.view(-1).mean().item()
        self._last_sparsity_target = s
        return x * mask

    def get_stats(self) -> dict:
        """Retorna métricas actuales de HSAQ."""
        return {
            'sparsity': self.sparsity,
            'sparsity_target': getattr(self, '_last_sparsity_target', self.sparsity),
            'actual_sparsity': getattr(self, '_last_sparsity', 0.0),
            'threshold': getattr(self, '_last_threshold', 0.0),
            'type': 'activation_quantization',
            'method': 'kthvalue_dynamic_threshold',
        }