"""Position-selection policies over real caches.

- sink_recent (B3): StreamingLLM-inspired *retention pattern* (first `sinks` positions + most recent). Unlike StreamingLLM,
  retained keys keep their ORIGINAL positions (StreamingLLM re-rotates to positions-within-cache); labelled 'inspired'.
- random (B4): uniform random positions, seeded; matched budget.
- attention_informed (B5): faithful re-implementation of SnapKV's selection rule (observation window = last W prompt
  tokens; per kv-head attention of window queries over prefix keys, averaged over the window queries, 1-D pooled with
  kernel k (avg or max), top-k per kv-head; window positions always kept). GQA handled by averaging query heads in each
  kv group. Requires the window queries: we capture them with a forward hook (an extra pass over the window tokens on a
  *copy* of the cache), so optimized fused attention is NOT disabled for the main computation; the extra pass cost is
  recorded as compression time.
- pyramid (B6): PyramidKV-style layer budget allocation (arithmetic decrease from bottom to top layer, total budget
  preserved) using the attention_informed scores within each layer; labelled 'faithful_reimplementation' of the budget
  rule (verified against KVCache-Factory code: base=cap-w, min=base//beta, max=2*base-min, steps=(max-min)//(L-1)).
"""
from __future__ import annotations
import math
import numpy as np
import mlx.core as mx
import mlx.nn as nn
from .sparse import SparseKVCache, select_positions, copy_cache, dense_kv


def sink_recent_indices(T: int, keep: int, sinks: int = 4) -> np.ndarray:
    keep = max(min(keep, T), 1)
    s = min(sinks, keep)
    rec = keep - s
    idx = np.concatenate([np.arange(s), np.arange(T - rec, T)]) if rec > 0 else np.arange(s)
    return np.unique(idx)


def random_indices(T: int, keep: int, seed: int) -> np.ndarray:
    rng = np.random.default_rng(seed)
    keep = max(min(keep, T), 1)
    return np.sort(rng.choice(T, keep, replace=False))


class _QueryCapture:
    """Context manager: temporarily patch the attention class __call__ to record post-RoPE queries per layer."""

    def __init__(self, model):
        self.model = model
        inner = model.model if hasattr(model, "model") and hasattr(model.model, "layers") else model
        self.layers = inner.layers
        self.queries: list = [None] * len(self.layers)
        self.index_of = {id(l.self_attn): i for i, l in enumerate(self.layers)}
        self.cls = type(self.layers[0].self_attn)
        self.orig = self.cls.__call__

    def __enter__(self):
        cap = self; orig = self.orig

        def patched(self_attn, x, mask=None, cache=None):
            idx = cap.index_of.get(id(self_attn))
            if idx is not None:
                B, L, D = x.shape
                q = self_attn.q_proj(x).reshape(B, L, self_attn.n_heads, -1).transpose(0, 2, 1, 3)
                q = self_attn.rope(q, offset=cache.offset) if cache is not None else self_attn.rope(q)
                cap.queries[idx] = q
            return orig(self_attn, x, mask=mask, cache=cache)
        self.cls.__call__ = patched
        return self

    def __exit__(self, *a):
        self.cls.__call__ = self.orig


def attention_window_scores(model, cache: list, window_token_ids: list[int], pool_kernel: int = 7, pooling: str = "max") -> list[np.ndarray]:
    """Return per-layer scores [H_kv, T_prefix] = pooled mean attention of the window queries onto prefix keys.
    The window tokens are processed on a *copy* of the cache (the original stays immutable)."""
    if hasattr(model, "model") and not hasattr(model, "layers") and hasattr(model, "info"):
        model = model.model   # MLXModel wrapper -> underlying nn.Module
    work = copy_cache(cache)
    T = work[0].offset
    toks = mx.array(window_token_ids, dtype=mx.int32)[None]
    with _QueryCapture(model) as qc:
        model(toks, cache=work)
        mx.eval([q for q in qc.queries])
    scores = []
    for layer_idx, c in enumerate(work):
        k, _ = c.keys_and_values()               # [B, H_kv, T+W, d]
        k = k[..., :T, :].astype(mx.float32)
        q = qc.queries[layer_idx].astype(mx.float32)   # [B, H_q, W, d]
        Hq, Hkv = q.shape[1], k.shape[1]
        rep = Hq // Hkv
        q = q.reshape(q.shape[0], Hkv, rep, q.shape[2], q.shape[3])   # [B, H_kv, rep, W, d]
        s = (q @ k[:, :, None, :, :].transpose(0, 1, 2, 4, 3)) / math.sqrt(k.shape[-1])  # [B, H_kv, rep, W, T]
        s = mx.softmax(s, axis=-1)
        s = s.mean(axis=(0, 2, 3))      # [H_kv, T]: mean over batch, group heads, window queries
        a = np.array(s)
        if pool_kernel > 1:
            pad = pool_kernel // 2
            ap = np.pad(a, ((0, 0), (pad, pad)), mode="edge")
            win = np.lib.stride_tricks.sliding_window_view(ap, pool_kernel, axis=1)
            a = win.max(axis=-1) if pooling == "max" else win.mean(axis=-1)
        scores.append(a)
    return scores


def topk_indices_per_head(score: np.ndarray, keep: int, always_keep: np.ndarray | None = None) -> np.ndarray:
    """score [H, T] -> indices [H, keep] sorted ascending per head; always_keep positions forced in."""
    H, T = score.shape
    keep = max(min(keep, T), 1)
    out = np.zeros((H, keep), dtype=np.int64)
    for h in range(H):
        s = score[h].copy()
        if always_keep is not None and len(always_keep):
            s[always_keep] = np.inf
        idx = np.argpartition(-s, keep - 1)[:keep]
        out[h] = np.sort(idx)
    return out


def pyramid_budgets(total_keep: int, n_layers: int, window: int, beta: int = 20) -> list[int]:
    """PyramidKV budget rule (KVCache-Factory): per-layer budgets decreasing linearly, summing ~ total."""
    cap = total_keep
    base = max(cap - window, 1)
    mn = max(base // beta, 1)
    mxb = 2 * base - mn
    steps = (mxb - mn) // max(n_layers - 1, 1)
    budgets = [max(mxb - l * steps, mn) + window for l in range(n_layers)]
    # rescale to preserve total budget exactly
    scale = total_keep * n_layers / sum(budgets)
    budgets = [max(1, int(round(b * scale))) for b in budgets]
    return budgets


def select_cache(model, cache: list, policy: str, keep_frac: float, window_token_ids: list[int] | None = None,
                 seed: int = 0, sinks: int = 4, pool_kernel: int = 7, pooling: str = "max", window_keep: int = 32) -> tuple[list, dict]:
    """Build a sparse cache under `policy` retaining round(keep_frac*T) positions. Returns (cache, info)."""
    T = cache[0].offset
    keep = max(1, int(round(keep_frac * T)))
    L = len(cache)
    info = {"policy": policy, "keep_frac": keep_frac, "T": T, "keep": keep}
    import time
    t0 = time.perf_counter()
    if policy == "sink_recent":
        idx = sink_recent_indices(T, keep, sinks)
        out = select_positions(cache, [idx] * L)
    elif policy == "random":
        idx = random_indices(T, keep, seed)
        out = select_positions(cache, [idx] * L)
    elif policy in ("attention", "pyramid"):
        if not window_token_ids:
            raise ValueError("attention policies need window_token_ids")
        scores = attention_window_scores(model, cache, window_token_ids, pool_kernel, pooling)
        always = np.arange(max(0, T - window_keep), T)
        if policy == "attention":
            per_layer = [topk_indices_per_head(s, keep, always) for s in scores]
        else:
            budgets = pyramid_budgets(keep, L, window_keep)
            per_layer = [topk_indices_per_head(s, b, always) for s, b in zip(scores, budgets)]
            info["layer_budgets"] = budgets
        out = select_positions(cache, per_layer)
        info["window_tokens"] = len(window_token_ids)
    else:
        raise ValueError(policy)
    mx.eval([x for c in out for x in (c.keys, c.values)])
    info["compress_ms"] = (time.perf_counter() - t0) * 1000.0
    return out, info