"""TurboQuant State Cache for GatedDeltaNet (Qwen3.5 linear attention).

Qwen3.5 uses a hybrid architecture:
  - 18 GatedDeltaNet layers with recurrent state (ArraysCache)
  - 6 standard self-attention layers with KV cache (KVCache)

The GatedDeltaNet state is a matrix of shape (B, n_heads, head_k_dim, head_v_dim)
= (B, 16, 128, 128) per layer. This is a fixed-size cache (doesn't grow with
sequence length), storing a compressed representation of all past tokens.

TurboQuant can quantize this state by treating each row (or slice) as a
128-dimensional vector and applying rotation + optimal scalar quantization.

IMPORTANT: Error accumulation warning
  The recurrent state is read-modify-written on every autoregressive step.
  Each quantize->dequantize round-trip adds small errors that compound over
  tokens. This makes delta state quantization:
    - Lossless for single-pass evaluation (perplexity) where the state is
      written once but never re-read within the same forward pass
    - Lossy for autoregressive generation (needle-in-haystack, chat) where
      errors accumulate over every generated token
  For generation tasks, keep the delta state at FP16 and only quantize
  the self-attention KV cache (which is write-once-read-many).

Design:
  - QuantizedArraysCache wraps ArraysCache
  - Intercepts state writes (cache[1] = state) to quantize
  - Intercepts state reads (state = cache[1]) to dequantize
  - Conv state (cache[0]) is left unquantized (tiny: kernel_size × conv_dim)
"""

from typing import Any, List, Optional

import mlx.core as mx

from .turbo_quant import TurboQuantMSE


class QuantizedArraysCache:
    """Drop-in replacement for ArraysCache that quantizes the recurrent state.

    The GatedDeltaNet cache has two slots:
      cache[0] = conv_state: (B, kernel-1, conv_dim) — small, left unquantized
      cache[1] = recurrent_state: (B, n_heads, k_dim, v_dim) — quantized

    We quantize along the last dimension (v_dim), treating each (v_dim,) slice
    as a vector for TurboQuant. This preserves the structure used by
    gated_delta_update which computes state = decay * state + beta * outer(k, v).
    """

    def __init__(self, size: int, bits: int = 4, v_dim: int = 128, seed: int = 42,
                 left_padding: Optional[List[int]] = None):
        self.cache = [None] * size
        self._bits = bits
        self._v_dim = v_dim
        self._quantizer = TurboQuantMSE(v_dim, bits, seed)
        self._quantized_state = None  # Stores (indices, norms) for slot 1
        self.left_padding = mx.array(left_padding) if left_padding else None
        self.lengths = None

    def __setitem__(self, idx, value):
        if idx == 1 and value is not None:
            # Quantize the recurrent state
            # value shape: (B, n_heads, k_dim, v_dim)
            indices, norms = self._quantizer.quantize(value)
            mx.eval(indices, norms)
            self._quantized_state = (indices, norms)
            # Store a sentinel so empty() checks work
            self.cache[idx] = True
        else:
            self.cache[idx] = value

    def __getitem__(self, idx):
        if idx == 1 and self._quantized_state is not None:
            # Dequantize on read
            indices, norms = self._quantized_state
            return self._quantizer.dequantize(indices, norms)
        return self.cache[idx]

    @property
    def state(self):
        # For serialization — return raw cache
        result = list(self.cache)
        if self._quantized_state is not None:
            indices, norms = self._quantized_state
            result[1] = self._quantizer.dequantize(indices, norms)
        return result

    @state.setter
    def state(self, v):
        self.cache = list(v)
        if len(v) > 1 and v[1] is not None:
            self[1] = v[1]  # Triggers quantization

    @property
    def meta_state(self):
        return ""

    @meta_state.setter
    def meta_state(self, v):
        pass

    def is_trimmable(self):
        return False

    def empty(self):
        return all(c is None for c in self.cache)

    def size(self):
        return 0

    def filter(self, batch_indices):
        self.cache = [
            c[batch_indices] if c is not None and c is not True else c
            for c in self.cache
        ]
        if self._quantized_state is not None:
            indices, norms = self._quantized_state
            self._quantized_state = (indices[batch_indices], norms[batch_indices])

    def extend(self, other):
        for i in range(len(self.cache)):
            if self.cache[i] is not None and self.cache[i] is not True:
                if isinstance(other, QuantizedArraysCache):
                    other_val = other.cache[i]
                else:
                    other_val = other[i] if hasattr(other, '__getitem__') else other.cache[i]
                if other_val is not None and other_val is not True:
                    self.cache[i] = mx.concatenate([self.cache[i], other_val])

    def extract(self, idx):
        new_cache = QuantizedArraysCache(
            len(self.cache), self._bits, self._v_dim
        )
        new_cache.cache = [
            c[idx:idx+1] if c is not None and c is not True else c
            for c in self.cache
        ]
        if self._quantized_state is not None:
            indices, norms = self._quantized_state
            new_cache._quantized_state = (indices[idx:idx+1], norms[idx:idx+1])
            new_cache.cache[1] = True
        return new_cache

    def prepare(self, lengths=None, **kwargs):
        self.lengths = mx.array(lengths) if lengths is not None else None

    def finalize(self):
        self.lengths = None
        self.left_padding = None

    def advance(self, N):
        if self.lengths is not None:
            self.lengths -= N
        if self.left_padding is not None:
            self.left_padding -= N

    def make_mask(self, N: int):
        if self.left_padding is not None:
            pos = mx.arange(N)
            return pos >= self.left_padding[:, None]
        elif self.lengths is not None:
            pos = mx.arange(N)
            return pos < self.lengths[:, None]
        return None

    @classmethod
    def from_state(cls, state, meta_state):
        obj = cls.__new__(cls)
        obj.cache = list(state)
        obj._quantized_state = None
        obj.left_padding = None
        obj.lengths = None
        return obj


def make_turbo_cache_for_qwen35(
    model,
    bits: int = 4,
    seed: int = 42,
    quantize_kv: bool = True,
) -> list:
    """Create a mixed cache for Qwen3.5 with TurboQuant quantization.

    - GatedDeltaNet layers: QuantizedArraysCache (quantized recurrent state)
    - Self-attention layers: TurboQuantKVCache or standard KVCache

    Args:
        model: The loaded Qwen3.5 model
        bits: Bit-width for quantization
        seed: Random seed
        quantize_kv: Whether to also quantize self-attention KV cache

    Returns:
        List of cache objects (one per layer)
    """
    from mlx_lm.models.cache import KVCache

    default_cache = model.make_cache()
    caches = []

    for i, (layer, default_c) in enumerate(zip(model.layers, default_cache)):
        cache_type = type(default_c).__name__

        if cache_type == "ArraysCache":
            # GatedDeltaNet layer — quantize recurrent state
            la = layer.linear_attn
            v_dim = la.head_v_dim
            size = len(default_c.cache) if hasattr(default_c, 'cache') else 2
            cache = QuantizedArraysCache(
                size=size,
                bits=bits,
                v_dim=v_dim,
                seed=seed + i,
            )
            caches.append(cache)

        elif cache_type == "KVCache" and quantize_kv:
            # Self-attention layer — quantize KV cache
            from .turbo_kv_cache import TurboQuantKVCache
            sa = layer.self_attn
            head_dim = sa.head_dim
            cache = TurboQuantKVCache(
                head_dim=head_dim,
                bits=bits,
                use_qjl=bits >= 2,
                seed=seed + i + 100,
            )
            caches.append(cache)

        else:
            # Keep default cache (KVCache without quantization, or unknown type)
            caches.append(default_c)

    n_quant_delta = sum(1 for c in caches if isinstance(c, QuantizedArraysCache))
    n_quant_kv = sum(
        1 for c in caches
        if hasattr(c, '_bits') and not isinstance(c, QuantizedArraysCache)
    )
    n_default = len(caches) - n_quant_delta - n_quant_kv

    print(f"  TurboQuant cache: {n_quant_delta} GatedDeltaNet layers (state), "
          f"{n_quant_kv} self-attn layers (KV), {n_default} default")

    return caches
