models/context_window.py

"""
Context Window — Janela deslizante com cache KV para sequências longas.

Implementa três estratégias (ver docs/MATH_ANALYSIS.md seção 3.4):
  - Sliding Window Pura
  - Sink + Sliding (StreamingLLM) — RECOMENDADO
  - Attention Recomputation (sem cache, máxima precisão)

Para CNN-BiGRU cooperativo, o contexto é gerenciado em modo:
  - Encoder (treino/classificação): processa janela inteira, BiGRU bidirecional
  - Decoder (geração): cache KV para CausalSelfAttention + GRU unidirecional

Autor: CNN-BiGRU Project
"""
from __future__ import annotations

import logging
from dataclasses import dataclass, field
from typing import Dict, List, Optional, Tuple

import torch
import torch.nn as nn

logger = logging.getLogger(__name__)


# ============================================================================
# Configuração
# ============================================================================

@dataclass
class ContextWindowConfig:
    """Configuração da janela de contexto."""
    # Comprimento máximo da janela (em tokens)
    max_window: int = 512
    # Estratégia de evicção: "sliding" | "sink_sliding" | "recompute"
    eviction_strategy: str = "sink_sliding"
    # Número de tokens "sink" (apenas para sink_sliding) — tipicamente BOS + system prompt
    n_sink_tokens: int = 4
    # Dimensão do embedding (para alocar cache KV)
    embed_dim: int = 256
    # Número de cabeças (para cache KV em multi-head attention)
    n_heads: int = 4
    # Dimensão por cabeça
    head_dim: Optional[int] = None  # default: embed_dim // n_heads
    # Número de camadas (para cache KV multicamada)
    n_layers: int = 2
    # Device padrão
    device: str = "cpu"
    # Dtype do cache (None = manter fp32)
    dtype: Optional[torch.dtype] = None


# ============================================================================
# KV Cache
# ============================================================================

class KVCache:
    """
    Cache de pares (Key, Value) para CausalSelfAttention.

    Shape por camada:
      K: [n_heads, seq_cached, head_dim]
      V: [n_heads, seq_cached, head_dim]

    Em modo batch:
      K: [batch, n_heads, seq_cached, head_dim]
      V: [batch, n_heads, seq_cached, head_dim]
    """

    def __init__(self, n_layers: int, batch_size: int, n_heads: int,
                 head_dim: int, device: torch.device, dtype: torch.dtype):
        self.n_layers = n_layers
        self.batch_size = batch_size
        self.n_heads = n_heads
        self.head_dim = head_dim
        self.device = device
        self.dtype = dtype
        # Pre-alocar listas vazias; preencher lazy no primeiro update
        self.keys: List[Optional[torch.Tensor]] = [None] * n_layers
        self.values: List[Optional[torch.Tensor]] = [None] * n_layers

    def update(self, layer_idx: int, new_k: torch.Tensor, new_v: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        Atualiza o cache da camada layer_idx com novos K, V.

        Args:
            layer_idx: índice da camada [0, n_layers)
            new_k: [batch, n_heads, seq_new, head_dim]
            new_v: [batch, n_heads, seq_new, head_dim]

        Returns:
            (cached_k, cached_v): [batch, n_heads, seq_total, head_dim]
        """
        if not (0 <= layer_idx < self.n_layers):
            raise IndexError(f"layer_idx {layer_idx} fora de range [0, {self.n_layers})")

        # Validar shapes
        if new_k.dim() != 4 or new_v.dim() != 4:
            raise ValueError(
                f"new_k e new_v devem ser 4D [batch, n_heads, seq, head_dim]; "
                f"recebido {new_k.shape}"
            )

        # Converter dtype/device
        new_k = new_k.to(device=self.device, dtype=self.dtype)
        new_v = new_v.to(device=self.device, dtype=self.dtype)

        if self.keys[layer_idx] is None:
            self.keys[layer_idx] = new_k
            self.values[layer_idx] = new_v
        else:
            self.keys[layer_idx] = torch.cat([self.keys[layer_idx], new_k], dim=2)
            self.values[layer_idx] = torch.cat([self.values[layer_idx], new_v], dim=2)

        return self.keys[layer_idx], self.values[layer_idx]

    def evict_sliding(self, max_keep: int) -> None:
        """
        Estratégia sliding window pura: mantém apenas os últimos max_keep tokens.
        """
        for layer_idx in range(self.n_layers):
            if self.keys[layer_idx] is None:
                continue
            seq = self.keys[layer_idx].size(2)
            if seq > max_keep:
                start = seq - max_keep
                self.keys[layer_idx] = self.keys[layer_idx][:, :, start:, :].contiguous()
                self.values[layer_idx] = self.values[layer_idx][:, :, start:, :].contiguous()

    def evict_sink_sliding(self, max_keep: int, n_sink: int) -> None:
        """
        Estratégia sink + sliding: mantém os primeiros n_sink + últimos (max_keep - n_sink).
        """
        if n_sink >= max_keep:
            logger.warning(
                f"n_sink ({n_sink}) >= max_keep ({max_keep}); usando sliding puro"
            )
            self.evict_sliding(max_keep)
            return

        n_sliding = max_keep - n_sink
        for layer_idx in range(self.n_layers):
            if self.keys[layer_idx] is None:
                continue
            seq = self.keys[layer_idx].size(2)
            if seq > max_keep:
                sink_k = self.keys[layer_idx][:, :, :n_sink, :]
                sink_v = self.values[layer_idx][:, :, :n_sink, :]
                sliding_k = self.keys[layer_idx][:, :, -n_sliding:, :]
                sliding_v = self.values[layer_idx][:, :, -n_sliding:, :]
                self.keys[layer_idx] = torch.cat([sink_k, sliding_k], dim=2).contiguous()
                self.values[layer_idx] = torch.cat([sink_v, sliding_v], dim=2).contiguous()

    def reset(self) -> None:
        """Limpa o cache completamente."""
        self.keys = [None] * self.n_layers
        self.values = [None] * self.n_layers

    def get_seq_len(self, layer_idx: int = 0) -> int:
        if self.keys[layer_idx] is None:
            return 0
        return self.keys[layer_idx].size(2)

    def total_tokens(self) -> int:
        return max(self.get_seq_len(i) for i in range(self.n_layers))


# ============================================================================
# Context Window Manager
# ============================================================================

class ContextWindowManager:
    """
    Gerencia a janela de contexto aplicando a política de evicção configurada.
    """

    def __init__(self, config: ContextWindowConfig):
        self.config = config
        self.head_dim = config.head_dim or (config.embed_dim // config.n_heads)
        if config.embed_dim % config.n_heads != 0:
            raise ValueError(
                f"embed_dim ({config.embed_dim}) deve ser divisível por "
                f"n_heads ({config.n_heads})"
            )
        self.kv_cache: Optional[KVCache] = None
        # Sequência de tokens brutos (para recomputação se necessário)
        self.token_history: List[torch.Tensor] = []

    # ----------------------------------------------------------------------
    # Inicialização
    # ----------------------------------------------------------------------

    def init_cache(self, batch_size: int, device: torch.device,
                   dtype: Optional[torch.dtype] = None) -> KVCache:
        """Cria novo cache KV para uma sessão de geração."""
        dtype = dtype or self.config.dtype or torch.float32
        self.kv_cache = KVCache(
            n_layers=self.config.n_layers,
            batch_size=batch_size,
            n_heads=self.config.n_heads,
            head_dim=self.head_dim,
            device=device,
            dtype=dtype,
        )
        self.token_history = []
        return self.kv_cache

    # ----------------------------------------------------------------------
    # Adicionar tokens
    # ----------------------------------------------------------------------

    def append_tokens(self, token_ids: torch.Tensor) -> torch.Tensor:
        """
        Adiciona tokens ao histórico e retorna a janela ativa.

        Args:
            token_ids: [batch, seq_new]

        Returns:
            window: [batch, seq_window] — tokens na janela ativa após evicção
        """
        if token_ids.dim() != 2:
            raise ValueError(f"token_ids deve ser [batch, seq]; recebido {token_ids.shape}")

        # Para simplificação, mantemos histórico apenas do batch[0]
        # (geração autoregressiva tipicamente batch=1)
        if token_ids.size(0) > 1:
            logger.warning("ContextWindowManager.append_tokens: batch>1, histórico rastreia apenas batch[0]")

        self.token_history.append(token_ids)

        # Concatenar todo o histórico
        full = torch.cat(self.token_history, dim=1)

        # Aplicar evicção se necessário
        return self._evict_tokens(full)

    def _evict_tokens(self, full_seq: torch.Tensor) -> torch.Tensor:
        """Aplica política de evicção à sequência completa."""
        seq_len = full_seq.size(1)
        max_keep = self.config.max_window

        if seq_len <= max_keep:
            return full_seq

        strategy = self.config.eviction_strategy
        if strategy == "sliding":
            return full_seq[:, -max_keep:]
        elif strategy == "sink_sliding":
            n_sink = self.config.n_sink_tokens
            n_sliding = max_keep - n_sink
            sink = full_seq[:, :n_sink]
            sliding = full_seq[:, -n_sliding:]
            return torch.cat([sink, sliding], dim=1)
        elif strategy == "recompute":
            # Manter tudo (atenção recomputada)
            return full_seq
        else:
            logger.warning(f"Estratégia de evicção desconhecida: {strategy}; usando sliding")
            return full_seq[:, -max_keep:]

    # ----------------------------------------------------------------------
    # Evict cache KV
    # ----------------------------------------------------------------------

    def evict_cache(self) -> None:
        """Aplica a política de evicção ao cache KV."""
        if self.kv_cache is None:
            return
        strategy = self.config.eviction_strategy
        max_keep = self.config.max_window
        if strategy == "sliding":
            self.kv_cache.evict_sliding(max_keep)
        elif strategy == "sink_sliding":
            self.kv_cache.evict_sink_sliding(max_keep, self.config.n_sink_tokens)
        elif strategy == "recompute":
            self.kv_cache.reset()
        else:
            self.kv_cache.evict_sliding(max_keep)

    # ----------------------------------------------------------------------
    # Reset
    # ----------------------------------------------------------------------

    def reset(self) -> None:
        """Reseta todo o estado (cache + histórico)."""
        if self.kv_cache is not None:
            self.kv_cache.reset()
        self.kv_cache = None
        self.token_history = []

    # ----------------------------------------------------------------------
    # Info
    # ----------------------------------------------------------------------

    def get_info(self) -> Dict:
        """Retorna informações sobre o estado atual."""
        cache_len = self.kv_cache.total_tokens() if self.kv_cache else 0
        history_len = sum(t.size(1) for t in self.token_history)
        return {
            "max_window": self.config.max_window,
            "strategy": self.config.eviction_strategy,
            "n_sink": self.config.n_sink_tokens,
            "cache_tokens": cache_len,
            "history_tokens": history_len,
            "active": self.kv_cache is not None,
        }


# ============================================================================
# Helpers
# ============================================================================

def make_context_window(
    max_window: int = 512,
    strategy: str = "sink_sliding",
    n_sink: int = 4,
    embed_dim: int = 256,
    n_heads: int = 4,
    n_layers: int = 2,
    device: str = "cpu",
) -> ContextWindowManager:
    """Factory para criar ContextWindowManager com defaults sensatos."""
    config = ContextWindowConfig(
        max_window=max_window,
        eviction_strategy=strategy,
        n_sink_tokens=n_sink,
        embed_dim=embed_dim,
        n_heads=n_heads,
        head_dim=embed_dim // n_heads,
        n_layers=n_layers,
        device=device,
    )
    return ContextWindowManager(config)


__all__ = [
    "ContextWindowConfig",
    "KVCache",
    "ContextWindowManager",
    "make_context_window",
]