"""
LMCODE: Language Model with Memory CODE
=======================================

A memory-augmented language model with dual memory systems:
- Short-term memory: Working memory for immediate context (like Transformer KV cache)
- Long-term memory: Persistent storage for knowledge and experiences

Inspired by:
- LongMem (2023): Augmenting LLMs with Long-Term Memory
- MemoRAG (2024): Dual-system RAG with Global and Local Memory
- CAMELoT (2024): Training-free Consolidated Associative Memory
"""

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


class ShortTermMemory(nn.Module):
    """
    Short-term memory module using recurrent state updates.
    Acts as working memory for immediate context (similar to Transformer KV cache).
    """
    
    def __init__(self, hidden_size: int, memory_size: int = 512):
        super().__init__()
        self.hidden_size = hidden_size
        self.memory_size = memory_size
        
        # Memory state
        self.memory = nn.Parameter(torch.zeros(1, memory_size, hidden_size))
        
        # Memory update gate
        self.update_gate = nn.Linear(hidden_size, 1)
        
        # Memory read/write projections
        self.read_proj = nn.Linear(hidden_size, hidden_size)
        self.write_proj = nn.Linear(hidden_size, hidden_size)
        
        # Initialize memory
        nn.init.normal_(self.memory, mean=0, std=0.02)
    
    def forward(self, hidden_states: torch.Tensor, 
                attention_mask: Optional[torch.Tensor] = None) -> Dict[str, torch.Tensor]:
        """
        Update and read from short-term memory.
        
        Args:
            hidden_states: Current hidden states [batch_size, seq_len, hidden_size]
            attention_mask: Optional attention mask
            
        Returns:
            Dictionary with memory outputs
        """
        batch_size, seq_len, _ = hidden_states.shape
        
        # Expand memory for batch
        memory = self.memory.expand(batch_size, -1, -1)  # [batch, memory_size, hidden]
        
        # Compute update gates
        update_scores = torch.sigmoid(self.update_gate(hidden_states))  # [batch, seq, 1]
        
        # Read from memory
        read_query = self.read_proj(hidden_states)  # [batch, seq, hidden]
        read_scores = torch.matmul(read_query, memory.transpose(-1, -2))  # [batch, seq, memory_size]
        read_scores = read_scores / math.sqrt(self.hidden_size)
        
        if attention_mask is not None:
            read_scores = read_scores.masked_fill(~attention_mask.unsqueeze(-1).bool(), -1e9)
        
        read_weights = F.softmax(read_scores, dim=-1)  # [batch, seq, memory_size]
        memory_read = torch.matmul(read_weights, memory)  # [batch, seq, hidden]
        
        # Write to memory (soft update)
        write_values = self.write_proj(hidden_states)  # [batch, seq, hidden]
        write_weights = update_scores.expand(-1, -1, self.memory_size)  # [batch, seq, memory_size]
        
        # Update memory (detached to prevent gradient flow through time)
        memory_delta = torch.matmul(write_weights.transpose(-1, -2), write_values)  # [batch, memory_size, hidden]
        memory = memory + 0.1 * memory_delta  # Small learning rate for memory update
        
        # Combine hidden states with memory read
        output = hidden_states + 0.5 * memory_read  # Residual connection
        
        return {
            'output': output,
            'memory_read': memory_read,
            'read_weights': read_weights,
            'updated_memory': memory
        }