import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import Optional, List

# ---------------------------------------------------------------------------
# p2.8 "Pure" Modules - Zero-Shot Optimized (No Trainable Params)
# ---------------------------------------------------------------------------

class ReadOnlyCache:
    """
    Wrapper for transformers.Cache that prevents updates.
    Used for recurrent loops t > 0 to maintain context without corruption.
    """
    def __init__(self, real_cache):
        self.__dict__["_real"] = real_cache
        
    def __getattr__(self, name):
        return getattr(self._real, name)
        
    def update(self, key_states, value_states, layer_idx, cache_kwargs=None):
        # Return the already-updated sequences from the real cache
        # Supporting Transformers 5.x DynamicLayer structure
        if hasattr(self._real, "layers"):
            layer = self._real.layers[layer_idx]
            return layer.keys, layer.values
        # Fallback for older versions
        return self._real.key_cache[layer_idx], self._real.value_cache[layer_idx]

class LTIInjection(nn.Module):
    """
    Pure-functional LTI injection. 
    Uses fixed coefficients to provide 'Computational Headroom' without training.
    """
    def __init__(self, dim: int):
        super().__init__()
        self.input_norm = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
        self.gamma = 0.08 # Winning Gamma from Sweep

    def forward(self, h: torch.Tensor, e: torch.Tensor, transformer_out: torch.Tensor) -> torch.Tensor:
        # Coefficients for stable zero-shot recursion
        # h_new = transformer_out + gamma * (e_norm - h)
        
        e_norm = self.input_norm(e).to(h.dtype)
        
        # Identity-centered update rule
        return transformer_out + self.gamma * (e_norm - h)

class StabilityMonitor:
    """Parameter-free heuristics for Phi (Φ) and Lambda (λ)."""
    
    @staticmethod
    def calculate_phi(h_new: torch.Tensor, h_old: torch.Tensor) -> torch.Tensor:
        """Measure internal state stability (Φ) via Cosine Similarity."""
        B = h_new.shape[0]
        return F.cosine_similarity(h_new.view(B, -1), h_old.view(B, -1), dim=-1).mean()

    @staticmethod
    def detect_lambda(hidden_states: torch.Tensor, e: torch.Tensor) -> torch.Tensor:
        """
        Detect embedding mode (λ).
        Measures how much the current state has diverged from the initial anchor.
        High divergence often indicates 'Self-Reflection' or 'Deep Reasoning'.
        """
        B = hidden_states.shape[0]
        # In p2.8, λ is high if we are in 'We' mode (deeply embedded)
        # Heuristic: 1 - cosine_similarity(h, e)
        dist = 1.0 - F.cosine_similarity(hidden_states.view(B, -1), e.view(B, -1), dim=-1).mean()
        return dist.clamp(0, 1)

# Note: ACTHalting, LoRAAdapter, and IntrospectiveDelta removed.
# They are 'useless leftovers' in a pure zero-shot context.