"""
Enhanced Hierarchical Memory Architecture for Toxo - Phase 2 Implementation

This module implements sophisticated memory management with:
- Multi-level memory hierarchy (Immediate, Working, Episodic, Semantic, Procedural, Meta)
- Intelligent memory consolidation and compression
- Cross-memory pattern recognition and transfer
- Temporal decay modeling and importance-based retention
- Advanced retrieval with associative activation
"""

import asyncio
import numpy as np
import sqlite3
import json
import pickle
from typing import Dict, List, Any, Optional, Tuple, Union, Set
from dataclasses import dataclass, field, asdict
from datetime import datetime, timedelta
from enum import Enum
from pathlib import Path
from collections import defaultdict, deque
import hashlib
import logging

# ML libraries
from sklearn.cluster import DBSCAN, KMeans
from sklearn.metrics.pairwise import cosine_similarity
from sklearn.preprocessing import StandardScaler
from sklearn.decomposition import PCA
import networkx as nx

# Toxo imports
from ..utils.logger import get_logger
from ..utils.exceptions import MemoryError, ValidationError
from ..core.config import ToxoConfig


class MemoryLevel(Enum):
    """Memory hierarchy levels."""
    IMMEDIATE = "immediate"      # Real-time context
    WORKING = "working"          # Session context
    EPISODIC = "episodic"        # Experience memory
    SEMANTIC = "semantic"        # Knowledge memory
    PROCEDURAL = "procedural"    # Skill memory
    META = "meta"                # Learning memory


class MemoryType(Enum):
    """Types of memory content."""
    INTERACTION = "interaction"
    KNOWLEDGE = "knowledge"
    PATTERN = "pattern"
    SKILL = "skill"
    FEEDBACK = "feedback"
    CONTEXT = "context"
    GOAL = "goal"
    STRATEGY = "strategy"


@dataclass
class MemoryItem:
    """Enhanced memory item with hierarchical metadata."""
    memory_id: str
    content: str
    memory_level: MemoryLevel
    memory_type: MemoryType
    embedding: Optional[np.ndarray] = None
    importance: float = 0.5
    access_count: int = 0
    last_accessed: Optional[datetime] = None
    created_at: datetime = field(default_factory=datetime.now)
    decay_rate: float = 0.01
    associations: Set[str] = field(default_factory=set)
    metadata: Dict[str, Any] = field(default_factory=dict)
    tags: Set[str] = field(default_factory=set)


@dataclass
class MemoryAssociation:
    """Association between memory items."""
    source_id: str
    target_id: str
    association_type: str
    strength: float
    created_at: datetime = field(default_factory=datetime.now)
    metadata: Dict[str, Any] = field(default_factory=dict)


@dataclass
class ConsolidationRule:
    """Rule for memory consolidation."""
    rule_id: str
    source_level: MemoryLevel
    target_level: MemoryLevel
    conditions: Dict[str, Any]
    consolidation_strategy: str
    priority: int = 1


class MemoryConsolidator:
    """Handles memory consolidation between levels."""
    
    def __init__(self):
        self.logger = get_logger(f"{__name__}.Consolidator")
        self.consolidation_rules: List[ConsolidationRule] = []
        self._setup_default_rules()
    
    def _setup_default_rules(self):
        """Setup default consolidation rules."""
        # Immediate -> Working
        self.consolidation_rules.append(ConsolidationRule(
            rule_id="immediate_to_working",
            source_level=MemoryLevel.IMMEDIATE,
            target_level=MemoryLevel.WORKING,
            conditions={"min_importance": 0.3, "max_age_minutes": 5},
            consolidation_strategy="importance_threshold"
        ))
        
        # Working -> Episodic
        self.consolidation_rules.append(ConsolidationRule(
            rule_id="working_to_episodic",
            source_level=MemoryLevel.WORKING,
            target_level=MemoryLevel.EPISODIC,
            conditions={"min_importance": 0.5, "min_access_count": 2},
            consolidation_strategy="pattern_extraction"
        ))
        
        # Episodic -> Semantic
        self.consolidation_rules.append(ConsolidationRule(
            rule_id="episodic_to_semantic",
            source_level=MemoryLevel.EPISODIC,
            target_level=MemoryLevel.SEMANTIC,
            conditions={"min_importance": 0.7, "pattern_frequency": 3},
            consolidation_strategy="abstraction"
        ))
    
    async def consolidate_memories(
        self,
        memories: Dict[MemoryLevel, List[MemoryItem]]
    ) -> Dict[str, Any]:
        """Perform memory consolidation across levels."""
        consolidation_results = {}
        
        for rule in sorted(self.consolidation_rules, key=lambda r: r.priority, reverse=True):
            if rule.source_level not in memories:
                continue
            
            source_memories = memories[rule.source_level]
            candidates = await self._find_consolidation_candidates(source_memories, rule)
            
            if candidates:
                consolidated = await self._apply_consolidation_strategy(candidates, rule)
                consolidation_results[rule.rule_id] = consolidated
        
        return consolidation_results
    
    async def _find_consolidation_candidates(
        self,
        memories: List[MemoryItem],
        rule: ConsolidationRule
    ) -> List[MemoryItem]:
        """Find memories that meet consolidation criteria."""
        candidates = []
        
        for memory in memories:
            if await self._meets_conditions(memory, rule.conditions):
                candidates.append(memory)
        
        return candidates
    
    async def _meets_conditions(self, memory: MemoryItem, conditions: Dict[str, Any]) -> bool:
        """Check if memory meets consolidation conditions."""
        # Check importance threshold
        if "min_importance" in conditions:
            if memory.importance < conditions["min_importance"]:
                return False
        
        # Check age
        if "max_age_minutes" in conditions:
            age_minutes = (datetime.now() - memory.created_at).total_seconds() / 60
            if age_minutes > conditions["max_age_minutes"]:
                return False
        
        # Check access count
        if "min_access_count" in conditions:
            if memory.access_count < conditions["min_access_count"]:
                return False
        
        return True
    
    async def _apply_consolidation_strategy(
        self,
        candidates: List[MemoryItem],
        rule: ConsolidationRule
    ) -> Dict[str, Any]:
        """Apply consolidation strategy to candidates."""
        if rule.consolidation_strategy == "importance_threshold":
            return await self._consolidate_by_importance(candidates, rule)
        elif rule.consolidation_strategy == "pattern_extraction":
            return await self._consolidate_by_patterns(candidates, rule)
        elif rule.consolidation_strategy == "abstraction":
            return await self._consolidate_by_abstraction(candidates, rule)
        else:
            return {"strategy": "unknown", "consolidated": 0}
    
    async def _consolidate_by_importance(
        self,
        candidates: List[MemoryItem],
        rule: ConsolidationRule
    ) -> Dict[str, Any]:
        """Consolidate memories based on importance."""
        consolidated_count = 0
        
        for memory in candidates:
            # Move to target level
            memory.memory_level = rule.target_level
            memory.importance *= 1.1  # Boost importance for surviving consolidation
            consolidated_count += 1
        
        return {
            "strategy": "importance_threshold",
            "consolidated": consolidated_count,
            "target_level": rule.target_level.value
        }
    
    async def _consolidate_by_patterns(
        self,
        candidates: List[MemoryItem],
        rule: ConsolidationRule
    ) -> Dict[str, Any]:
        """Consolidate memories by extracting patterns."""
        # Group similar memories
        if not candidates:
            return {"strategy": "pattern_extraction", "consolidated": 0}
        
        # Simple clustering based on content similarity
        embeddings = [mem.embedding for mem in candidates if mem.embedding is not None]
        
        if len(embeddings) < 2:
            return {"strategy": "pattern_extraction", "consolidated": 0}
        
        # Cluster memories
        embeddings_array = np.array(embeddings)
        clustering = DBSCAN(eps=0.3, min_samples=2).fit(embeddings_array)
        
        patterns_extracted = 0
        for cluster_id in set(clustering.labels_):
            if cluster_id == -1:  # Noise
                continue
            
            cluster_indices = np.where(clustering.labels_ == cluster_id)[0]
            cluster_memories = [candidates[i] for i in cluster_indices]
            
            # Create pattern memory
            pattern_content = await self._extract_pattern_from_cluster(cluster_memories)
            patterns_extracted += 1
        
        return {
            "strategy": "pattern_extraction",
            "consolidated": len(candidates),
            "patterns_extracted": patterns_extracted,
            "target_level": rule.target_level.value
        }
    
    async def _consolidate_by_abstraction(
        self,
        candidates: List[MemoryItem],
        rule: ConsolidationRule
    ) -> Dict[str, Any]:
        """Consolidate memories through abstraction."""
        # Create abstract concepts from specific memories
        abstractions_created = 0
        
        # Group by memory type
        type_groups = defaultdict(list)
        for memory in candidates:
            type_groups[memory.memory_type].append(memory)
        
        for memory_type, memories in type_groups.items():
            if len(memories) >= 3:  # Need multiple instances for abstraction
                abstract_concept = await self._create_abstract_concept(memories)
                abstractions_created += 1
        
        return {
            "strategy": "abstraction",
            "consolidated": len(candidates),
            "abstractions_created": abstractions_created,
            "target_level": rule.target_level.value
        }
    
    async def _extract_pattern_from_cluster(self, memories: List[MemoryItem]) -> str:
        """Extract common pattern from clustered memories."""
        # Simple pattern extraction - find common words
        all_words = []
        for memory in memories:
            all_words.extend(memory.content.lower().split())
        
        word_counts = defaultdict(int)
        for word in all_words:
            word_counts[word] += 1
        
        # Find most common words
        common_words = sorted(word_counts.items(), key=lambda x: x[1], reverse=True)[:10]
        pattern = " ".join([word for word, count in common_words if count > 1])
        
        return f"Pattern: {pattern}"
    
    async def _create_abstract_concept(self, memories: List[MemoryItem]) -> str:
        """Create abstract concept from specific memories."""
        # Extract key concepts and create abstraction
        concepts = []
        for memory in memories:
            # Simple concept extraction
            words = memory.content.split()
            concepts.extend([word for word in words if len(word) > 4])
        
        unique_concepts = list(set(concepts))[:5]
        return f"Abstract concept: {' + '.join(unique_concepts)}"


class AssociativeRetriever:
    """Handles associative memory retrieval."""
    
    def __init__(self):
        self.logger = get_logger(f"{__name__}.Retriever")
        self.association_graph = nx.Graph()
    
    def add_association(self, association: MemoryAssociation):
        """Add association to the graph."""
        self.association_graph.add_edge(
            association.source_id,
            association.target_id,
            weight=association.strength,
            type=association.association_type,
            created_at=association.created_at
        )
    
    async def retrieve_associated_memories(
        self,
        memory_id: str,
        max_depth: int = 3,
        min_strength: float = 0.3
    ) -> List[Tuple[str, float]]:
        """Retrieve memories associated with given memory."""
        if memory_id not in self.association_graph:
            return []
        
        associated = []
        
        # Use breadth-first search with depth limit
        for target_id in nx.single_source_shortest_path_length(
            self.association_graph, memory_id, cutoff=max_depth
        ):
            if target_id == memory_id:
                continue
            
            try:
                path = nx.shortest_path(self.association_graph, memory_id, target_id)
                # Calculate path strength
                path_strength = 1.0
                for i in range(len(path) - 1):
                    edge_data = self.association_graph[path[i]][path[i + 1]]
                    path_strength *= edge_data.get('weight', 0.5)
                
                if path_strength >= min_strength:
                    associated.append((target_id, path_strength))
            except nx.NetworkXNoPath:
                continue
        
        # Sort by strength
        associated.sort(key=lambda x: x[1], reverse=True)
        return associated


class HierarchicalMemorySystem:
    """
    Enhanced hierarchical memory system with multi-level architecture.
    """
    
    def __init__(self, config: Optional[ToxoConfig] = None):
        self.config = config or ToxoConfig()
        self.logger = get_logger(__name__)
        
        # Memory storage by level
        self.memories: Dict[MemoryLevel, List[MemoryItem]] = {
            level: [] for level in MemoryLevel
        }
        
        # Memory associations
        self.associations: List[MemoryAssociation] = []
        
        # Components
        self.consolidator = MemoryConsolidator()
        self.retriever = AssociativeRetriever()
        
        # Memory management
        self.memory_limits = {
            MemoryLevel.IMMEDIATE: 50,
            MemoryLevel.WORKING: 200,
            MemoryLevel.EPISODIC: 1000,
            MemoryLevel.SEMANTIC: 5000,
            MemoryLevel.PROCEDURAL: 2000,
            MemoryLevel.META: 1000
        }
        
        # Performance tracking
        self.access_patterns: Dict[str, List[datetime]] = defaultdict(list)
        self.consolidation_history: List[Dict[str, Any]] = []
        
        # Database for persistence
        self.db_path: Optional[Path] = None
        self.connection: Optional[sqlite3.Connection] = None
        
        self.logger.info("Hierarchical Memory System initialized")
    
    async def store_memory(
        self,
        content: str,
        memory_level: MemoryLevel,
        memory_type: MemoryType,
        importance: float = 0.5,
        metadata: Optional[Dict[str, Any]] = None,
        tags: Optional[Set[str]] = None,
        embedding: Optional[np.ndarray] = None
    ) -> str:
        """
        Store memory item in specified level.
        
        Args:
            content: Memory content
            memory_level: Target memory level
            memory_type: Type of memory
            importance: Importance score (0-1)
            metadata: Additional metadata
            tags: Memory tags
            embedding: Content embedding
            
        Returns:
            Memory ID
        """
        try:
            # Generate memory ID
            memory_id = hashlib.md5(
                f"{content}_{memory_level.value}_{datetime.now().isoformat()}".encode()
            ).hexdigest()
            
            # Create memory item
            memory_item = MemoryItem(
                memory_id=memory_id,
                content=content,
                memory_level=memory_level,
                memory_type=memory_type,
                embedding=embedding,
                importance=importance,
                metadata=metadata or {},
                tags=tags or set()
            )
            
            # Store in appropriate level
            self.memories[memory_level].append(memory_item)
            
            # Manage memory limits
            await self._manage_memory_limits(memory_level)
            
            # Create associations
            await self._create_automatic_associations(memory_item)
            
            # Trigger consolidation if needed
            if await self._should_trigger_consolidation(memory_level):
                await self.consolidate_memories()
            
            self.logger.debug(f"Stored memory: {memory_id} in {memory_level.value}")
            return memory_id
            
        except Exception as e:
            self.logger.error(f"Error storing memory: {str(e)}")
            raise MemoryError(f"Failed to store memory: {str(e)}")
    
    async def retrieve_memories(
        self,
        query: str,
        memory_levels: Optional[List[MemoryLevel]] = None,
        memory_types: Optional[List[MemoryType]] = None,
        limit: int = 10,
        similarity_threshold: float = 0.1,
        use_associations: bool = True,
        query_embedding: Optional[np.ndarray] = None
    ) -> List[Tuple[MemoryItem, float]]:
        """
        Retrieve memories using multi-level search with associations.
        
        Args:
            query: Search query
            memory_levels: Levels to search (all if None)
            memory_types: Types to search (all if None)
            limit: Maximum results
            similarity_threshold: Minimum similarity
            use_associations: Whether to use associative retrieval
            query_embedding: Pre-computed query embedding
            
        Returns:
            List of (memory_item, similarity_score) tuples
        """
        try:
            search_levels = memory_levels or list(MemoryLevel)
            search_types = memory_types or list(MemoryType)
            
            all_candidates = []
            
            # Search each level
            for level in search_levels:
                level_memories = self.memories[level]
                
                for memory in level_memories:
                    # Filter by type
                    if memory.memory_type not in search_types:
                        continue
                    
                    # Calculate similarity
                    similarity = await self._calculate_similarity(
                        query, memory, query_embedding
                    )
                    
                    if similarity >= similarity_threshold:
                        # Update access tracking
                        await self._update_access_tracking(memory)
                        
                        all_candidates.append((memory, similarity))
            
            # Sort by similarity and importance
            all_candidates.sort(
                key=lambda x: (x[1] * x[0].importance), reverse=True
            )
            
            # Get top candidates
            top_candidates = all_candidates[:limit]
            
            # Add associative memories if requested
            if use_associations and top_candidates:
                associative_memories = await self._get_associative_memories(
                    [mem for mem, _ in top_candidates[:3]], limit // 2
                )
                top_candidates.extend(associative_memories)
                
                # Re-sort and limit
                top_candidates.sort(
                    key=lambda x: (x[1] * x[0].importance), reverse=True
                )
                top_candidates = top_candidates[:limit]
            
            self.logger.debug(f"Retrieved {len(top_candidates)} memories for query: {query[:50]}")
            return top_candidates
            
        except Exception as e:
            self.logger.error(f"Error retrieving memories: {str(e)}")
            raise MemoryError(f"Failed to retrieve memories: {str(e)}")
    
    async def consolidate_memories(self) -> Dict[str, Any]:
        """Perform memory consolidation across all levels."""
        try:
            self.logger.info("Starting memory consolidation")
            
            consolidation_results = await self.consolidator.consolidate_memories(self.memories)
            
            # Update consolidation history
            self.consolidation_history.append({
                'timestamp': datetime.now().isoformat(),
                'results': consolidation_results
            })
            
            # Clean up after consolidation
            await self._cleanup_after_consolidation()
            
            self.logger.info(f"Memory consolidation completed: {consolidation_results}")
            return consolidation_results
            
        except Exception as e:
            self.logger.error(f"Error during consolidation: {str(e)}")
            raise MemoryError(f"Failed to consolidate memories: {str(e)}")
    
    async def create_association(
        self,
        source_memory_id: str,
        target_memory_id: str,
        association_type: str,
        strength: float = 0.5,
        metadata: Optional[Dict[str, Any]] = None
    ) -> None:
        """Create association between memories."""
        association = MemoryAssociation(
            source_id=source_memory_id,
            target_id=target_memory_id,
            association_type=association_type,
            strength=strength,
            metadata=metadata or {}
        )
        
        self.associations.append(association)
        self.retriever.add_association(association)
        
        self.logger.debug(f"Created association: {source_memory_id} -> {target_memory_id}")
    
    async def get_memory_statistics(self) -> Dict[str, Any]:
        """Get comprehensive memory statistics."""
        stats = {
            'memory_counts': {},
            'total_memories': 0,
            'total_associations': len(self.associations),
            'consolidation_count': len(self.consolidation_history),
            'access_patterns': {},
            'importance_distribution': {},
            'memory_age_distribution': {}
        }
        
        # Count memories by level
        for level, memories in self.memories.items():
            count = len(memories)
            stats['memory_counts'][level.value] = count
            stats['total_memories'] += count
            
            if memories:
                # Importance distribution
                importances = [mem.importance for mem in memories]
                stats['importance_distribution'][level.value] = {
                    'mean': np.mean(importances),
                    'std': np.std(importances),
                    'min': np.min(importances),
                    'max': np.max(importances)
                }
                
                # Age distribution
                ages = [(datetime.now() - mem.created_at).total_seconds() / 3600 for mem in memories]
                stats['memory_age_distribution'][level.value] = {
                    'mean_hours': np.mean(ages),
                    'max_hours': np.max(ages),
                    'min_hours': np.min(ages)
                }
        
        # Access patterns
        for memory_id, accesses in self.access_patterns.items():
            if len(accesses) > 1:
                stats['access_patterns'][memory_id] = {
                    'access_count': len(accesses),
                    'last_access': accesses[-1].isoformat(),
                    'access_frequency': len(accesses) / max(1, (datetime.now() - accesses[0]).total_seconds() / 3600)
                }
        
        return stats
    
    async def optimize_memory_structure(self) -> Dict[str, Any]:
        """Optimize memory structure for better performance."""
        optimization_results = {
            'memories_pruned': 0,
            'associations_strengthened': 0,
            'importance_updates': 0
        }
        
        # Prune low-importance, old memories
        for level, memories in self.memories.items():
            if level in [MemoryLevel.IMMEDIATE, MemoryLevel.WORKING]:
                continue  # Don't prune short-term memories
            
            pruned = []
            for memory in memories[:]:
                age_hours = (datetime.now() - memory.created_at).total_seconds() / 3600
                
                # Prune if very old and low importance
                if age_hours > 168 and memory.importance < 0.2 and memory.access_count < 2:
                    memories.remove(memory)
                    pruned.append(memory.memory_id)
                    optimization_results['memories_pruned'] += 1
        
        # Strengthen frequently accessed associations
        for association in self.associations:
            source_accesses = len(self.access_patterns.get(association.source_id, []))
            target_accesses = len(self.access_patterns.get(association.target_id, []))
            
            if source_accesses > 5 and target_accesses > 5:
                association.strength = min(1.0, association.strength * 1.1)
                optimization_results['associations_strengthened'] += 1
        
        # Update importance based on access patterns
        for level, memories in self.memories.items():
            for memory in memories:
                access_count = len(self.access_patterns.get(memory.memory_id, []))
                if access_count > memory.access_count:
                    memory.importance = min(1.0, memory.importance * (1 + 0.1 * access_count))
                    memory.access_count = access_count
                    optimization_results['importance_updates'] += 1
        
        self.logger.info(f"Memory optimization completed: {optimization_results}")
        return optimization_results
    
    async def _calculate_similarity(
        self,
        query: str,
        memory: MemoryItem,
        query_embedding: Optional[np.ndarray] = None
    ) -> float:
        """Calculate similarity between query and memory."""
        # Improved text-based similarity
        query_lower = query.lower()
        memory_lower = memory.content.lower()
        
        # Check for direct substring matches first
        if query_lower in memory_lower or memory_lower in query_lower:
            return 0.8
        
        # Word-based similarity
        query_words = set(query_lower.split())
        memory_words = set(memory_lower.split())
        
        if not query_words or not memory_words:
            return 0.0
        
        # Calculate multiple similarity metrics
        intersection = query_words & memory_words
        union = query_words | memory_words
        jaccard_similarity = len(intersection) / len(union) if union else 0.0
        
        # Common words ratio (more generous than Jaccard)
        common_ratio = len(intersection) / min(len(query_words), len(memory_words)) if query_words and memory_words else 0.0
        
        # Keyword boost for important terms
        important_keywords = {"machine", "learning", "algorithm", "model", "data", "ai", "artificial", "intelligence"}
        keyword_matches = len(query_words & memory_words & important_keywords)
        keyword_boost = keyword_matches * 0.2
        
        # Tag matching boost
        query_tags = set(query_lower.split())
        tag_boost = len(query_tags & memory.tags) * 0.2
        
        # Use the maximum of different similarity measures for better recall
        similarity = max(jaccard_similarity, common_ratio * 0.6) + keyword_boost + tag_boost
        
        return min(1.0, similarity)
    
    async def _update_access_tracking(self, memory: MemoryItem) -> None:
        """Update access tracking for memory."""
        memory.access_count += 1
        memory.last_accessed = datetime.now()
        self.access_patterns[memory.memory_id].append(datetime.now())
    
    async def _get_associative_memories(
        self,
        seed_memories: List[MemoryItem],
        limit: int
    ) -> List[Tuple[MemoryItem, float]]:
        """Get memories associated with seed memories."""
        associative_results = []
        
        for seed_memory in seed_memories:
            associated_ids = await self.retriever.retrieve_associated_memories(
                seed_memory.memory_id, max_depth=2, min_strength=0.3
            )
            
            for memory_id, strength in associated_ids[:limit]:
                # Find the actual memory
                found_memory = await self._find_memory_by_id(memory_id)
                if found_memory:
                    associative_results.append((found_memory, strength * 0.8))  # Reduce score for associative
        
        return associative_results[:limit]
    
    async def _find_memory_by_id(self, memory_id: str) -> Optional[MemoryItem]:
        """Find memory by ID across all levels."""
        for level_memories in self.memories.values():
            for memory in level_memories:
                if memory.memory_id == memory_id:
                    return memory
        return None
    
    async def _manage_memory_limits(self, level: MemoryLevel) -> None:
        """Manage memory limits for a specific level."""
        memories = self.memories[level]
        limit = self.memory_limits[level]
        
        if len(memories) > limit:
            # Sort by importance and recency
            memories.sort(key=lambda m: (m.importance, m.last_accessed or m.created_at), reverse=True)
            
            # Keep only the top memories
            self.memories[level] = memories[:limit]
            
            self.logger.debug(f"Pruned {len(memories) - limit} memories from {level.value}")
    
    async def _create_automatic_associations(self, new_memory: MemoryItem) -> None:
        """Create automatic associations for new memory."""
        # Find similar memories for association
        similar_memories = []
        
        for level_memories in self.memories.values():
            for memory in level_memories:
                if memory.memory_id == new_memory.memory_id:
                    continue
                
                similarity = await self._calculate_similarity(
                    new_memory.content, memory
                )
                
                if similarity > 0.6:
                    similar_memories.append((memory, similarity))
        
        # Create associations with most similar memories
        similar_memories.sort(key=lambda x: x[1], reverse=True)
        
        for memory, similarity in similar_memories[:5]:
            await self.create_association(
                new_memory.memory_id,
                memory.memory_id,
                "similarity",
                similarity
            )
    
    async def _should_trigger_consolidation(self, level: MemoryLevel) -> bool:
        """Determine if consolidation should be triggered."""
        memories = self.memories[level]
        
        # Trigger based on memory count
        if level == MemoryLevel.IMMEDIATE and len(memories) > 30:
            return True
        elif level == MemoryLevel.WORKING and len(memories) > 150:
            return True
        
        # Trigger based on time since last consolidation
        if self.consolidation_history:
            last_consolidation = datetime.fromisoformat(self.consolidation_history[-1]['timestamp'])
            hours_since = (datetime.now() - last_consolidation).total_seconds() / 3600
            
            if hours_since > 24:  # Daily consolidation
                return True
        
        return False
    
    async def _cleanup_after_consolidation(self) -> None:
        """Clean up after consolidation process."""
        # Remove duplicate memories
        for level, memories in self.memories.items():
            seen_content = set()
            unique_memories = []
            
            for memory in memories:
                content_hash = hashlib.md5(memory.content.encode()).hexdigest()
                if content_hash not in seen_content:
                    seen_content.add(content_hash)
                    unique_memories.append(memory)
            
            self.memories[level] = unique_memories
    
    def save_state(self, path: Path) -> None:
        """Save memory system state."""
        state = {
            'memories': {
                level.value: [asdict(mem) for mem in memories]
                for level, memories in self.memories.items()
            },
            'associations': [asdict(assoc) for assoc in self.associations],
            'consolidation_history': self.consolidation_history,
            'access_patterns': {
                memory_id: [dt.isoformat() for dt in accesses]
                for memory_id, accesses in self.access_patterns.items()
            }
        }
        
        with open(path, 'w') as f:
            json.dump(state, f, indent=2, default=str)
        
        self.logger.info(f"Memory system state saved to {path}")
    
    def load_state(self, path: Path) -> None:
        """Load memory system state."""
        with open(path, 'r') as f:
            state = json.load(f)
        
        # Restore memories
        self.memories = {}
        for level_str, memory_dicts in state['memories'].items():
            level = MemoryLevel(level_str)
            self.memories[level] = []
            
            for mem_dict in memory_dicts:
                memory = MemoryItem(
                    memory_id=mem_dict['memory_id'],
                    content=mem_dict['content'],
                    memory_level=MemoryLevel(mem_dict['memory_level']),
                    memory_type=MemoryType(mem_dict['memory_type']),
                    embedding=np.array(mem_dict['embedding']) if mem_dict.get('embedding') else None,
                    importance=mem_dict['importance'],
                    access_count=mem_dict['access_count'],
                    last_accessed=datetime.fromisoformat(mem_dict['last_accessed']) if mem_dict.get('last_accessed') else None,
                    created_at=datetime.fromisoformat(mem_dict['created_at']),
                    decay_rate=mem_dict.get('decay_rate', 0.01),
                    associations=set(mem_dict.get('associations', [])),
                    metadata=mem_dict.get('metadata', {}),
                    tags=set(mem_dict.get('tags', []))
                )
                self.memories[level].append(memory)
        
        # Restore associations
        self.associations = [
            MemoryAssociation(
                source_id=assoc['source_id'],
                target_id=assoc['target_id'],
                association_type=assoc['association_type'],
                strength=assoc['strength'],
                created_at=datetime.fromisoformat(assoc['created_at']),
                metadata=assoc.get('metadata', {})
            ) for assoc in state.get('associations', [])
        ]
        
        # Rebuild association graph
        for association in self.associations:
            self.retriever.add_association(association)
        
        # Restore other state
        self.consolidation_history = state.get('consolidation_history', [])
        self.access_patterns = {
            memory_id: [datetime.fromisoformat(dt) for dt in accesses]
            for memory_id, accesses in state.get('access_patterns', {}).items()
        }
        
        self.logger.info(f"Memory system state loaded from {path}")
    
    def get_state(self) -> Dict[str, Any]:
        """Get the current state for serialization."""
        return {
            "config_available": self.config is not None,
            "total_memories": sum(len(memories) for memories in self.memories.values()),
            "memories_by_level": {level.value: len(memories) for level, memories in self.memories.items()},
            "associations_count": len(getattr(self, 'associations', [])),
            "system_type": "HierarchicalMemorySystem"
        }


# Factory function
def create_hierarchical_memory_system(config: Optional[ToxoConfig] = None) -> HierarchicalMemorySystem:
    """Create and return a Hierarchical Memory System instance."""
    return HierarchicalMemorySystem(config) 