src/cognitive/working_memory.py

"""
Working Memory Management System with Attention-based Focus and Context Switching
Implements dynamic attention mechanisms and context management for cognitive agents
"""
import sqlite3
import json
import uuid
from datetime import datetime, timedelta
from typing import Dict, List, Any, Optional, Tuple
from dataclasses import dataclass, asdict
import logging
from pathlib import Path
import heapq
import threading
import time

logger = logging.getLogger(__name__)

@dataclass
class WorkingMemoryItem:
    """Individual item in working memory"""
    id: str
    content: str
    item_type: str  # goal, observation, hypothesis, plan, etc.
    priority: float  # 0.0-1.0, higher is more important
    activation_level: float  # 0.0-1.0, current activation
    created_at: datetime
    last_accessed: datetime
    access_count: int
    decay_rate: float  # how quickly activation decays
    context_tags: List[str]
    source_agent: str
    related_items: List[str]  # IDs of related items

@dataclass
class AttentionFocus:
    """Current attention focus with weighted priorities"""
    focus_id: str
    focus_type: str  # task, threat, goal, etc.
    focus_items: List[str]  # Working memory item IDs
    attention_weight: float  # 0.0-1.0
    duration: timedelta
    created_at: datetime
    metadata: Dict[str, Any]

class WorkingMemoryManager:
    """Advanced working memory with attention-based focus management"""
    
    def __init__(self, db_path: str = "data/cognitive/working_memory.db",
                 capacity: int = 50, decay_interval: float = 30.0):
        """Initialize working memory system"""
        self.db_path = Path(db_path)
        self.db_path.parent.mkdir(parents=True, exist_ok=True)
        self.capacity = capacity  # Maximum items in working memory
        self.decay_interval = decay_interval  # Seconds between decay updates
        
        self._init_database()
        self._memory_items = {}  # In-memory cache
        self._attention_focus = None
        self._attention_history = []
        
        # Start background decay process
        self._decay_thread = threading.Thread(target=self._decay_loop, daemon=True)
        self._decay_running = True
        self._decay_thread.start()
        
        # Load existing items
        self._load_working_memory()
    
    def _init_database(self):
        """Initialize database schemas"""
        with sqlite3.connect(self.db_path) as conn:
            conn.execute("""
                CREATE TABLE IF NOT EXISTS working_memory_items (
                    id TEXT PRIMARY KEY,
                    content TEXT NOT NULL,
                    item_type TEXT NOT NULL,
                    priority REAL NOT NULL,
                    activation_level REAL NOT NULL,
                    created_at TEXT NOT NULL,
                    last_accessed TEXT NOT NULL,
                    access_count INTEGER DEFAULT 0,
                    decay_rate REAL DEFAULT 0.1,
                    context_tags TEXT,
                    source_agent TEXT,
                    related_items TEXT,
                    updated_at TEXT DEFAULT CURRENT_TIMESTAMP
                )
            """)
            
            conn.execute("""
                CREATE TABLE IF NOT EXISTS attention_focus_log (
                    id TEXT PRIMARY KEY,
                    focus_type TEXT NOT NULL,
                    focus_items TEXT,
                    attention_weight REAL NOT NULL,
                    duration_seconds REAL,
                    created_at TEXT NOT NULL,
                    ended_at TEXT,
                    metadata TEXT,
                    agent_id TEXT
                )
            """)
            
            conn.execute("""
                CREATE TABLE IF NOT EXISTS context_switches (
                    id TEXT PRIMARY KEY,
                    from_focus TEXT,
                    to_focus TEXT,
                    switch_reason TEXT,
                    switch_cost REAL,
                    timestamp TEXT DEFAULT CURRENT_TIMESTAMP,
                    agent_id TEXT
                )
            """)
            
            # Create indices
            conn.execute("CREATE INDEX IF NOT EXISTS idx_wm_priority ON working_memory_items(priority DESC)")
            conn.execute("CREATE INDEX IF NOT EXISTS idx_wm_activation ON working_memory_items(activation_level DESC)")
            conn.execute("CREATE INDEX IF NOT EXISTS idx_wm_type ON working_memory_items(item_type)")
            conn.execute("CREATE INDEX IF NOT EXISTS idx_wm_agent ON working_memory_items(source_agent)")
    
    def add_item(self, content: str, item_type: str, priority: float = 0.5,
                source_agent: str = "", context_tags: List[str] = None) -> str:
        """Add item to working memory with attention-based priority"""
        try:
            item_id = str(uuid.uuid4())
            
            item = WorkingMemoryItem(
                id=item_id,
                content=content,
                item_type=item_type,
                priority=priority,
                activation_level=priority,  # Initial activation equals priority
                created_at=datetime.now(),
                last_accessed=datetime.now(),
                access_count=0,
                decay_rate=0.1,  # Default decay rate
                context_tags=context_tags or [],
                source_agent=source_agent,
                related_items=[]
            )
            
            # Check capacity and evict if necessary
            if len(self._memory_items) >= self.capacity:
                self._evict_lowest_activation()
            
            # Store in memory and database
            self._memory_items[item_id] = item
            self._store_item_to_db(item)
            
            # Update attention focus if this is high priority
            if priority > 0.7 and (not self._attention_focus or 
                                  priority > self._attention_focus.attention_weight):
                self._update_attention_focus(item_id, item_type, priority)
            
            logger.info(f"Added working memory item: {item_type} (priority: {priority:.2f})")
            return item_id
            
        except Exception as e:
            logger.error(f"Error adding working memory item: {e}")
            return ""
    
    def get_item(self, item_id: str) -> Optional[WorkingMemoryItem]:
        """Retrieve item from working memory and update activation"""
        try:
            if item_id in self._memory_items:
                item = self._memory_items[item_id]
                
                # Update access statistics
                item.last_accessed = datetime.now()
                item.access_count += 1
                
                # Boost activation on access (but cap at 1.0)
                activation_boost = min(0.2, 1.0 - item.activation_level)
                item.activation_level = min(1.0, item.activation_level + activation_boost)
                
                # Update in database
                self._update_item_in_db(item)
                
                logger.debug(f"Retrieved working memory item: {item_id[:8]}...")
                return item
            
            logger.warning(f"Working memory item not found: {item_id}")
            return None
            
        except Exception as e:
            logger.error(f"Error retrieving working memory item: {e}")
            return None
    
    def get_active_items(self, min_activation: float = 0.3,
                        item_type: str = "", limit: int = 20) -> List[WorkingMemoryItem]:
        """Get currently active items above activation threshold"""
        try:
            active_items = []
            
            for item in self._memory_items.values():
                if (item.activation_level >= min_activation and
                    (not item_type or item.item_type == item_type)):
                    active_items.append(item)
            
            # Sort by activation level (highest first)
            active_items.sort(key=lambda x: x.activation_level, reverse=True)
            
            logger.info(f"Retrieved {len(active_items[:limit])} active items")
            return active_items[:limit]
            
        except Exception as e:
            logger.error(f"Error getting active items: {e}")
            return []
    
    def focus_attention(self, focus_type: str, item_ids: List[str],
                       attention_weight: float = 0.8, agent_id: str = "") -> str:
        """Focus attention on specific items"""
        try:
            focus_id = str(uuid.uuid4())
            
            # End current focus if exists
            if self._attention_focus:
                self._end_attention_focus()
            
            # Create new attention focus
            new_focus = AttentionFocus(
                focus_id=focus_id,
                focus_type=focus_type,
                focus_items=item_ids,
                attention_weight=attention_weight,
                duration=timedelta(0),
                created_at=datetime.now(),
                metadata={'agent_id': agent_id}
            )
            
            self._attention_focus = new_focus
            
            # Boost activation of focused items
            for item_id in item_ids:
                if item_id in self._memory_items:
                    item = self._memory_items[item_id]
                    item.activation_level = min(1.0, item.activation_level + 0.3)
                    self._update_item_in_db(item)
            
            # Store focus in database
            self._store_attention_focus(new_focus)
            
            logger.info(f"Focused attention on {focus_type}: {len(item_ids)} items")
            return focus_id
            
        except Exception as e:
            logger.error(f"Error focusing attention: {e}")
            return ""
    
    def switch_context(self, new_focus_type: str, new_item_ids: List[str],
                      switch_reason: str = "", agent_id: str = "") -> Dict[str, Any]:
        """Switch attention context with cost calculation"""
        try:
            switch_result = {
                'switch_id': str(uuid.uuid4()),
                'from_focus': None,
                'to_focus': new_focus_type,
                'switch_cost': 0.0,
                'success': False
            }
            
            # Calculate switch cost based on current focus
            if self._attention_focus:
                switch_result['from_focus'] = self._attention_focus.focus_type
                
                # Cost factors:
                # 1. How long we've been in current focus
                current_duration = datetime.now() - self._attention_focus.created_at
                duration_cost = min(current_duration.total_seconds() / 300.0, 0.3)  # Max 5min
                
                # 2. Number of active items being abandoned
                abandoned_items = len(self._attention_focus.focus_items)
                abandonment_cost = min(abandoned_items * 0.1, 0.4)
                
                # 3. Similarity between old and new focus
                similarity_discount = self._calculate_focus_similarity(
                    self._attention_focus.focus_items, new_item_ids
                )
                
                total_cost = duration_cost + abandonment_cost - similarity_discount
                switch_result['switch_cost'] = max(0.0, min(total_cost, 1.0))
                
                # Record context switch
                self._record_context_switch(
                    self._attention_focus.focus_type,
                    new_focus_type,
                    switch_reason,
                    switch_result['switch_cost'],
                    agent_id
                )
            
            # Perform the switch
            focus_id = self.focus_attention(new_focus_type, new_item_ids, agent_id=agent_id)
            switch_result['success'] = bool(focus_id)
            
            logger.info(f"Context switch: {switch_result['from_focus']} -> {new_focus_type} (cost: {switch_result['switch_cost']:.3f})")
            return switch_result
            
        except Exception as e:
            logger.error(f"Error switching context: {e}")
            return {'error': str(e)}
    
    def get_current_focus(self) -> Optional[AttentionFocus]:
        """Get current attention focus"""
        return self._attention_focus
    
    def decay_memory(self):
        """Apply decay to all working memory items"""
        try:
            decayed_count = 0
            evicted_items = []
            
            for item_id, item in list(self._memory_items.items()):
                # Apply decay based on time since last access
                time_since_access = datetime.now() - item.last_accessed
                decay_amount = item.decay_rate * (time_since_access.total_seconds() / 60.0)
                
                item.activation_level = max(0.0, item.activation_level - decay_amount)
                decayed_count += 1
                
                # Evict items with very low activation
                if item.activation_level < 0.05:
                    evicted_items.append(item_id)
                else:
                    # Update in database
                    self._update_item_in_db(item)
            
            # Remove evicted items
            for item_id in evicted_items:
                del self._memory_items[item_id]
                self._remove_item_from_db(item_id)
            
            if evicted_items:
                logger.info(f"Memory decay: {decayed_count} items decayed, {len(evicted_items)} evicted")
                
        except Exception as e:
            logger.error(f"Error during memory decay: {e}")
    
    def find_related_items(self, item_id: str, max_items: int = 5) -> List[WorkingMemoryItem]:
        """Find items related to the given item"""
        try:
            if item_id not in self._memory_items:
                return []
            
            source_item = self._memory_items[item_id]
            related_items = []
            
            for other_id, other_item in self._memory_items.items():
                if other_id == item_id:
                    continue
                
                # Calculate relatedness score
                relatedness = 0.0
                
                # Same type bonus
                if source_item.item_type == other_item.item_type:
                    relatedness += 0.3
                
                # Shared context tags
                shared_tags = set(source_item.context_tags) & set(other_item.context_tags)
                if shared_tags:
                    relatedness += len(shared_tags) * 0.2
                
                # Same source agent
                if source_item.source_agent == other_item.source_agent:
                    relatedness += 0.2
                
                # Temporal proximity
                time_diff = abs((source_item.created_at - other_item.created_at).total_seconds())
                if time_diff < 300:  # Within 5 minutes
                    relatedness += 0.3 * (300 - time_diff) / 300
                
                if relatedness > 0.1:  # Minimum relatedness threshold
                    related_items.append((other_item, relatedness))
            
            # Sort by relatedness and return top items
            related_items.sort(key=lambda x: x[1], reverse=True)
            
            return [item for item, score in related_items[:max_items]]
            
        except Exception as e:
            logger.error(f"Error finding related items: {e}")
            return []
    
    def _update_attention_focus(self, item_id: str, item_type: str, priority: float):
        """Update current attention focus"""
        if self._attention_focus:
            self._end_attention_focus()
        
        self.focus_attention(item_type, [item_id], priority)
    
    def _end_attention_focus(self):
        """End current attention focus"""
        if self._attention_focus:
            # Update duration
            self._attention_focus.duration = datetime.now() - self._attention_focus.created_at
            
            # Add to history
            self._attention_history.append(self._attention_focus)
            
            # Update in database
            with sqlite3.connect(self.db_path) as conn:
                conn.execute("""
                    UPDATE attention_focus_log SET
                        ended_at = ?,
                        duration_seconds = ?
                    WHERE id = ?
                """, (
                    datetime.now().isoformat(),
                    self._attention_focus.duration.total_seconds(),
                    self._attention_focus.focus_id
                ))
            
            self._attention_focus = None
    
    def _evict_lowest_activation(self):
        """Evict item with lowest activation to make space"""
        if not self._memory_items:
            return
        
        lowest_item_id = min(
            self._memory_items.keys(),
            key=lambda x: self._memory_items[x].activation_level
        )
        
        del self._memory_items[lowest_item_id]
        self._remove_item_from_db(lowest_item_id)
        
        logger.debug(f"Evicted working memory item: {lowest_item_id[:8]}...")
    
    def _calculate_focus_similarity(self, items1: List[str], items2: List[str]) -> float:
        """Calculate similarity between two sets of focus items"""
        if not items1 or not items2:
            return 0.0
        
        set1 = set(items1)
        set2 = set(items2)
        
        intersection = len(set1 & set2)
        union = len(set1 | set2)
        
        return intersection / union if union > 0 else 0.0
    
    def _record_context_switch(self, from_focus: str, to_focus: str,
                              reason: str, cost: float, agent_id: str):
        """Record context switch in database"""
        try:
            switch_id = str(uuid.uuid4())
            
            with sqlite3.connect(self.db_path) as conn:
                conn.execute("""
                    INSERT INTO context_switches (
                        id, from_focus, to_focus, switch_reason,
                        switch_cost, agent_id
                    ) VALUES (?, ?, ?, ?, ?, ?)
                """, (switch_id, from_focus, to_focus, reason, cost, agent_id))
                
        except Exception as e:
            logger.error(f"Error recording context switch: {e}")
    
    def _decay_loop(self):
        """Background thread for memory decay"""
        while self._decay_running:
            try:
                time.sleep(self.decay_interval)
                self.decay_memory()
            except Exception as e:
                logger.error(f"Error in decay loop: {e}")
    
    def _load_working_memory(self):
        """Load working memory items from database"""
        try:
            with sqlite3.connect(self.db_path) as conn:
                cursor = conn.execute("""
                    SELECT * FROM working_memory_items 
                    ORDER BY activation_level DESC 
                    LIMIT ?
                """, (self.capacity,))
                
                for row in cursor.fetchall():
                    item = WorkingMemoryItem(
                        id=row[0],
                        content=row[1],
                        item_type=row[2],
                        priority=row[3],
                        activation_level=row[4],
                        created_at=datetime.fromisoformat(row[5]),
                        last_accessed=datetime.fromisoformat(row[6]),
                        access_count=row[7],
                        decay_rate=row[8],
                        context_tags=json.loads(row[9]) if row[9] else [],
                        source_agent=row[10] or "",
                        related_items=json.loads(row[11]) if row[11] else []
                    )
                    self._memory_items[item.id] = item
                
                logger.info(f"Loaded {len(self._memory_items)} working memory items")
                
        except Exception as e:
            logger.error(f"Error loading working memory: {e}")
    
    def _store_item_to_db(self, item: WorkingMemoryItem):
        """Store item to database"""
        try:
            with sqlite3.connect(self.db_path) as conn:
                conn.execute("""
                    INSERT INTO working_memory_items (
                        id, content, item_type, priority, activation_level,
                        created_at, last_accessed, access_count, decay_rate,
                        context_tags, source_agent, related_items
                    ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
                """, (
                    item.id, item.content, item.item_type, item.priority,
                    item.activation_level, item.created_at.isoformat(),
                    item.last_accessed.isoformat(), item.access_count,
                    item.decay_rate, json.dumps(item.context_tags),
                    item.source_agent, json.dumps(item.related_items)
                ))
                
        except Exception as e:
            logger.error(f"Error storing item to database: {e}")
    
    def _update_item_in_db(self, item: WorkingMemoryItem):
        """Update item in database"""
        try:
            with sqlite3.connect(self.db_path) as conn:
                conn.execute("""
                    UPDATE working_memory_items SET
                        activation_level = ?, last_accessed = ?,
                        access_count = ?, updated_at = CURRENT_TIMESTAMP
                    WHERE id = ?
                """, (
                    item.activation_level, item.last_accessed.isoformat(),
                    item.access_count, item.id
                ))
                
        except Exception as e:
            logger.error(f"Error updating item in database: {e}")
    
    def _remove_item_from_db(self, item_id: str):
        """Remove item from database"""
        try:
            with sqlite3.connect(self.db_path) as conn:
                conn.execute("DELETE FROM working_memory_items WHERE id = ?", (item_id,))
                
        except Exception as e:
            logger.error(f"Error removing item from database: {e}")
    
    def _store_attention_focus(self, focus: AttentionFocus):
        """Store attention focus in database"""
        try:
            with sqlite3.connect(self.db_path) as conn:
                conn.execute("""
                    INSERT INTO attention_focus_log (
                        id, focus_type, focus_items, attention_weight,
                        created_at, metadata, agent_id
                    ) VALUES (?, ?, ?, ?, ?, ?, ?)
                """, (
                    focus.focus_id, focus.focus_type,
                    json.dumps(focus.focus_items), focus.attention_weight,
                    focus.created_at.isoformat(), json.dumps(focus.metadata),
                    focus.metadata.get('agent_id', '')
                ))
                
        except Exception as e:
            logger.error(f"Error storing attention focus: {e}")
    
    def get_working_memory_statistics(self) -> Dict[str, Any]:
        """Get comprehensive working memory statistics"""
        try:
            with sqlite3.connect(self.db_path) as conn:
                stats = {
                    'current_capacity': len(self._memory_items),
                    'max_capacity': self.capacity,
                    'utilization': len(self._memory_items) / self.capacity
                }
                
                # Activation distribution
                if self._memory_items:
                    activations = [item.activation_level for item in self._memory_items.values()]
                    stats['avg_activation'] = sum(activations) / len(activations)
                    stats['max_activation'] = max(activations)
                    stats['min_activation'] = min(activations)
                
                # Item type distribution
                type_counts = {}
                for item in self._memory_items.values():
                    type_counts[item.item_type] = type_counts.get(item.item_type, 0) + 1
                stats['item_types'] = type_counts
                
                # Context switch statistics
                cursor = conn.execute("""
                    SELECT COUNT(*), AVG(switch_cost) 
                    FROM context_switches 
                    WHERE timestamp > datetime('now', '-1 hour')
                """)
                row = cursor.fetchone()
                stats['recent_switches'] = row[0] or 0
                stats['avg_switch_cost'] = row[1] or 0.0
                
                # Current focus
                if self._attention_focus:
                    stats['current_focus'] = {
                        'type': self._attention_focus.focus_type,
                        'items': len(self._attention_focus.focus_items),
                        'weight': self._attention_focus.attention_weight,
                        'duration_seconds': (datetime.now() - self._attention_focus.created_at).total_seconds()
                    }
                else:
                    stats['current_focus'] = None
                
                return stats
                
        except Exception as e:
            logger.error(f"Error getting working memory statistics: {e}")
            return {'error': str(e)}
    
    def cleanup(self):
        """Cleanup resources"""
        self._decay_running = False
        if self._decay_thread.is_alive():
            self._decay_thread.join(timeout=1.0)

# Export the main classes
__all__ = ['WorkingMemoryManager', 'WorkingMemoryItem', 'AttentionFocus']