"""TurboQuant: Vector quantization via random rotation + optimal scalar quantization.

Implements the TurboQuant algorithm from "TurboQuant: Online Vector Quantization"
(arxiv 2504.19874, ICLR 2026) for KV cache compression in LLMs.

The key insight: randomly rotating high-dimensional vectors makes their coordinates
nearly independent with a known distribution (Beta → approx Gaussian). This allows
optimal scalar quantization per coordinate, preserving inner products needed by
attention's Q·K^T computation.

Two variants:
  - TurboQuantMSE: MSE-optimal, (b)-bit scalar quantization after rotation
  - TurboQuantProd: Inner-product optimal, (b-1)-bit MSE + 1-bit QJL residual
"""

import math

import mlx.core as mx
import numpy as np


# ============================================================================
# Precomputed optimal centroids for the Beta distribution
# ============================================================================

def _compute_gaussian_centroids(bits: int) -> list[float]:
    """Compute optimal Lloyd-Max centroids for N(0,1) distribution.

    For high-dimensional vectors after random rotation, each coordinate
    is approximately N(0, 1/d). We quantize the normalized version (N(0,1))
    and scale by 1/sqrt(d) at runtime.

    Uses the well-known optimal centroids for Gaussian distributions.
    """
    from scipy.stats import norm

    n_levels = 2 ** bits

    if bits == 1:
        # Optimal for N(0,1): +/- sqrt(2/pi) ≈ +/- 0.7979
        c = math.sqrt(2.0 / math.pi)
        return [-c, c]

    # Lloyd-Max iteration for N(0,1)
    # Start with uniform quantile initialization
    quantiles = np.linspace(0, 1, n_levels + 1)[1:-1]
    boundaries = norm.ppf(quantiles)
    boundaries = np.concatenate([[-np.inf], boundaries, [np.inf]])

    for _ in range(200):  # Lloyd-Max iterations
        # Compute centroids as conditional expectations
        centroids = []
        for i in range(n_levels):
            lo, hi = boundaries[i], boundaries[i + 1]
            # E[X | lo < X < hi] for N(0,1)
            p = norm.cdf(hi) - norm.cdf(lo)
            if p < 1e-15:
                centroids.append((lo + hi) / 2 if np.isfinite(lo) and np.isfinite(hi)
                                 else lo if np.isfinite(lo) else hi)
            else:
                centroids.append((norm.pdf(lo) - norm.pdf(hi)) / p)
        centroids = np.array(centroids)

        # Update boundaries as midpoints between centroids
        new_boundaries = np.concatenate([
            [-np.inf],
            (centroids[:-1] + centroids[1:]) / 2,
            [np.inf],
        ])

        if np.max(np.abs(new_boundaries[1:-1] - boundaries[1:-1])) < 1e-10:
            break
        boundaries = new_boundaries

    return sorted(centroids.tolist())


# Precompute and cache centroids for common bit-widths
_CENTROID_CACHE: dict[int, list[float]] = {}


def get_centroids(bits: int) -> np.ndarray:
    """Get precomputed optimal centroids for N(0,1) at given bit-width."""
    if bits not in _CENTROID_CACHE:
        _CENTROID_CACHE[bits] = _compute_gaussian_centroids(bits)
    return np.array(_CENTROID_CACHE[bits], dtype=np.float32)


# ============================================================================
# Random matrix generation
# ============================================================================

def generate_rotation_matrix(d: int, seed: int = 42) -> mx.array:
    """Generate a d×d random orthogonal matrix via QR decomposition.

    The rotation is data-oblivious — same matrix used for all vectors.
    Memory cost: d² × 4 bytes (e.g., 128² × 4 = 64 KB for head_dim=128).
    """
    key = mx.random.key(seed)
    G = mx.random.normal(shape=(d, d), key=key)
    # QR decomposition must run on CPU in MLX
    Q, R = mx.linalg.qr(G, stream=mx.cpu)
    mx.eval(Q, R)
    # Ensure proper rotation (det = +1) by adjusting signs
    diag_signs = mx.sign(mx.diag(R))
    Q = Q * diag_signs[None, :]
    mx.eval(Q)
    return Q


def generate_qjl_matrix(d: int, seed: int = 137) -> mx.array:
    """Generate a d×d random Gaussian matrix for QJL transform.

    Used in TurboQuantProd for 1-bit residual quantization.
    Stored as float16 to save memory (~32 KB for d=128).
    """
    key = mx.random.key(seed)
    S = mx.random.normal(shape=(d, d), key=key)
    return S.astype(mx.float16)


# ============================================================================
# TurboQuantMSE: MSE-optimal vector quantizer
# ============================================================================

class TurboQuantMSE:
    """MSE-optimal vector quantizer via random rotation + scalar quantization.

    Algorithm:
      1. Normalize input vectors (store norms separately)
      2. Rotate: y = x_unit @ R^T  (coordinates become ~N(0, 1/d))
      3. Scale to N(0,1): z = y * sqrt(d)
      4. Quantize each coordinate: find nearest centroid
      5. Dequantize: look up centroids, scale back, rotate back, restore norms
    """

    def __init__(self, d: int, bits: int, seed: int = 42):
        self.d = d
        self.bits = bits
        self.rotation = generate_rotation_matrix(d, seed)
        centroids_np = get_centroids(bits)
        self.centroids = mx.array(centroids_np)  # (2^bits,)
        self.n_centroids = len(centroids_np)
        # Precompute boundaries for fast quantization
        boundaries_np = (centroids_np[:-1] + centroids_np[1:]) / 2
        self.boundaries = mx.array(boundaries_np)

    def quantize(self, x: mx.array) -> tuple[mx.array, mx.array]:
        """Quantize vectors.

        Args:
            x: (..., d) float vectors

        Returns:
            indices: (..., d) uint8 centroid indices
            norms: (..., 1) vector norms
        """
        # Store norms
        norms = mx.linalg.norm(x, axis=-1, keepdims=True)
        x_unit = x / mx.maximum(norms, mx.array(1e-8))

        # Rotate
        y = x_unit @ self.rotation.T  # (..., d)

        # Scale to N(0,1) for centroid matching
        z = y * math.sqrt(self.d)

        # Find nearest centroid using boundaries (searchsorted)
        # For small n_centroids this is efficient
        indices = mx.zeros(z.shape, dtype=mx.uint8)
        for i, b in enumerate(self.boundaries.tolist()):
            indices = mx.where(z > b, mx.array(i + 1, dtype=mx.uint8), indices)

        return indices, norms

    def dequantize(self, indices: mx.array, norms: mx.array) -> mx.array:
        """Dequantize from indices back to vectors.

        Args:
            indices: (..., d) uint8 centroid indices
            norms: (..., 1) vector norms

        Returns:
            x_hat: (..., d) reconstructed vectors
        """
        # Look up centroids
        z_hat = self.centroids[indices.astype(mx.int32)]

        # Scale back from N(0,1) to N(0, 1/d)
        y_hat = z_hat / math.sqrt(self.d)

        # Inverse rotation
        x_hat = y_hat @ self.rotation

        # Restore norms
        return x_hat * norms


# ============================================================================
# TurboQuantProd: Inner-product optimal quantizer
# ============================================================================

class TurboQuantProd:
    """Inner-product optimal quantizer: MSE(b-1 bits) + QJL(1 bit) on residual.

    The QJL (Quantized Johnson-Lindenstrauss) stage provides an unbiased
    inner product estimator, which is crucial for attention's Q·K^T computation.

    Total budget: b bits per element (b-1 for MSE + 1 for QJL signs).
    """

    def __init__(self, d: int, bits: int, seed: int = 42):
        if bits < 2:
            raise ValueError("TurboQuantProd requires at least 2 bits (1 MSE + 1 QJL)")
        self.d = d
        self.bits = bits
        self.mse = TurboQuantMSE(d, bits - 1, seed)
        self.qjl = generate_qjl_matrix(d, seed + 1000)

    def quantize(self, x: mx.array) -> tuple[mx.array, mx.array, mx.array, mx.array]:
        """Quantize vectors with MSE + QJL residual correction.

        Args:
            x: (..., d) float vectors

        Returns:
            mse_indices: (..., d) uint8
            qjl_signs: (..., d) int8 in {-1, +1}
            residual_norms: (..., 1)
            norms: (..., 1)
        """
        norms = mx.linalg.norm(x, axis=-1, keepdims=True)
        x_unit = x / mx.maximum(norms, mx.array(1e-8))

        # MSE stage
        mse_indices, _ = self.mse.quantize(x_unit)
        x_mse = self.mse.dequantize(mse_indices, mx.ones_like(norms))

        # Residual
        residual = x_unit - x_mse
        residual_norms = mx.linalg.norm(residual, axis=-1, keepdims=True)

        # QJL: sign(S · residual)
        qjl_signs = mx.sign(residual @ self.qjl.astype(x.dtype).T)
        # Ensure no zeros (map 0 to +1)
        qjl_signs = mx.where(qjl_signs == 0, mx.array(1.0), qjl_signs)

        return mse_indices, qjl_signs.astype(mx.int8), residual_norms, norms

    def dequantize(
        self,
        mse_indices: mx.array,
        qjl_signs: mx.array,
        residual_norms: mx.array,
        norms: mx.array,
    ) -> mx.array:
        """Dequantize with MSE + QJL correction.

        Args:
            mse_indices: (..., d) uint8
            qjl_signs: (..., d) int8
            residual_norms: (..., 1)
            norms: (..., 1)

        Returns:
            x_hat: (..., d) reconstructed vectors
        """
        # MSE reconstruction (on unit sphere)
        x_mse = self.mse.dequantize(mse_indices, mx.ones_like(norms))

        # QJL correction: (sqrt(pi/2) / d) * ||r|| * S^T @ signs
        scale = math.sqrt(math.pi / 2) / self.d
        qjl_correction = (
            scale
            * residual_norms
            * (qjl_signs.astype(mx.float16) @ self.qjl)
        )

        return (x_mse + qjl_correction) * norms


# ============================================================================
# Convenience functions
# ============================================================================

def create_quantizer(
    head_dim: int,
    bits: int,
    use_qjl: bool = True,
    seed: int = 42,
) -> TurboQuantMSE | TurboQuantProd:
    """Create the appropriate TurboQuant quantizer.

    Args:
        head_dim: Dimension of each attention head
        bits: Total bits per element
        use_qjl: Whether to use QJL residual correction (recommended for KV cache)
        seed: Random seed for rotation matrix

    Returns:
        TurboQuantMSE or TurboQuantProd instance
    """
    if use_qjl and bits >= 2:
        return TurboQuantProd(head_dim, bits, seed)
    else:
        return TurboQuantMSE(head_dim, bits, seed)
