import math
import torch
import torch.nn as nn
import torch.nn.functional as F


class KeyValueMemory(nn.Module):
    """A differentiable Key-Value memory store with gradient-friendly design.
    
    Think of this as a dynamically updatable database "beside" the weights.
    Structure:
    - Keys: Vectors representing "what" the info is about (indexing).
    - Values: Vectors representing the actual info content.
    
    Args:
        key_dim: Dimension of key vectors.
        value_dim: Dimension of value vectors.
        memory_size: Number of memory slots. Defaults to 1024.
        init_std: Standard deviation for random initialization. Defaults to 0.02.
    """
    def __init__(self, key_dim, value_dim, memory_size=1024, init_std=0.02):
        super().__init__()
        self.key_dim = key_dim
        self.value_dim = value_dim
        self.memory_size = memory_size
        self.scale = math.sqrt(key_dim)  # For scaled dot-product attention
        
        # Initialize with small random values instead of zeros to prevent
        # uniform softmax distribution which causes gradient issues
        self.register_buffer("keys", torch.randn(memory_size, key_dim) * init_std)
        self.register_buffer("values", torch.randn(memory_size, value_dim) * init_std)
        self.register_buffer("usage", torch.zeros(memory_size))
        self.register_buffer("num_writes", torch.tensor(0, dtype=torch.long))

    def write(self, keys, values, detach=True):
        """Write new facts into the memory.
        
        Args:
            keys: (batch_size, key_dim) - Embedding of the prompt/context
            values: (batch_size, value_dim) - Embedding of the answer/fact
            detach: Whether to detach gradients. Defaults to True.
        """
        batch_size = keys.size(0)
        
        if not hasattr(self, 'write_pointer'):
            self.write_pointer = 0
        
        # Compute indices for batch write (handles wrap-around)
        indices = (torch.arange(batch_size, device=keys.device) + self.write_pointer) % self.memory_size
        
        # Use functional updates via index_copy_ instead of in-place indexing
        # This prevents autograd errors when detach=False and gradients need to flow
        if detach:
            keys = keys.detach()
            values = values.detach()
        
        # Note: index_copy_ is still in-place on the buffer, but it properly
        # handles the source tensor's gradient tracking without breaking autograd
        with torch.no_grad():
            self.keys.index_copy_(0, indices, keys)
            self.values.index_copy_(0, indices, values)
            self.usage.index_fill_(0, indices, 1.0)
        
        self.write_pointer = (self.write_pointer + batch_size) % self.memory_size
        self.num_writes += batch_size

    def read(self, query, temperature=1.0):
        """Retrieve from memory using scaled dot-product attention.
        
        Args:
            query: (batch_size, key_dim) or (batch_size, seq_len, key_dim)
            temperature: Softmax temperature. Lower = sharper attention. Defaults to 1.0.
        
        Returns:
            retrieved_values: Weighted sum of relevant memories
            weights: Attention weights over memory slots
        """
        # Scaled dot-product attention for better gradient flow
        scores = torch.matmul(query, self.keys.t()) / self.scale
        
        # Apply temperature scaling
        if temperature != 1.0:
            scores = scores / temperature
        
        weights = F.softmax(scores, dim=-1)
        output = torch.matmul(weights, self.values)
        
        return output, weights

    def get_memory_utilization(self):
        """Return fraction of memory slots that have been written to."""
        return self.usage.mean().item()
    
    def save_memory(self, path: str):
        """Save memory contents to disk (separate from model weights).
        
        Args:
            path: File path to save memory.
            
        Example:
            >>> memory.save_memory("knowledge_base.mem")
        """
        torch.save({
            "keys": self.keys,
            "values": self.values,
            "usage": self.usage,
            "num_writes": self.num_writes,
            "write_pointer": getattr(self, 'write_pointer', 0),
            "config": {
                "key_dim": self.key_dim,
                "value_dim": self.value_dim,
                "memory_size": self.memory_size,
            }
        }, path)
    
    def load_memory(self, path: str, device: str = None):
        """Load memory contents from disk.
        
        Args:
            path: File path to load from.
            device: Device to load to. Uses current device if None.
            
        Example:
            >>> memory.load_memory("knowledge_base.mem")
        """
        if device is None:
            device = self.keys.device
        
        data = torch.load(path, map_location=device)
        
        # Resize if needed
        loaded_size = data["keys"].size(0)
        if loaded_size != self.memory_size:
            self.extend(loaded_size)
        
        self.keys.copy_(data["keys"])
        self.values.copy_(data["values"])
        self.usage.copy_(data["usage"])
        self.num_writes.copy_(data["num_writes"])
        self.write_pointer = data.get("write_pointer", 0)
    
    def extend(self, new_size: int):
        """Grow memory capacity dynamically.
        
        Args:
            new_size: New memory size (must be >= current size).
            
        Example:
            >>> memory.extend(8192)  # Double capacity
        """
        if new_size <= self.memory_size:
            return
        
        device = self.keys.device
        
        # Create new larger buffers
        new_keys = torch.randn(new_size, self.key_dim, device=device) * 0.02
        new_values = torch.randn(new_size, self.value_dim, device=device) * 0.02
        new_usage = torch.zeros(new_size, device=device)
        
        # Copy existing data
        new_keys[:self.memory_size] = self.keys
        new_values[:self.memory_size] = self.values
        new_usage[:self.memory_size] = self.usage
        
        # Replace buffers
        self.register_buffer("keys", new_keys)
        self.register_buffer("values", new_values)
        self.register_buffer("usage", new_usage)
        self.memory_size = new_size
    
    def clear(self):
        """Clear all memory contents."""
        self.keys.zero_()
        self.values.zero_()
        self.usage.zero_()
        self.num_writes.zero_()
        self.write_pointer = 0


class RetrievalBlock(nn.Module):
    """A Transformer Block that includes a request to External Memory.
    
    Features gradient-friendly design with:
    - Scaled attention for stable gradients
    - Value projection for dimension compatibility
    - Learnable memory gate with proper initialization
    - Pre-norm architecture for better gradient flow
    
    .. warning::
        **Scalability**: This block computes full attention over the memory bank,
        which is O(Batch × Seq × Memory_Size). For memory_size > ~10,000, this
        becomes a bottleneck. For large knowledge bases (100K+ entries), use
        :class:`~composennent.modules.internal_rag.InternalRAGBlock` with
        :class:`~composennent.modules.internal_rag.ScalableMemory` instead,
        which provides O(k) retrieval via FAISS-backed top-k search.
    
    Args:
        hidden_dim: Hidden dimension of the transformer.
        memory_component: KeyValueMemory instance.
        dropout: Dropout probability. Defaults to 0.1.
        temperature: Softmax temperature for memory attention. Defaults to 1.0.
    """
    def __init__(self, hidden_dim, memory_component, dropout=0.1, temperature=1.0):
        super().__init__()
        self.hidden_dim = hidden_dim
        self.memory = memory_component
        self.temperature = temperature
        self.scale = math.sqrt(memory_component.key_dim)
        
        # Query projection
        self.query_proj = nn.Linear(hidden_dim, memory_component.key_dim)
        
        # Value projection if dimensions don't match
        if memory_component.value_dim != hidden_dim:
            self.value_proj = nn.Linear(memory_component.value_dim, hidden_dim)
        else:
            self.value_proj = nn.Identity()
        
        # Learnable memory gate - initialized to small values so memory 
        # contribution starts small and grows during training
        self.memory_gate = nn.Sequential(
            nn.Linear(hidden_dim * 2, hidden_dim),
            nn.Sigmoid()
        )
        # Initialize gate bias to negative value so initial gate output is small
        # This prevents memory from dominating early in training
        nn.init.constant_(self.memory_gate[0].bias, -2.0)
        
        self.layer_norm = nn.LayerNorm(hidden_dim)
        self.dropout = nn.Dropout(dropout)

    def forward(self, hidden_states):
        """Forward pass with gradient-friendly memory integration.
        
        Args:
            hidden_states: (batch_size, seq_len, hidden_dim)
            
        Returns:
            Output tensor of shape (batch_size, seq_len, hidden_dim)
        """
        # Pre-norm for better gradient flow
        normed = self.layer_norm(hidden_states)
        
        # Generate queries
        queries = self.query_proj(normed)  # (batch, seq, key_dim)
        
        # Scaled attention over memory
        scores = torch.matmul(queries, self.memory.keys.t()) / self.scale
        
        if self.temperature != 1.0:
            scores = scores / self.temperature
        
        weights = F.softmax(scores, dim=-1)
        
        # Retrieve and project values
        retrieved = torch.matmul(weights, self.memory.values)
        retrieved = self.value_proj(retrieved)  # Now (batch, seq, hidden_dim)
        retrieved = self.dropout(retrieved)
        
        # Gated integration - concatenate hidden and retrieved
        combined = torch.cat([normed, retrieved], dim=-1)
        gate = self.memory_gate(combined)  # (batch, seq, hidden_dim) with values in (0, 1)
        
        # Gated residual: output = hidden + gate * retrieved
        # This allows gradient to flow through residual connection even if gate is small
        memory_contribution = gate * retrieved
        output = hidden_states + memory_contribution
        
        return output


class MemoryHead(nn.Module):
    """Prediction head for Self-Modifying Memory.
    
    Decides: 
    1. WHETHER to write (Gate)
    2. WHAT to address (Key)
    3. WHAT content to write (Value)
    
    Args:
        hidden_dim: Hidden dimension of the model.
        key_dim: Dimension of memory keys.
        value_dim: Dimension of memory values.
    """
    def __init__(self, hidden_dim, key_dim, value_dim):
        super().__init__()
        self.write_gate = nn.Linear(hidden_dim, 1)
        self.write_key = nn.Linear(hidden_dim, key_dim)
        self.write_value = nn.Linear(hidden_dim, value_dim)
        
        # Initialize gate bias to negative so writes are rare initially
        nn.init.constant_(self.write_gate.bias, -3.0)
        
    def forward(self, hidden_state):
        """Generate write decision, key, and value.
        
        Args:
            hidden_state: (batch, hidden_dim) - usually the last token's state
            
        Returns:
            gate: Write probability (batch, 1)
            key: Key to write (batch, key_dim)
            value: Value to write (batch, value_dim)
        """
        gate = torch.sigmoid(self.write_gate(hidden_state))
        key = self.write_key(hidden_state)
        value = self.write_value(hidden_state)
        
        return gate, key, value
