"""TurboQuant KV Cache for MLX.

Drop-in replacement for mlx-lm's QuantizedKVCache that uses TurboQuant
(random rotation + optimal scalar quantization) instead of affine quantization.

Key design:
  - Stores KV vectors as compressed centroid indices + norms
  - Returns dequantized plain arrays from update_and_fetch()
  - Does NOT expose a `bits` attribute → mlx-lm uses standard SDPA path
  - Compatible with mlx-lm's generate pipeline without modification

Usage:
    from optiq.core.turbo_kv_cache import TurboQuantKVCache
    cache = [TurboQuantKVCache(head_dim=128, bits=4) for _ in range(n_layers)]
    # Pass to model as usual — works with mlx-lm's generate_step
"""

import mlx.core as mx

from .turbo_quant import TurboQuantMSE, TurboQuantProd, create_quantizer


class TurboQuantKVCache:
    """KV cache using TurboQuant for inner-product preserving quantization.

    Instead of mlx-lm's affine quantization (per-group scale/bias that distorts
    inner products), this uses random rotation + optimal scalar quantization
    which provably preserves the inner products needed by attention's Q·K^T.

    The cache stores compressed representations but returns dequantized float16
    tensors from update_and_fetch(), so attention uses the standard SDPA path.
    Memory savings come from the compressed storage allowing longer sequences.
    """

    step = 256

    def __init__(
        self,
        head_dim: int = 128,
        bits: int = 4,
        use_qjl: bool = True,
        seed: int = 42,
    ):
        self.head_dim = head_dim
        self._bits = bits
        self.use_qjl = use_qjl and bits >= 2
        self.offset = 0
        self.seed = seed

        # Create quantizers (one for keys, one for values — same config)
        self.k_quantizer = create_quantizer(head_dim, bits, self.use_qjl, seed)
        self.v_quantizer = create_quantizer(head_dim, bits, self.use_qjl, seed + 500)

        # Storage for quantized data
        self._k_store = None  # Will hold quantized key data
        self._v_store = None  # Will hold quantized value data

    # Note: we intentionally do NOT have a `bits` property.
    # This ensures mlx-lm's scaled_dot_product_attention uses the standard
    # SDPA path (mx.fast.scaled_dot_product_attention) instead of the
    # quantized path (which expects affine format).

    def _quantize_vectors(self, x: mx.array, quantizer) -> dict:
        """Quantize a batch of vectors.

        Args:
            x: (B, n_heads, num_steps, head_dim)
            quantizer: TurboQuantMSE or TurboQuantProd

        Returns:
            dict with quantized components
        """
        if isinstance(quantizer, TurboQuantProd):
            mse_idx, qjl_signs, res_norms, norms = quantizer.quantize(x)
            return {
                "mse_indices": mse_idx,
                "qjl_signs": qjl_signs,
                "residual_norms": res_norms,
                "norms": norms,
            }
        else:
            indices, norms = quantizer.quantize(x)
            return {
                "indices": indices,
                "norms": norms,
            }

    def _dequantize_vectors(self, store: dict, quantizer) -> mx.array:
        """Dequantize stored vectors back to float.

        Returns: (B, n_heads, seq_len, head_dim) float16
        """
        if isinstance(quantizer, TurboQuantProd):
            return quantizer.dequantize(
                store["mse_indices"],
                store["qjl_signs"],
                store["residual_norms"],
                store["norms"],
            )
        else:
            return quantizer.dequantize(store["indices"], store["norms"])

    def _init_store(self, quantizer, B, n_heads, n_steps, head_dim, dtype):
        """Initialize empty storage for quantized vectors."""
        if isinstance(quantizer, TurboQuantProd):
            return {
                "mse_indices": mx.zeros((B, n_heads, n_steps, head_dim), dtype=mx.uint8),
                "qjl_signs": mx.zeros((B, n_heads, n_steps, head_dim), dtype=mx.int8),
                "residual_norms": mx.zeros((B, n_heads, n_steps, 1), dtype=dtype),
                "norms": mx.zeros((B, n_heads, n_steps, 1), dtype=dtype),
            }
        else:
            return {
                "indices": mx.zeros((B, n_heads, n_steps, head_dim), dtype=mx.uint8),
                "norms": mx.zeros((B, n_heads, n_steps, 1), dtype=dtype),
            }

    def _expand_store(self, store: dict, n_new_steps: int):
        """Expand storage by appending zeros."""
        expanded = {}
        for k, v in store.items():
            shape = list(v.shape)
            shape[2] = n_new_steps
            new = mx.zeros(shape, dtype=v.dtype)
            expanded[k] = mx.concatenate([v, new], axis=2)
        return expanded

    def _trim_store(self, store: dict, end: int):
        """Trim storage to given sequence position."""
        return {k: v[..., :end, :] for k, v in store.items()}

    def _write_store(self, store: dict, data: dict, start: int, end: int):
        """Write quantized data into storage at given range."""
        for k in store:
            store[k][..., start:end, :] = data[k]

    def update_and_fetch(self, keys, values):
        """Quantize new KV entries and return all cached entries (dequantized).

        This is the critical method called by attention layers.

        Args:
            keys: (B, n_kv_heads, num_steps, head_dim) new keys
            values: (B, n_kv_heads, num_steps, head_dim) new values

        Returns:
            (dequantized_keys, dequantized_values) covering full sequence
        """
        B, n_heads, num_steps, head_dim = keys.shape
        v_head_dim = values.shape[-1]
        prev = self.offset

        # Grow storage if needed
        if self._k_store is None or (prev + num_steps) > list(self._k_store.values())[0].shape[2]:
            new_steps = (self.step + num_steps - 1) // self.step * self.step

            if self._k_store is not None:
                # Trim excess, then expand
                if prev % self.step != 0:
                    self._k_store = self._trim_store(self._k_store, prev)
                    self._v_store = self._trim_store(self._v_store, prev)
                self._k_store = self._expand_store(self._k_store, new_steps)
                self._v_store = self._expand_store(self._v_store, new_steps)
            else:
                self._k_store = self._init_store(
                    self.k_quantizer, B, n_heads, new_steps, head_dim, keys.dtype
                )
                self._v_store = self._init_store(
                    self.v_quantizer, B, n_heads, new_steps, v_head_dim, values.dtype
                )

        self.offset += num_steps

        # Quantize new entries
        k_data = self._quantize_vectors(keys, self.k_quantizer)
        v_data = self._quantize_vectors(values, self.v_quantizer)

        # Write into storage
        self._write_store(self._k_store, k_data, prev, self.offset)
        self._write_store(self._v_store, v_data, prev, self.offset)

        # Dequantize full sequence and return plain arrays
        k_slice = {k: v[..., :self.offset, :] for k, v in self._k_store.items()}
        v_slice = {k: v[..., :self.offset, :] for k, v in self._v_store.items()}

        keys_out = self._dequantize_vectors(k_slice, self.k_quantizer)
        values_out = self._dequantize_vectors(v_slice, self.v_quantizer)

        return keys_out, values_out

    @property
    def state(self):
        """Return serializable state (for prompt cache saving)."""
        if self._k_store is None:
            return [], []

        k = self._trim_store(self._k_store, self.offset)
        v = self._trim_store(self._v_store, self.offset)
        return (
            list(k.values()),
            list(v.values()),
        )

    @state.setter
    def state(self, v):
        pass  # Not implemented for TurboQuant caches

    @property
    def meta_state(self):
        return tuple(map(str, (self.offset, self.head_dim, self._bits)))

    @meta_state.setter
    def meta_state(self, v):
        self.offset, self.head_dim, self._bits = map(int, v)

    def is_trimmable(self):
        return True

    def trim(self, n):
        n = min(self.offset, n)
        self.offset -= n
        return n

    def make_mask(self, N, return_array=True, window_size=None):
        """Create attention mask compatible with mlx-lm."""
        from mlx_lm.models.base import create_causal_mask

        if window_size is not None:
            return create_causal_mask(N, self.offset, window_size=window_size)
        elif N == 1:
            return None
        elif return_array:
            return create_causal_mask(N, self.offset, window_size=window_size)
        else:
            return "causal"

    def empty(self):
        return self._k_store is None

    def size(self):
        return self.offset


def make_turbo_kv_caches(
    n_layers: int,
    head_dim: int = 128,
    bits: int = 4,
    use_qjl: bool = True,
    seed: int = 42,
    per_layer_bits: dict[int, int] | None = None,
) -> list[TurboQuantKVCache]:
    """Create a list of TurboQuant KV caches, one per layer.

    Args:
        n_layers: Number of transformer layers
        head_dim: Dimension per attention head
        bits: Default bits per element
        use_qjl: Use QJL residual correction
        seed: Random seed
        per_layer_bits: Optional dict mapping layer_idx → bits for mixed precision

    Returns:
        List of TurboQuantKVCache instances
    """
    caches = []
    for i in range(n_layers):
        layer_bits = per_layer_bits.get(i, bits) if per_layer_bits else bits
        caches.append(TurboQuantKVCache(
            head_dim=head_dim,
            bits=layer_bits,
            use_qjl=use_qjl,
            seed=seed + i,  # Different rotation per layer for diversity
        ))
    return caches
