"""Mixed-precision KV cache quantization for MLX.

The KV cache stores key/value projections for all past tokens during
autoregressive generation. For long contexts, this dominates memory:

  KV memory = 2 * n_layers * seq_len * n_kv_heads * head_dim * dtype_bytes

For Qwen3-0.6B-base at 4K context:
  = 2 * 28 * 4096 * 8 * 64 * 2 bytes = 234 MB (float16)

Quantizing to 4-bit: 234 * 4/16 = 58 MB (75% reduction)

But uniform quantization hurts quality — some layers' KV caches are much
more sensitive than others. OptiQ measures per-layer KV sensitivity and
assigns different bit-widths, preserving quality where it matters.

This module provides:
  1. Per-layer KV cache sensitivity analysis
  2. Mixed-precision KV cache bit assignment
  3. Integration with mlx-lm's generate pipeline
"""

import time
from dataclasses import dataclass

import mlx.core as mx
import mlx.nn as nn
import numpy as np


@dataclass
class KVSensitivityResult:
    """Per-layer KV cache sensitivity measurement."""
    layer_idx: int
    layer_name: str
    sensitivities: dict[int, float]  # bits -> KL divergence
    n_kv_heads: int
    head_dim: int
    cache_bytes_per_token: float  # bytes per token at full precision


def measure_kv_sensitivity(
    model_path: str,
    candidate_bits: list[int] = [4, 8],
    n_samples: int = 5,
    seq_len: int = 512,
    group_size: int = 64,
) -> list[KVSensitivityResult]:
    """Measure per-layer KV cache sensitivity to quantization.

    For each layer, quantizes only that layer's KV cache and measures
    the KL divergence of the output logits vs reference (full-precision
    KV cache).

    This is analogous to weight sensitivity analysis but for the KV cache:
    some layers' KV projections are robust to quantization (attention
    patterns are simple), while others are highly sensitive (complex
    multi-head interactions).

    Args:
        model_path: Path to MLX model directory
        candidate_bits: Bit-widths to test (e.g., [4, 8])
        n_samples: Number of calibration sequences
        seq_len: Length of each calibration sequence
        group_size: Quantization group size
    """
    from mlx_lm import load

    print(f"  Loading model from {model_path}...")
    model, tokenizer = load(model_path)
    mx.eval(model.parameters())

    n_layers = len(model.layers)

    # Load calibration data
    calibration_tokens = _get_calibration_tokens(tokenizer, n_samples, seq_len)
    print(f"  Calibration: {len(calibration_tokens)} sequences x {seq_len} tokens")

    # Step 1: Get reference logits with full-precision KV cache
    print(f"  Computing reference logits (full-precision KV cache)...")
    ref_logits_list = []
    for tokens in calibration_tokens:
        input_ids = mx.array([tokens])
        logits = model(input_ids)
        ref_logits_list.append(logits)
        mx.eval(logits)

    # Step 2: For each layer, quantize only that layer's KV cache
    results = []

    for layer_idx in range(n_layers):
        layer = model.layers[layer_idx]
        # Get KV head dimensions
        n_kv_heads = layer.self_attn.n_kv_heads
        # head_dim not always exposed; compute from scale = 1/sqrt(head_dim)
        head_dim = getattr(layer.self_attn, 'head_dim', None)
        if head_dim is None:
            head_dim = round(1.0 / (layer.self_attn.scale ** 2))
        bytes_per_token = 2 * n_kv_heads * head_dim * 2  # K+V, float16

        sensitivities = {}

        for bits in candidate_bits:
            kl_divs = []

            for sample_idx, tokens in enumerate(calibration_tokens):
                input_ids = mx.array([tokens])

                # Create cache with this layer quantized
                from mlx_lm.models.cache import KVCache, QuantizedKVCache

                cache = [KVCache() for _ in range(n_layers)]

                # First pass: fill the cache
                _ = model(input_ids, cache=cache)
                mx.eval([c.keys for c in cache if c.keys is not None])
                mx.eval([c.values for c in cache if c.keys is not None])

                # Quantize only the target layer's cache
                cache[layer_idx] = cache[layer_idx].to_quantized(
                    group_size=group_size, bits=bits
                )
                mx.eval(cache[layer_idx].keys)
                mx.eval(cache[layer_idx].values)

                # Second pass with the mixed cache: generate next-token logits
                # We need to do a fresh forward pass but using the cached KV
                # For this, we use the last token as input to get next-token logits
                last_token = mx.array([[tokens[-1]]])
                # Reset cache offsets to seq_len-1 so the model reads the cached KV
                # Actually, we can just compare the logits from the cached forward pass
                # Simpler approach: compare the logits from the initial pass
                # with a fresh pass using the quantized cache

                # Fresh pass with quantized cache
                cache_q = [KVCache() for _ in range(n_layers)]
                for i in range(n_layers):
                    if i == layer_idx:
                        cache_q[i] = QuantizedKVCache(
                            group_size=group_size, bits=bits
                        )
                    else:
                        cache_q[i] = KVCache()

                logits_q = model(input_ids, cache=cache_q)
                mx.eval(logits_q)

                # KL divergence between reference and quantized-cache logits
                ref_logits = ref_logits_list[sample_idx]
                kl = _compute_kl_divergence(ref_logits, logits_q)
                kl_divs.append(kl)

            mean_kl = float(np.mean(kl_divs))
            # Clamp to 0 — negative values are numerical noise
            sensitivities[bits] = max(mean_kl, 0.0)

        layer_name = f"model.layers.{layer_idx}"
        results.append(KVSensitivityResult(
            layer_idx=layer_idx,
            layer_name=layer_name,
            sensitivities=sensitivities,
            n_kv_heads=n_kv_heads,
            head_dim=head_dim,
            cache_bytes_per_token=bytes_per_token,
        ))

        # Print progress
        sens_str = ", ".join(f"{b}b:{sensitivities[b]:.4f}" for b in candidate_bits)
        print(f"  Layer {layer_idx:>2}/{n_layers}: {sens_str}")

    del model
    import gc
    gc.collect()

    return results


def _compute_kl_divergence(ref_logits: mx.array, test_logits: mx.array) -> float:
    """Compute KL divergence between two sets of logits."""
    # Use last position's logits for efficiency
    ref = ref_logits[0, -1, :]  # (vocab_size,)
    test = test_logits[0, -1, :]

    # Log-softmax
    ref_log_probs = ref - mx.logsumexp(ref, keepdims=True)
    test_log_probs = test - mx.logsumexp(test, keepdims=True)

    # KL(ref || test) = sum(ref_probs * (log_ref - log_test))
    ref_probs = mx.softmax(ref)
    kl = mx.sum(ref_probs * (ref_log_probs - test_log_probs))
    return float(kl.item())


def _get_calibration_tokens(tokenizer, n_samples, seq_len):
    """Load calibration tokens from WikiText-2."""
    try:
        from datasets import load_dataset
        ds = load_dataset("wikitext", "wikitext-2-raw-v1", split="validation")
        texts = [t for t in ds["text"] if len(t.strip()) > 100]
    except Exception:
        texts = [
            "The transformer architecture has revolutionized natural language processing. "
            "Self-attention mechanisms allow models to capture long-range dependencies "
            "in text sequences, enabling breakthrough performance on many tasks. "
        ] * n_samples * 5

    all_tokens = []
    for text in texts:
        tokens = tokenizer.encode(text, add_special_tokens=False)
        if len(tokens) >= seq_len:
            all_tokens.append(tokens[:seq_len])
        if len(all_tokens) >= n_samples:
            break

    # Pad if needed
    while len(all_tokens) < n_samples:
        # Combine texts to reach seq_len
        combined = []
        for text in texts:
            combined.extend(tokenizer.encode(text, add_special_tokens=False))
            if len(combined) >= seq_len:
                all_tokens.append(combined[:seq_len])
                break
        if len(all_tokens) >= n_samples:
            break

    return all_tokens[:n_samples]


@dataclass
class KVCacheConfig:
    """Per-layer KV cache quantization configuration."""
    layer_idx: int
    bits: int  # 0 = full precision (no quantization)
    group_size: int = 64


def optimize_kv_cache(
    sensitivity_results: list[KVSensitivityResult],
    target_kv_bits: float = 6.0,
    candidate_bits: list[int] = [4, 8],
    group_size: int = 64,
) -> list[KVCacheConfig]:
    """Assign per-layer KV cache bit-widths via greedy knapsack.

    Same algorithm as weight optimization: start all layers at min bits,
    greedily upgrade the most sensitive layers to higher bits until the
    average KV bits target is reached.

    Args:
        sensitivity_results: Per-layer KV sensitivity measurements
        target_kv_bits: Target average bits for KV cache
        candidate_bits: Available bit-widths
        group_size: Quantization group size

    Returns:
        List of per-layer KV cache configurations
    """
    import heapq

    candidate_bits = sorted(candidate_bits)
    min_bits = candidate_bits[0]
    max_bits = candidate_bits[-1]
    n_layers = len(sensitivity_results)

    # All layers have the same KV cache size per token (same heads/dim)
    # So we optimize purely on sensitivity
    allocation = {r.layer_idx: min_bits for r in sensitivity_results}

    # Build upgrade heap
    upgrade_heap = []
    for r in sensitivity_results:
        _push_kv_upgrade(upgrade_heap, r, min_bits, candidate_bits)

    # Budget: how many total "bit upgrades" we can do
    total_layers = n_layers
    bit_budget = (target_kv_bits - min_bits) * total_layers

    while upgrade_heap and bit_budget > 0:
        neg_eff, layer_idx, old_bits, new_bits, kl_reduction = heapq.heappop(upgrade_heap)

        if allocation[layer_idx] != old_bits:
            continue

        bit_cost = new_bits - old_bits
        if bit_cost > bit_budget * 1.05:
            continue

        allocation[layer_idx] = new_bits
        bit_budget -= bit_cost

        r = sensitivity_results[layer_idx]
        _push_kv_upgrade(upgrade_heap, r, new_bits, candidate_bits)

    configs = [
        KVCacheConfig(layer_idx=r.layer_idx, bits=allocation[r.layer_idx], group_size=group_size)
        for r in sensitivity_results
    ]

    achieved_bits = sum(c.bits for c in configs) / len(configs)
    n_high = sum(1 for c in configs if c.bits == max_bits)
    n_low = sum(1 for c in configs if c.bits == min_bits)

    print(f"  KV cache optimization: {n_low} layers @ {min_bits}-bit, "
          f"{n_high} layers @ {max_bits}-bit")
    print(f"  Target KV bits: {target_kv_bits:.1f}, achieved: {achieved_bits:.2f}")

    return configs


def _push_kv_upgrade(heap, r, current_bits, candidate_bits):
    """Push next KV cache upgrade onto the heap."""
    import heapq

    idx = candidate_bits.index(current_bits) if current_bits in candidate_bits else -1
    if idx < 0 or idx >= len(candidate_bits) - 1:
        return

    next_bits = candidate_bits[idx + 1]
    kl_current = r.sensitivities.get(current_bits, 0)
    kl_next = r.sensitivities.get(next_bits, 0)
    kl_reduction = kl_current - kl_next

    if kl_reduction <= 0:
        return

    # All layers have the same cache size, so efficiency = kl_reduction / bit_cost
    bit_cost = next_bits - current_bits
    efficiency = kl_reduction / max(bit_cost, 1)

    heapq.heappush(heap, (-efficiency, r.layer_idx, current_bits, next_bits, kl_reduction))


def make_mixed_kv_cache(kv_configs: list[KVCacheConfig]) -> list:
    """Create a per-layer mixed-precision KV cache.

    Returns a list of cache objects compatible with mlx-lm's generate pipeline.
    Each layer gets either a KVCache (full precision) or a QuantizedKVCache
    (at the assigned bit-width).
    """
    from mlx_lm.models.cache import KVCache, QuantizedKVCache

    caches = []
    for config in kv_configs:
        cache = QuantizedKVCache(group_size=config.group_size, bits=config.bits)
        caches.append(cache)

    return caches


def maybe_quantize_kv_mixed(prompt_cache, kv_configs, quantized_kv_start=0):
    """Per-layer KV cache quantization (replaces mlx-lm's uniform version).

    Args:
        prompt_cache: List of KVCache objects (one per layer)
        kv_configs: Per-layer bit-width assignments
        quantized_kv_start: Minimum offset before quantization triggers
    """
    from mlx_lm.models.cache import KVCache

    for config in kv_configs:
        idx = config.layer_idx
        if idx >= len(prompt_cache):
            continue
        c = prompt_cache[idx]
        # Only convert KVCache (not already quantized)
        if isinstance(c, KVCache) and c.offset >= quantized_kv_start:
            prompt_cache[idx] = c.to_quantized(
                group_size=config.group_size,
                bits=config.bits,
            )


def generate_with_mixed_kv(
    model,
    tokenizer,
    prompt: str,
    max_tokens: int = 100,
    kv_configs: list[KVCacheConfig] | None = None,
    temp: float = 0.0,
    verbose: bool = False,
) -> dict:
    """Generate text with per-layer mixed-precision KV cache.

    Uses mlx-lm's generate_step but with a custom KV cache that has
    different bit-widths per layer.

    Returns dict with generated text, tokens, timing info.
    """
    from mlx_lm.generate import generate_step
    from mlx_lm.sample_utils import make_sampler

    tokens = tokenizer.encode(prompt)
    prompt_tokens = mx.array(tokens)

    sampler = make_sampler(temp=temp)

    # Create mixed-precision cache from the start
    if kv_configs is not None:
        prompt_cache = make_mixed_kv_cache(kv_configs)
    else:
        prompt_cache = None

    generated_tokens = []
    token_times = []

    tic = time.perf_counter()
    for token, logprobs in generate_step(
        prompt_tokens,
        model,
        max_tokens=max_tokens,
        sampler=sampler,
        prompt_cache=prompt_cache,
    ):
        token_times.append(time.perf_counter() - tic)
        generated_tokens.append(token.item() if hasattr(token, 'item') else int(token))
        tic = time.perf_counter()

    text = tokenizer.decode(generated_tokens)

    # Memory estimate for KV cache
    n_total_tokens = len(tokens) + len(generated_tokens)
    kv_memory = 0
    if kv_configs:
        for config in kv_configs:
            # Each layer: 2 (K+V) * n_kv_heads * head_dim * n_tokens * bits/8
            # We don't know head_dim here, estimate from model
            layer = model.layers[config.layer_idx]
            n_kv_heads = layer.self_attn.n_kv_heads
            head_dim = getattr(layer.self_attn, 'head_dim', None)
            if head_dim is None:
                head_dim = round(1.0 / (layer.self_attn.scale ** 2))
            kv_memory += 2 * n_kv_heads * head_dim * n_total_tokens * config.bits / 8

    return {
        "text": text,
        "n_prompt_tokens": len(tokens),
        "n_generated_tokens": len(generated_tokens),
        "total_time": sum(token_times),
        "tokens_per_sec": len(generated_tokens) / sum(token_times) if token_times else 0,
        "kv_memory_bytes": kv_memory,
        "kv_memory_mb": kv_memory / (1024 * 1024),
    }


def estimate_kv_memory(
    n_layers: int,
    n_kv_heads: int,
    head_dim: int,
    seq_len: int,
    bits: float = 16.0,
) -> float:
    """Estimate KV cache memory in bytes.

    Args:
        n_layers: Number of transformer layers
        n_kv_heads: Number of KV heads per layer
        head_dim: Dimension per head
        seq_len: Sequence length
        bits: Average bits per element (16 for float16, 8/4 for quantized)

    Returns: Memory in bytes
    """
    # K + V = 2 tensors per layer
    elements = 2 * n_layers * seq_len * n_kv_heads * head_dim
    return elements * bits / 8
