memory/memory_graph.py

"""
Memory Graph Module
Graph-based memory storage system with compression, retrieval, and RAG integration.
"""
import os
import json
import gzip
import base64
import time
import torch
import torch.nn.functional as F
import networkx as nx
import numpy as np
import re
from collections import deque
from typing import List, Optional, Any
from networkx.readwrite import json_graph
from sklearn.decomposition import PCA
try:
    import config as _config_module  # optional top-level config module
except ModuleNotFoundError:
    _config_module = None  # graceful — MemoryGraph works without it
try:
    from rag.rag_generator import RAGGenerator
except ImportError:
    RAGGenerator = None
import sys

# Try to import shared model for encoder
SHARED_MODEL_AVAILABLE = False
get_shared_model = None
get_shared_tokenizer = None
get_shared_adapter = None
try:
    from shared_model import get_shared_adapter, get_shared_model, get_shared_tokenizer
    SHARED_MODEL_AVAILABLE = True
except Exception:
    pass

class MemoryGraph:
    __slots__ = (
        'device', 'graph_file', 'max_nodes', 'save_interval', 'similarity_threshold',
        'compression_level', 'quantization_bits', 'pca', 'original_dim',
        'graph', 'id_counter', 'addition_count', 'node_queue',
        'access_counts', 'last_accessed', 'importance_scores', 'encoder', 'rag_generator',
        'use_learnable_lru', 'importance_decay', 'access_boost', 'connectivity_weight',
        'vae_compressor', 'use_vae_compression', 'vae_latent_dim', 'memory_pressure_threshold',
        'node_types', 'temporal_connections', 'consolidation_buffer', 'forgetting_curves',
        'use_combined_compression', 'pca_on_vae', 'memory_hierarchies',
        'enable_memory_types', 'enable_temporal_connections', 'enable_consolidation',
        'shared_model', 'shared_tokenizer', 'enable_vae_training', 'vae_learning_rate', 'vae_weights_path',
        'embedding_cache', 'cache_hits', 'cache_misses', 'cache_max_size',
        'forgetation_model', 'forgetation_history', 'forgetation_learning_rate', 'forgetation_buffer', 'forgetation_optimizer',
        'forgetation_shared_model', 'forgetation_shared_inner', 'forgetation_shared_dim', 'shared_adapter',
        'last_save_time', 'save_interval_seconds', 'min_save_interval'
    )

    def __init__(
        self,
        device: str = "cuda",
        graph_file: str = 'memory_graph.json.gz',
        max_nodes: Optional[int] = None,  # None = infinite growth
        save_interval: int = 100,
        save_interval_seconds: int = 120,
        min_save_interval: int = 10,
        similarity_threshold: float = 0.75,
        compression_level: float = 0.5,
        quantization_bits: int = 8,
        rag_generator_instance: Optional[RAGGenerator] = None,
        use_vae_compression: bool = True,  # Use VAE + PCA combined compression
        vae_latent_dim: int = 128,  # VAE compression target (512 -> 128 = 4x compression)
        memory_pressure_threshold: float = 0.85,  # Only prune when GPU memory > 85%
        use_combined_compression: bool = True,  # Use VAE -> PCA pipeline (no fallbacks)
        enable_memory_types: bool = True,  # Enable different memory node types
        enable_temporal_connections: bool = True,  # Enable temporal memory connections
        enable_consolidation: bool = True,  # Enable memory consolidation
        shared_model: Optional[Any] = None,  # Shared model for memory efficiency
        shared_tokenizer: Optional[Any] = None,  # Shared tokenizer
        enable_vae_training: bool = True,  # Enable VAE learning over time
        vae_learning_rate: float = 1e-4  # VAE learning rate
    ):
        # Core setup - FORCE CUDA if available
        if torch.cuda.is_available() and device != 'cpu':
            self.device = torch.device("cuda")
        else:
            self.device = torch.device(device if device == 'cpu' else 'cpu')

        graph_file = self._resolve_graph_path(graph_file)
        # Normalize graph_file to always use .gz extension for compressed storage
        if graph_file.endswith('.gz'):
            self.graph_file = graph_file
        elif graph_file.endswith('.json'):
            # Convert .json to .json.gz
            self.graph_file = graph_file + '.gz'
            # Migrate old uncompressed file if it exists
            if os.path.exists(graph_file) and not os.path.exists(self.graph_file):
                print(f"[MIGRATE] Found uncompressed {graph_file}, will migrate to compressed format on next save")
        else:
            # Add .gz extension if not present
            self.graph_file = graph_file + '.gz'
        self.max_nodes = max_nodes  # None = infinite growth
        self.save_interval = save_interval
        self.save_interval_seconds = int(save_interval_seconds)
        self.min_save_interval = max(1, int(min_save_interval))
        self.last_save_time = time.time()
        self.similarity_threshold = similarity_threshold
        self.compression_level = compression_level
        self.quantization_bits = quantization_bits
        self.use_vae_compression = use_vae_compression
        self.vae_latent_dim = vae_latent_dim
        self.memory_pressure_threshold = memory_pressure_threshold
        self.use_combined_compression = use_combined_compression  # VAE -> PCA pipeline
        self.enable_memory_types = enable_memory_types
        self.enable_temporal_connections = enable_temporal_connections
        self.enable_consolidation = enable_consolidation
        
        # Shared model attributes for memory efficiency
        self.shared_model = shared_model
        self.shared_tokenizer = shared_tokenizer
        self.shared_adapter = None
        if SHARED_MODEL_AVAILABLE and get_shared_adapter is not None:
            try:
                self.shared_adapter = get_shared_adapter()
            except Exception:
                self.shared_adapter = None
        if self.shared_adapter is not None:
            self.shared_model = self.shared_adapter.get_model()
            self.shared_tokenizer = self.shared_adapter.get_tokenizer()
        if self.shared_model is None and SHARED_MODEL_AVAILABLE and get_shared_model is not None:
            try:
                self.shared_model = get_shared_model()
                if self.shared_model is not None:
                    print("[MEMORY] Using shared model for memory efficiency (zero overhead)")
            except Exception as e:
                print(f"[WARN] Could not get shared model: {e}")
        
        # VAE training settings
        self.enable_vae_training = enable_vae_training
        self.vae_learning_rate = vae_learning_rate

        # Internal state
        self.pca = None
        self.pca_on_vae = None  # PCA applied to VAE latent space
        self.original_dim = None
        self.vae_compressor = None  # VAE compressor for better compression
        self.vae_weights_path = self.graph_file.replace('.json.gz', '_vae_weights.pt') if self.graph_file.endswith('.gz') else self.graph_file.replace('.json', '_vae_weights.pt')  # Path to save/load VAE weights
        
        # ⚡ INFINITE INTELLIGENCE: Embedding cache for zero-overhead reuse
        self.embedding_cache = {}  # Cache compressed embeddings by hash
        self.cache_hits = 0
        self.cache_misses = 0
        self.cache_max_size = 1000  # Maximum cache size (LRU eviction)
        
        self.graph = nx.Graph()
        self.id_counter = 0
        self.addition_count = 0
        self.node_queue = deque()
        self.access_counts = {}
        self.last_accessed = {}
        self.importance_scores = {}  # Learnable importance scores for each node
        self.encoder = None  # Shared model will be used for query encoding
        self.rag_generator = rag_generator_instance  # Assign the passed RAG instance
        
        # AGI Memory Features: Human-like memory types and structures
        self.node_types = {}  # Track memory node types: 'episodic', 'semantic', 'procedural', 'working'
        self.temporal_connections = {}  # Temporal sequence connections between memories
        self.consolidation_buffer = []  # Buffer for memory consolidation
        self.forgetting_curves = {}  # Ebbinghaus forgetting curves for each node
        self.memory_hierarchies = {}  # Hierarchical memory organization
        
        # Learnable LRU parameters
        self.use_learnable_lru = True  # Enable learnable importance-based eviction
        self.importance_decay = 0.95  # Decay factor for importance over time
        self.access_boost = 0.1  # Boost to importance when accessed
        self.connectivity_weight = 0.3  # Weight for graph connectivity in importance
        
        # ⚡ EFFICIENT FORGETATIONS: Learnable forgetting system (better than human forgetting!)
        # This learns what to forget and what to keep, adapting over time
        self.forgetation_model = None  # Neural network that learns forgetting patterns
        self.forgetation_history = []  # History of what was forgotten and if it was a mistake
        self.forgetation_learning_rate = 1e-3  # Learning rate for forgetation model
        self.forgetation_buffer = deque(maxlen=100)  # Buffer for training forgetation model
        self.forgetation_optimizer = None  # Optimizer for forgetation model
        self._initialize_forgetation_model()
        
        # Initialize combined VAE + PCA compression (AGI breakthrough approach)
        if self.use_vae_compression and self.use_combined_compression:
            try:
                from neuro_fusion import VAECompressor

                # VAE will be initialized when we get first embedding dimension
                # PCA will be applied to VAE latent space for maximum compression
                print(f"[AGI] Combined VAE+PCA compression enabled (target latent dim: {vae_latent_dim})")
                print(f"[AGI] Compression pipeline: Original -> VAE -> PCA -> Quantization (no fallbacks)")
            except ImportError:
                if self.use_combined_compression:
                    print("[ERROR] VAE compressor required for combined compression - cannot proceed without it")
                    raise ImportError("VAE compressor is required for AGI memory system")
                else:
                    print("[WARN] VAE compressor not available, but combined compression disabled - will use legacy mode")
                    self.use_vae_compression = False

        # Load existing graph if available (but preserve max_nodes setting)
        saved_max_nodes = self.max_nodes
        # Try loading compressed version first, then fallback to uncompressed for migration
        if os.path.exists(self.graph_file):
            self.load_graph()
            # Restore max_nodes setting (don't let loaded graph override infinite growth)
            if saved_max_nodes is None:
                self.max_nodes = None  # Ensure infinite growth is preserved
        else:
            # Check for old uncompressed version for migration
            old_file = self.graph_file.replace('.gz', '')
            if old_file.endswith('.json') and os.path.exists(old_file):
                print(f"[MIGRATE] Found old uncompressed file {old_file}, loading and will convert to compressed format")
                # Temporarily use old file for loading
                temp_file = self.graph_file
                self.graph_file = old_file
                self.load_graph()
                self.graph_file = temp_file  # Restore .gz path
                # Restore max_nodes setting
                if saved_max_nodes is None:
                    self.max_nodes = None
                # Save immediately in compressed format
                self.save_graph()
                # Clean up old uncompressed file
                try:
                    os.remove(old_file)
                    print(f"[CLEANUP] Removed old uncompressed file: {old_file}")
                except Exception as e:
                    print(f"[WARN] Could not remove old file {old_file}: {e}")

        max_nodes_str = "INFINITE" if self.max_nodes is None else str(self.max_nodes)
        compression_str = "VAE" if self.use_vae_compression else "PCA"
        print(f"[NETWORK] Initialized MemoryGraph on {self.device} | "
              f"Max Nodes={max_nodes_str}, Save Interval={self.save_interval}, Compression={compression_str}")
        
        # Debug: Verify max_nodes is correct
        if self.max_nodes is not None:
            print(f"[WARNING] max_nodes is set to {self.max_nodes} - this will limit growth!")
        else:
            print(f"[CONFIG] Infinite growth enabled - graph will grow without hard limit (using VAE compression for efficiency)")

    def _alpha_ratio(self, text: str) -> float:
        if not text:
            return 0.0
        alpha = sum(1 for ch in text if ch.isalpha())
        total = max(len(text), 1)
        return alpha / total

    def _sanitize_text(self, text: str, max_chars: int = 2000) -> str:
        if not text:
            return ""
        cleaned = text.strip()
        cleaned = re.sub(r"```.*?```", "", cleaned, flags=re.DOTALL)
        cleaned = cleaned.replace("```", "")
        cleaned = re.sub(r"(?i)\b(final answer|this is the final answer|answer:)\b", "", cleaned)
        cleaned = re.sub(r"(?i)here'?s my response:\s*", "", cleaned)
        cleaned = re.sub(r"<[^>]+>", "", cleaned)
        cleaned = re.sub(r"\s+\n", "\n", cleaned).strip()
        if len(cleaned) > max_chars:
            cleaned = cleaned[:max_chars].rstrip() + "..."
        return cleaned

    def _is_low_quality_chunk(self, text: str) -> bool:
        if not text or len(text) < 20:
            return True
        return self._alpha_ratio(text) < 0.6

    def _score_context(self, similarity: float, importance: float, last_accessed: float) -> float:
        recency_hours = max((time.time() - last_accessed) / 3600.0, 0.0)
        recency = 1.0 / (1.0 + recency_hours / 24.0)
        return float(similarity) * (0.7 + 0.3 * float(importance)) * recency

    def _resolve_graph_path(self, graph_file: str) -> str:
        if os.path.isabs(graph_file):
            path = graph_file
        else:
            project_root = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
            memory_root = os.path.join(project_root, "memory")
            path = os.path.join(memory_root, graph_file)
        os.makedirs(os.path.dirname(path), exist_ok=True)
        return path

    def _dynamic_save_interval(self) -> int:
        size = max(self.graph.number_of_nodes(), 1)
        target = int(size * 0.02)
        target = max(self.min_save_interval, target)
        return min(self.save_interval, target)

    def _should_save(self) -> bool:
        additions_trigger = self.addition_count >= self._dynamic_save_interval()
        time_trigger = (time.time() - self.last_save_time) >= self.save_interval_seconds
        return additions_trigger or time_trigger

    def _initialize_vae_compressor(self, embedding: torch.Tensor):
        """Initializes the AGI VAE compressor with shared backbone references for zero-overhead."""
        if self.vae_compressor is not None:
            return
        
        try:
            from neuro_fusion import VAECompressor
            input_dim = embedding.size(-1)
            self.original_dim = input_dim
            # AGI breakthrough: use 2x hidden for better latent mapping
            hidden_dim = input_dim * 2
            
            self.vae_compressor = VAECompressor(
                input_dim=input_dim,
                hidden_dim=hidden_dim,
                latent_dim=self.vae_latent_dim,
                shared_model=self.shared_model,
                shared_tokenizer=self.shared_tokenizer,
                enable_training=self.enable_vae_training,
                learning_rate=self.vae_learning_rate,
                device=self.device
            )
            print(f"[MEMORY] VAE Compressor initialized (input_dim={input_dim}, latent_dim={self.vae_latent_dim}) [Unified: {self.shared_model is not None}]")
        except Exception as e:
            print(f"[WARN] VAE Initialization failed: {e}. Falling back to legacy PCA.")
            self.use_vae_compression = False

    def _initialize_compressor(self, embedding: torch.Tensor):
        """Legacy PCA-only initialization."""
        self.original_dim = embedding.size(-1)
        target_dim = max(int(self.original_dim * self.compression_level), 1)
        try:
            self.pca = PCA(n_components=target_dim)
            print(f"[CONFIG] Legacy PCA initialized with {target_dim} components.")
        except Exception as e:
            print(f"[WARN] PCA Initialization failed: {e}")
            self.pca = None

    def _initialize_forgetation_model(self):
        """Initializes the neural network that learns the Ebbinghaus forgetting patterns."""
        try:
            import torch.nn as nn
            from torch.optim import Adam
            
            # 5 features: [access, time, connectivity, type, consolidation]
            self.forgetation_model = nn.Sequential(
                nn.Linear(5, 16),
                nn.ReLU(),
                nn.Linear(16, 8),
                nn.ReLU(),
                nn.Linear(8, 1),
                nn.Sigmoid()
            ).to_empty(device=self.device)
            # ⚡ Fix: Since to_empty doesn't initialize parameters, we must call reset_parameters or similar if needed.
            # But nn.Sequential doesn't have it, so we iterate.
            for layer in self.forgetation_model:
                if hasattr(layer, 'reset_parameters'):
                    layer.reset_parameters()
            self.forgetation_optimizer = Adam(self.forgetation_model.parameters(), lr=self.forgetation_learning_rate)
            print("[FORGETATIONS] Neural forgetting system active.")
        except Exception as e:
            print(f"[WARN] Failed to initialize forgetation system: {e}")

    def _compress_embedding(self, embedding: torch.Tensor, embedding_hash: Optional[str] = None) -> dict:
        """Compresses embedding via AGI pipeline: VAE -> PCA -> Quantization -> Base64."""
        # ⚡ INFINITE INTELLIGENCE: Zero-overhead cache check
        if embedding_hash is not None and embedding_hash in self.embedding_cache:
            self.cache_hits += 1
            return self.embedding_cache[embedding_hash]
        
        self.cache_misses += 1
        
        # Auto-init if needed
        if self.use_vae_compression and self.vae_compressor is None:
            self._initialize_vae_compressor(embedding)
        
        if self.use_vae_compression and self.vae_compressor is not None:
            # ⚡ OPTIMIZATION: Process in float32 for VAE precision
            embedding_f32 = embedding.to(torch.float32).unsqueeze(0)
            
            # Learn from this experience (incremental training)
            if self.enable_vae_training and self.vae_compressor.is_trainable and hasattr(self.vae_compressor, "train_incremental"):
                self.vae_compressor.train_incremental(embedding_f32, beta=1.0)
            
            with torch.no_grad():
                # Stage 1: VAE Encode
                mu, _ = self.vae_compressor._encode(embedding_f32)
                vae_latent = mu.squeeze(0).cpu().numpy().reshape(1, -1)
                
                # Stage 2: PCA on VAE space
                if self.pca_on_vae is not None:
                    if not hasattr(self.pca_on_vae, 'components_'):
                         self.consolidation_buffer.append(vae_latent[0])
                         if len(self.consolidation_buffer) >= 20:
                             self.pca_on_vae.fit(np.array(self.consolidation_buffer))
                             compressed = self.pca_on_vae.transform(vae_latent)[0]
                             self.consolidation_buffer = []
                         else:
                             compressed = vae_latent[0]
                    else:
                        compressed = self.pca_on_vae.transform(vae_latent)[0]
                else:
                    compressed = vae_latent[0]
                
                # Stage 3: Quantization
                scale = (2 ** (self.quantization_bits - 1) - 1)
                mn, mx = compressed.min(), compressed.max()
                rng = max(abs(mx - mn), 1e-5)
                norm = (compressed - mn) / rng
                quant = np.round(norm * scale).astype(np.int8)
                
                # Stage 4: Base64
                quant_b64 = base64.b64encode(quant.tobytes()).decode('ascii')
                
                result = {
                    'data': quant_b64,
                    'format': 'base64_int8',
                    'min': float(mn),
                    'max': float(mx),
                    'bits': self.quantization_bits,
                    'compressed': True,
                    'vae': True,
                    'pca_on_vae': self.pca_on_vae is not None,
                    'latent_dim': self.vae_latent_dim,
                    'final_dim': len(compressed),
                    'shape': list(quant.shape)
                }
                
                # Cache results
                if embedding_hash is not None:
                    if len(self.embedding_cache) >= self.cache_max_size:
                        # FIFO eviction
                        del self.embedding_cache[next(iter(self.embedding_cache))]
                    self.embedding_cache[embedding_hash] = result
                
                return result
        else:
            # Legacy PCA Fallback
            if self.pca is None:
                self._initialize_compressor(embedding)
            
            emb_np = embedding.cpu().detach().numpy().reshape(1, -1)
            if hasattr(self.pca, 'components_'):
                compressed = self.pca.transform(emb_np)[0]
            else:
                compressed = emb_np[0]
            
            # Simple Quant/B64
            scale = (2 ** (self.quantization_bits - 1) - 1)
            mn, mx = compressed.min(), compressed.max()
            rng = max(abs(mx - mn), 1e-5)
            norm = (compressed - mn) / rng
            quant = np.round(norm * scale).astype(np.int8)
            quant_b64 = base64.b64encode(quant.tobytes()).decode('ascii')
            
            return {
                'data': quant_b64, 'format': 'base64_int8', 'min': float(mn), 'max': float(mx),
                'bits': self.quantization_bits, 'compressed': True, 'vae': False,
                'shape': list(quant.shape)
            }

    def _decompress_embedding(self, compressed_data):
        """Decompress embedding using AGI VAE or heritage PCA."""
        if isinstance(compressed_data, dict) and compressed_data.get('vae', False):
            # Prep VAE
            if self.vae_compressor is None:
                 latent_dim = compressed_data.get('latent_dim', self.vae_latent_dim)
                 if self.original_dim is None: self.original_dim = 512 # Heuristic
                 self._initialize_vae_compressor(torch.randn(self.original_dim))
                 
            try:
                data = compressed_data['data']
                if isinstance(data, str) and compressed_data.get('format') == 'base64_int8':
                    arr = np.frombuffer(base64.b64decode(data), dtype=np.int8)
                    if 'shape' in compressed_data: arr = arr.reshape(compressed_data['shape'])
                    # Dequantize
                    scale = (2 ** (compressed_data['bits'] - 1) - 1)
                    mn, mx = compressed_data['min'], compressed_data['max']
                    rng = max(abs(mx - mn), 1e-5)
                    arr = arr.astype(np.float32) / scale * rng + mn
                else:
                    arr = np.array(data)
                # Inverse PCA on VAE
                if compressed_data.get('pca_on_vae', False) and self.pca_on_vae is not None:
                     arr = self.pca_on_vae.inverse_transform(arr.reshape(1, -1))[0]
                # VAE Decode
                z_tensor = torch.tensor(arr, device=self.device, dtype=torch.float32).unsqueeze(0)
                with torch.no_grad():
                    reconstructed = self.vae_compressor._decode(z_tensor)
                return reconstructed.squeeze(0).to(self.device).to(torch.float16)
            except Exception as e:
                print(f"[WARN] VAE Decompression error: {e}")
        
        # Standard Fallback
        if isinstance(compressed_data, dict) and 'data' in compressed_data:
            data = compressed_data['data']
            if isinstance(data, str):
                arr = np.frombuffer(base64.b64decode(data), dtype=np.int8 if 'int8' in compressed_data.get('format','') else np.float32)
                if 'shape' in compressed_data: arr = arr.reshape(compressed_data['shape'])
                if 'bits' in compressed_data and 'int8' in compressed_data.get('format',''):
                    scale = (2 ** (compressed_data['bits'] - 1) - 1)
                    mn, mx = compressed_data['min'], compressed_data['max']
                    rng = max(abs(mx - mn), 1e-5)
                    arr = arr.astype(np.float32) / scale * rng + mn
            else:
                arr = np.array(data)
        else:
            arr = np.array(compressed_data)
            
        if self.pca is not None and hasattr(self.pca, 'components_'):
             if arr.shape[0] < self.pca.n_components_:
                 arr = np.pad(arr, (0, self.pca.n_components_ - arr.shape[0]))
             arr = self.pca.inverse_transform(arr.reshape(1, -1))[0]
             
        return torch.tensor(arr, device=self.device, dtype=torch.float16)

    def add_experience(self, embedding: torch.Tensor, metadata=None, memory_type=None):
        """
        AGI Memory Addition: Add experience with human-like memory types and features.

        Args:
            embedding (torch.Tensor): The embedding of the experience.
            metadata (dict, optional): Additional metadata for the experience.
                                       Must include 'original_text' for RAG to function properly.
            memory_type (str, optional): Memory type - 'episodic', 'semantic', 'procedural', 'working'.
                                         Auto-detected if not provided.
        """
        if metadata is None:
            metadata = {}
        if 'original_text' not in metadata:
            print("[WARN] Warning: 'original_text' not found in metadata. RAG retrieval might be limited for this experience.")
            metadata['original_text'] = ""  # Provide a default empty string
        else:
            metadata['original_text'] = self._sanitize_text(metadata.get('original_text', ''), max_chars=2000)

        # AGI Feature: Auto-detect memory type if not provided
        if memory_type is None and self.enable_memory_types:
            memory_type = self._detect_memory_type(metadata, embedding)
        elif memory_type is None:
            memory_type = 'episodic'  # Default
        
        # Store memory type
        if self.enable_memory_types:
            metadata['memory_type'] = memory_type

        compressed = self._compress_embedding(embedding)
        nid = self.id_counter
        current_time = time.time()
        
        self.graph.add_node(nid, embedding=compressed, metadata=metadata)
        self.node_queue.append(nid)
        self.access_counts[nid] = 0
        self.last_accessed[nid] = current_time
        
        # AGI Features: Initialize human-like memory properties
        if self.enable_memory_types:
            self.node_types[nid] = memory_type
        
        # Initialize importance score (new nodes start with moderate importance)
        self.importance_scores[nid] = 0.5
        
        # AGI Feature: Initialize forgetting curve (Ebbinghaus model)
        if self.enable_consolidation:
            self.forgetting_curves[nid] = {
                'initial_strength': 1.0,
                'decay_rate': 0.1,  # Slower decay for important memories
                'last_accessed': current_time,
                'access_count': 0,
                'consolidation_level': 0.0  # Increases with repeated access
            }
        
        # AGI Feature: Temporal connection to previous memory
        if self.enable_temporal_connections and len(self.node_queue) > 1:
            prev_nid = self.node_queue[-2]  # Previous node
            if prev_nid in self.graph.nodes:
                # Add temporal edge (sequence connection)
                if prev_nid not in self.temporal_connections:
                    self.temporal_connections[prev_nid] = []
                self.temporal_connections[prev_nid].append(nid)
                # Also add to graph as temporal edge
                self.graph.add_edge(prev_nid, nid, weight=0.9, edge_type='temporal')

        # Link to the last 50 nodes if similarity > threshold
        nodes_to_check = list(self.graph.nodes)[-50:]
        if nid in nodes_to_check:  # Remove self if it somehow got included (shouldn't happen with [-50:] if nid is newest)
            nodes_to_check.remove(nid)

        for other_nid in nodes_to_check:
            if other_nid == nid:  # Redundant check but safe
                continue
            try:
                other_emb = self.graph.nodes[other_nid]['embedding']
                sim = self._cosine_similarity(
                    embedding,
                    self._decompress_embedding(other_emb)
                ).item()
                if sim > self.similarity_threshold:
                    self.graph.add_edge(nid, other_nid, weight=sim)
            except Exception as e:
                print(f"[ERROR] Error creating edge between {nid} and {other_nid}: {e}")
                continue

        self.id_counter += 1
        self.addition_count += 1

        # Learnable LRU Prune - only prune if:
        # 1. max_nodes is set AND exceeded, OR
        # 2. Memory pressure is high (GPU memory > threshold)
        should_prune = False
        prune_reason = ""
        
        # Debug: Check if max_nodes is unexpectedly set (likely from old saved graph)
        if self.max_nodes is not None and self.max_nodes <= 200:
            print(f"[WARNING] max_nodes={self.max_nodes} detected - this may be from old saved graph. Current nodes: {self.graph.number_of_nodes()}")
            print(f"[WARNING] Forcing max_nodes=None to enable infinite growth...")
            self.max_nodes = None  # Force infinite growth if old limit detected
        
        if self.max_nodes is not None and self.graph.number_of_nodes() > self.max_nodes:
            should_prune = True
            prune_reason = f"exceeded max_nodes ({self.max_nodes})"
        elif self.max_nodes is None:  # Infinite growth mode - check memory pressure
            # Check GPU memory pressure
            if torch.cuda.is_available():
                try:
                    # Better memory pressure calculation
                    total_memory = torch.cuda.get_device_properties(0).total_memory
                    reserved_memory = torch.cuda.memory_reserved(self.device)
                    allocated_memory = torch.cuda.memory_allocated(self.device)
                    
                    # Calculate actual memory usage as percentage of total GPU memory
                    memory_usage = reserved_memory / total_memory if total_memory > 0 else 0
                    
                    # Only prune if memory usage exceeds threshold AND we have enough nodes to prune
                    # Don't prune if we have fewer than 1000 nodes (allow growth)
                    if memory_usage > self.memory_pressure_threshold and self.graph.number_of_nodes() > 1000:
                        should_prune = True
                        prune_reason = f"high memory pressure ({memory_usage:.1%} > {self.memory_pressure_threshold:.1%})"
                except Exception as e:
                    # If memory check fails, use node count as fallback (prune if > 50k nodes)
                    # Much higher threshold for infinite growth
                    if self.graph.number_of_nodes() > 50000:
                        should_prune = True
                        prune_reason = f"large graph size ({self.graph.number_of_nodes()} nodes)"
        
        if should_prune:
            evicted = self.prune(1)
            if evicted:
                print(f"[REMOVE] Evicted node {evicted} ({'Importance-based' if self.use_learnable_lru else 'LRU'}, {prune_reason}). Now {self.graph.number_of_nodes()} nodes remain.")

        # Async Save using dynamic interval + time-based trigger
        if self._should_save():
            dynamic_interval = self._dynamic_save_interval()
            print(f"[SAVE] Async saving MemoryGraph (after {dynamic_interval} adds or {self.save_interval_seconds}s)...")
            try:
                from thread_manager import create_managed_thread
                create_managed_thread(target=self.save_graph, name="MemoryGraphSave")
            except ImportError:
                # Fallback to synchronous save if thread_manager not available
                self.save_graph()
            self.addition_count = 0
            self.last_save_time = time.time()

        print(f"[PACKAGE] MemoryGraph size: {self.graph.number_of_nodes()}")

    # alias for backward compatibility
    add_node = add_experience

    def prune(self, k: int = 1):
        """
        ⚡ EFFICIENT FORGETATIONS: Learnable forgetting system.
        Removes k nodes using intelligent, adaptive forgetting that learns what to forget.
        Better than human forgetting because it learns from mistakes and adapts over time!
        
        Uses:
        1. Learnable neural network to predict importance
        2. Learns from mistakes (if forgotten memory is later needed, learn from it)
        3. Adapts forgetting curves based on experience
        4. Uses semantic understanding from shared model when available
        """
        if not self.graph.nodes:
            return None
        
        removed = []
        
        if self.use_learnable_lru:
            # ⚡ EFFICIENT FORGETATIONS: Use learnable model if available
            node_importances = {}
            current_time = time.time()
            
            for nid in list(self.graph.nodes):
                # Extract features for forgetation model
                access_count = self.access_counts.get(nid, 0)
                time_since_access = current_time - self.last_accessed.get(nid, current_time)
                connectivity = self.graph.degree(nid)
                max_degree = max([self.graph.degree(n) for n in self.graph.nodes()] + [1])
                normalized_connectivity = connectivity / max_degree if max_degree > 0 else 0
                
                # Memory type encoding (episodic=0.0, semantic=0.33, procedural=0.66, working=1.0)
                memory_type = self.node_types.get(nid, 'episodic')
                type_encoding = {'episodic': 0.0, 'semantic': 0.33, 'procedural': 0.66, 'working': 1.0}.get(memory_type, 0.0)
                
                # Consolidation level from forgetting curve
                forgetting_curve = self.forgetting_curves.get(nid, {})
                consolidation = forgetting_curve.get('consolidation_level', 0.0)
                
                # Normalize features
                normalized_access = min(access_count / 100.0, 1.0)  # Cap at 100 accesses
                normalized_time = min(time_since_access / (365 * 24 * 3600), 1.0)  # Normalize to 1 year
                
                # Use forgetation model if available
                if self.forgetation_model is not None:
                    try:
                        # Prepare input features
                        features = torch.tensor([
                            normalized_access,
                            normalized_time,
                            normalized_connectivity,
                            type_encoding,
                            consolidation
                        ], device=self.device, dtype=torch.float32).unsqueeze(0)
                        
                        # Predict importance using learnable model
                        self.forgetation_model.eval()
                        with torch.no_grad():
                            predicted_importance = self.forgetation_model(features).item()
                        
                        # Combine with base importance (weighted average)
                        base_importance = self.importance_scores.get(nid, 0.5)
                        importance = 0.7 * predicted_importance + 0.3 * base_importance
                    except Exception as e:
                        # Fallback to rule-based if model fails
                        base_importance = self.importance_scores.get(nid, 0.5)
                        access_boost = min(access_count * self.access_boost, 1.0)
                        connectivity_boost = normalized_connectivity * self.connectivity_weight
                        time_decay = self.importance_decay ** (time_since_access / 3600)
                        importance = (base_importance + access_boost + connectivity_boost) * time_decay
                else:
                    # Rule-based fallback (backward compatible)
                    base_importance = self.importance_scores.get(nid, 0.5)
                    access_boost = min(access_count * self.access_boost, 1.0)
                    connectivity_boost = normalized_connectivity * self.connectivity_weight
                    time_decay = self.importance_decay ** (time_since_access / 3600)
                    importance = (base_importance + access_boost + connectivity_boost) * time_decay
                
                node_importances[nid] = importance
            
            # Sort by importance (lowest first) and remove k nodes
            sorted_nodes = sorted(node_importances.items(), key=lambda x: x[1])
            nodes_to_remove = [nid for nid, _ in sorted_nodes[:k]]
            
            for nid in nodes_to_remove:
                if nid in self.graph:
                    # ⚡ EFFICIENT FORGETATIONS: Record what we're forgetting for learning
                    forgetation_record = {
                        'node_id': nid,
                        'importance': node_importances[nid],
                        'access_count': self.access_counts.get(nid, 0),
                        'connectivity': self.graph.degree(nid),
                        'memory_type': self.node_types.get(nid, 'episodic'),
                        'timestamp': current_time
                    }
                    self.forgetation_history.append(forgetation_record)
                    
                    # Remove node
                    self.graph.remove_node(nid)
                    self.access_counts.pop(nid, None)
                    self.last_accessed.pop(nid, None)
                    self.importance_scores.pop(nid, None)
                    # Remove from queue if present
                    if nid in self.node_queue:
                        self.node_queue.remove(nid)
                    removed.append(nid)
        else:
            # Simple LRU: remove oldest nodes
            for _ in range(k):
                if not self.node_queue:
                    break
                old = self.node_queue.popleft()
                if old in self.graph:
                    self.graph.remove_node(old)
                    self.access_counts.pop(old, None)
                    self.last_accessed.pop(old, None)
                    self.importance_scores.pop(old, None)
                    removed.append(old)
                else:
                    print(f"Node {old} not found in graph during prune, skipping.")
        
        # ⚡ EFFICIENT FORGETATIONS: Train model if we have enough history
        if len(self.forgetation_history) >= 10 and self.forgetation_model is not None:
            self._train_forgetation_model()
        
        return removed[-1] if removed else None
    
    def _train_forgetation_model(self):
        """
        ⚡ EFFICIENT FORGETATIONS: Train the learnable forgetting model.
        Learns from mistakes - if a forgotten memory is later needed, we learn from it!
        """
        if self.forgetation_model is None or len(self.forgetation_history) < 10:
            return
        
        try:
            import torch.nn as nn
            
            # Prepare training data from history
            # Check if any forgotten memories were later accessed (mistakes)
            training_data = []
            
            for record in self.forgetation_history[-50:]:  # Use last 50 records
                nid = record['node_id']
                # Check if this node was later accessed (mistake - we shouldn't have forgotten it)
                was_mistake = False
                if nid in self.access_counts:  # If it's still accessible, it wasn't forgotten
                    was_mistake = False
                else:
                    # Check if it was accessed after being forgotten (would need to track this)
                    # For now, use importance as proxy - if importance was high, it was a mistake
                    was_mistake = record['importance'] > 0.3  # If importance > 0.3, probably a mistake
                
                # Features: [access_count, time_since_access, connectivity, memory_type, consolidation]
                features = torch.tensor([
                    min(record['access_count'] / 100.0, 1.0),
                    min((time.time() - record['timestamp']) / (365 * 24 * 3600), 1.0),
                    min(record['connectivity'] / 10.0, 1.0),
                    {'episodic': 0.0, 'semantic': 0.33, 'procedural': 0.66, 'working': 1.0}.get(record['memory_type'], 0.0),
                    0.0  # Consolidation (would need to track)
                ], device=self.device, dtype=torch.float32)
                
                # Target: 1.0 if it was a mistake (shouldn't have forgotten), 0.0 if correct
                target = torch.tensor([1.0 if was_mistake else 0.0], device=self.device, dtype=torch.float32)
                
                training_data.append((features, target))
            
            if len(training_data) < 5:
                return  # Need at least 5 samples
            
            # Train for a few steps
            self.forgetation_model.train()
            criterion = nn.MSELoss()
            
            for epoch in range(3):  # Quick training
                total_loss = 0.0
                for features, target in training_data:
                    self.forgetation_optimizer.zero_grad()
                    prediction = self.forgetation_model(features.unsqueeze(0))
                    loss = criterion(prediction, target.unsqueeze(0))
                    loss.backward()
                    torch.nn.utils.clip_grad_norm_(self.forgetation_model.parameters(), max_norm=1.0)
                    self.forgetation_optimizer.step()
                    total_loss += loss.item()
            
            self.forgetation_model.eval()
            
            if self.addition_count % 100 == 0:
                avg_loss = total_loss / len(training_data)
                print(f"[FORGETATIONS] Trained model (avg loss: {avg_loss:.4f}, samples: {len(training_data)})")
        except Exception as e:
            # Silently fail - don't break the system if training fails
            pass
    
    def learn_from_forgetting_mistake(self, forgotten_node_id: int, actual_importance: float):
        """
        ⚡ EFFICIENT FORGETATIONS: Learn from a mistake.
        Call this when we realize we shouldn't have forgotten a memory.
        This makes the system better than human forgetting - it learns from errors!
        
        Args:
            forgotten_node_id: ID of the node that was forgotten but shouldn't have been
            actual_importance: The actual importance (higher = bigger mistake)
        """
        if self.forgetation_model is None:
            return
        
        # Find the forgetation record
        for record in reversed(self.forgetation_history):
            if record['node_id'] == forgotten_node_id:
                # Add to training buffer as a mistake
                self.forgetation_buffer.append({
                    'record': record,
                    'was_mistake': True,
                    'actual_importance': actual_importance
                })
                break
        
        # Train immediately if buffer is full
        if len(self.forgetation_buffer) >= 10:
            self._train_forgetation_model()
            self.forgetation_buffer.clear()

    def _detect_memory_type(self, metadata: dict, embedding: torch.Tensor) -> str:
        """
        AGI Feature: Auto-detect memory type based on metadata and content.
        Human-like memory classification.
        """
        # Check metadata for explicit type hints
        text = metadata.get('original_text', '').lower()
        metadata_type = metadata.get('type', '').lower()
        
        # Procedural: Contains action words, steps, instructions
        procedural_keywords = ['step', 'how to', 'procedure', 'process', 'method', 'algorithm', 'instruction']
        if any(kw in text for kw in procedural_keywords) or 'procedural' in metadata_type:
            return 'procedural'
        
        # Semantic: Facts, concepts, definitions
        semantic_keywords = ['definition', 'concept', 'fact', 'knowledge', 'meaning', 'is a', 'refers to']
        if any(kw in text for kw in semantic_keywords) or 'semantic' in metadata_type:
            return 'semantic'
        
        # Working: Short-term, recent, active
        if metadata.get('stage', 0) == 0 or 'working' in metadata_type:
            return 'working'
        
        # Episodic: Default - personal experiences, events, narratives
        return 'episodic'
    
    def _update_forgetting_curve(self, nid: int):
        """
        AGI Feature: Update Ebbinghaus forgetting curve for memory consolidation.
        """
        if nid not in self.forgetting_curves:
            return
        
        curve = self.forgetting_curves[nid]
        current_time = time.time()
        time_since_access = current_time - curve['last_accessed']
        
        # Ebbinghaus forgetting curve: R = e^(-t/S)
        # where R is retention, t is time, S is strength
        # Update strength based on access frequency
        access_count = curve['access_count']
        consolidation = curve['consolidation_level']
        
        # More consolidated memories decay slower
        effective_decay = curve['decay_rate'] * (1.0 - consolidation * 0.5)
        retention = np.exp(-time_since_access / (3600 * (1.0 + consolidation)))  # Hours scale
        
        # Update consolidation based on repeated access
        if access_count > 0:
            consolidation = min(1.0, consolidation + 0.1 * (1.0 - consolidation))
        
        curve['consolidation_level'] = consolidation
        curve['last_accessed'] = current_time
        curve['access_count'] += 1
        
        # Update importance based on retention and consolidation
        self.importance_scores[nid] = max(0.1, min(1.0, retention * (0.5 + consolidation * 0.5)))
    
    def consolidate_memories(self, threshold: float = 0.3):
        """
        AGI Feature: Memory consolidation - strengthen important memories, weaken forgotten ones.
        Human-like memory reconsolidation process.
        """
        if not self.enable_consolidation:
            return
        
        consolidated = 0
        weakened = 0
        
        for nid in list(self.graph.nodes):
            if nid not in self.forgetting_curves:
                continue
            
            self._update_forgetting_curve(nid)
            importance = self.importance_scores.get(nid, 0.5)
            
            # Strengthen well-consolidated memories
            if importance > 0.7:
                self.importance_scores[nid] = min(1.0, importance * 1.05)
                consolidated += 1
            # Weaken forgotten memories (but don't delete - they might be recalled)
            elif importance < threshold:
                self.importance_scores[nid] = max(0.05, importance * 0.95)
                weakened += 1
        
        if consolidated > 0 or weakened > 0:
            print(f"[CONSOLIDATE] Strengthened {consolidated} memories, weakened {weakened} forgotten memories")
    
    def _cosine_similarity(self, emb1: torch.Tensor, emb2: torch.Tensor):
        # Ensure tensors are on the same device and are float32 for F.cosine_similarity
        emb1 = emb1.to(self.device).to(torch.float32)
        emb2 = emb2.to(self.device).to(torch.float32)

        v1, v2 = emb1.view(-1), emb2.view(-1)
        # Handle differing dimensions by padding the smaller one or truncating the larger one
        min_len = min(v1.size(0), v2.size(0))
        if v1.size(0) > min_len:
            v1 = v1[:min_len]
        elif v2.size(0) > min_len:
            v2 = v2[:min_len]
        
        # Ensure tensors are not empty
        if v1.numel() == 0 or v2.numel() == 0:
            return torch.tensor(0.0, device=self.device)

        return F.cosine_similarity(
            v1.unsqueeze(0),
            v2.unsqueeze(0)
        ).squeeze()

    def retrieve_experience(self, query_embedding: torch.Tensor, top_k: int = 1) -> List[dict]:
        """
        Retrieve the top_k most similar experiences based on an embedding.
        Updates importance scores for retrieved nodes (learnable LRU).
        
        Args:
            query_embedding (torch.Tensor): The embedding to query with.
            top_k (int): Number of top results to return.

        Returns:
            List[dict]: Retrieved memory entries, each with 'embedding' and 'metadata' (including original_text).
        """
        sims = []
        for nid in list(self.graph.nodes):  # Convert to list to avoid RuntimeError if graph changes during iteration
            try:
                node_data = self.graph.nodes[nid]
                node_emb_data = node_data.get('embedding')
                
                if node_emb_data is None:
                    print(f"Warning: Node {nid} has no embedding data, skipping.")
                    continue

                node_emb_tensor = self._decompress_embedding(node_emb_data)
                
                sim = self._cosine_similarity(
                    query_embedding,
                    node_emb_tensor
                ).item()
                sims.append((nid, sim))
            except Exception as e:
                print(f"[ERROR] Error retrieving experience for node {nid}: {e}")
                continue
        sims.sort(key=lambda x: x[1], reverse=True)
        results = []
        for nid, sim_score in sims[:top_k]:
            node_data = self.graph.nodes[nid]
            try:
                # AGI Feature: Update importance and access tracking (learnable LRU + forgetting curves)
                if self.use_learnable_lru:
                    # ⚡ EFFICIENT FORGETATIONS: Check if this was a forgotten memory (mistake detection)
                    # If this node was in forgetation history, we made a mistake!
                    was_forgotten = False
                    for record in self.forgetation_history:
                        if record['node_id'] == nid:
                            was_forgotten = True
                            # Learn from mistake - this memory was important but we forgot it
                            actual_importance = sim_score  # Use similarity as proxy for importance
                            self.learn_from_forgetting_mistake(nid, actual_importance)
                            if self.addition_count % 50 == 0:
                                print(f"[FORGETATIONS] Learned from mistake: node {nid} was forgotten but is important (sim: {sim_score:.3f})")
                            break
                    
                    # Boost importance when accessed
                    current_importance = self.importance_scores.get(nid, 0.5)
                    # Extra boost if it was previously forgotten (bigger mistake = bigger boost)
                    boost = self.access_boost * (2.0 if was_forgotten else 1.0)
                    self.importance_scores[nid] = min(current_importance + boost, 1.0)
                    # Update access counts and time
                    self.access_counts[nid] = self.access_counts.get(nid, 0) + 1
                    self.last_accessed[nid] = time.time()
                    
                    # AGI Feature: Update forgetting curve on access (memory consolidation)
                    if self.enable_consolidation and nid in self.forgetting_curves:
                        self._update_forgetting_curve(nid)
                
                emb = self._decompress_embedding(node_data['embedding']).cpu().tolist()
                metadata = dict(node_data.get('metadata', {}))
                metadata['similarity_score'] = sim_score  # Add similarity score to metadata
                metadata['node_id'] = nid
                results.append({'embedding': emb, 'metadata': metadata})
            except Exception as e:
                print(f"Error processing retrieved node {nid} for results: {e}")
                continue
        return results

    def retrieve(self, query: str, top_k: int = 3) -> str:
        """
        Encode a text query, retrieve the top_k most similar experiences
        from the MemoryGraph, and then use RAG to generate a response.

        Args:
            query (str): The text string to query against memory.
            top_k (int): Number of top results to return to the RAG model as context.

        Returns:
            str: The RAG-generated answer based on the query and retrieved context.
        """
        # Check if RAG has CUDA errors - if so, skip RAG and return direct results
        if hasattr(self, '_rag_cuda_error') and self._rag_cuda_error:
            # Skip RAG, return direct memory results
            try:
                # Try shared model first
                q_emb = None
                if self.shared_model is not None:
                    try:
                        # Get tokenizer
                        if hasattr(self.shared_model, 'get_tokenizer'):
                            tokenizer = self.shared_model.get_tokenizer()
                        else:
                            tokenizer = getattr(self, 'shared_tokenizer', None)
                        
                        if tokenizer is not None:
                            inputs = tokenizer(query, return_tensors="pt", padding=True, truncation=True, max_length=512)
                            if hasattr(self.shared_model, 'get_model'):
                                model = self.shared_model.get_model()
                            else:
                                model = self.shared_model
                            device = next(model.parameters()).device
                            inputs = {k: v.to(device) for k, v in inputs.items()}
                            with torch.no_grad():
                                if hasattr(model, 'get_input_embeddings'):
                                    input_embeddings = model.get_input_embeddings()
                                    token_embeddings = input_embeddings(inputs['input_ids'])
                                    q_emb = token_embeddings.mean(dim=1)
                                    q_emb = F.normalize(q_emb, p=2, dim=1)
                    except:
                        pass
                
                if q_emb is None:
                    if self.shared_adapter is not None:
                        q_emb = self.shared_adapter.encode_text(query, max_length=512, use_hidden=False)
                    else:
                        return None
                
                if isinstance(q_emb, torch.Tensor) and hasattr(self, 'device'):
                    try:
                        q_emb = q_emb.to(self.device)
                    except:
                        pass
                retrieved_experiences = self.retrieve_experience(q_emb, top_k)
                # Return formatted results instead of RAG response
                results = []
                for exp in retrieved_experiences:
                    text = exp['metadata'].get('original_text', '')
                    if text:
                        results.append({'text': text, 'embedding': exp.get('embedding'), 'metadata': exp.get('metadata', {})})
                return results if results else None
            except:
                return None
        
        if self.rag_generator is None:
            # Don't print warning - RAG is optional and handled by AGIMemorySystem
            return None  # Return None instead of error string to indicate RAG not available
        
        # CRITICAL: Use shared Qwen model for encoding (zero memory overhead, best coherence)
        q_emb = None
        if self.shared_model is not None or (SHARED_MODEL_AVAILABLE and get_shared_model is not None):
            try:
                # Get shared model if not already set
                shared_model = self.shared_model
                if shared_model is None and SHARED_MODEL_AVAILABLE and get_shared_model is not None:
                    shared_model = get_shared_model()
                
                if shared_model is not None:
                    # Get tokenizer - try multiple methods for maximum compatibility
                    tokenizer = None
                    # Method 1: Direct method from shared_model
                    if hasattr(shared_model, 'get_tokenizer'):
                        tokenizer = shared_model.get_tokenizer()
                    # Method 2: Attribute access
                    elif hasattr(shared_model, 'tokenizer'):
                        tokenizer = shared_model.tokenizer
                    # Method 3: From instance attribute
                    elif hasattr(self, 'shared_tokenizer') and self.shared_tokenizer is not None:
                        tokenizer = self.shared_tokenizer
                    # Method 4: Use global get_shared_tokenizer function (most reliable)
                    elif SHARED_MODEL_AVAILABLE and get_shared_tokenizer is not None:
                        try:
                            tokenizer = get_shared_tokenizer()
                        except Exception:
                            pass
                    elif self.shared_adapter is not None:
                        tokenizer = self.shared_adapter.get_tokenizer()
                    
                    # Get model - try multiple methods
                    model = None
                    if hasattr(shared_model, 'get_model'):
                        model = shared_model.get_model()
                    elif hasattr(shared_model, 'model'):
                        model = shared_model.model
                    elif hasattr(shared_model, 'base_model'):
                        model = shared_model.base_model
                    else:
                        model = shared_model
                    
                    if tokenizer is not None and model is not None:
                        # CRITICAL: Use efficient tokenization (shorter sequences for speed)
                        inputs = tokenizer(
                            query,
                            return_tensors="pt",
                            padding=False,  # No padding for speed
                            truncation=True,
                            max_length=256  # Reduced from 512 for faster processing
                        )
                        
                        # Get device from model (handle wrapped models)
                        try:
                            # Try to get device from model parameters
                            if hasattr(model, 'parameters'):
                                device = next(model.parameters()).device
                            elif hasattr(model, 'device'):
                                device = model.device
                            else:
                                device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
                        except:
                            device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
                        
                        inputs = {k: v.to(device) for k, v in inputs.items()}
                        
                        # CRITICAL: Use embedding layer directly (fastest, zero overhead)
                        with torch.no_grad():
                            try:
                                # Method 1: Direct embedding layer (fastest)
                                if hasattr(model, 'get_input_embeddings'):
                                    input_embeddings = model.get_input_embeddings()
                                    token_embeddings = input_embeddings(inputs['input_ids'])
                                    # Mean pool and normalize
                                    q_emb = token_embeddings.mean(dim=1)
                                    q_emb = F.normalize(q_emb, p=2, dim=1)
                                # Method 2: Qwen-style embed_tokens (direct access)
                                elif hasattr(model, 'embed_tokens'):
                                    token_embeddings = model.embed_tokens(inputs['input_ids'])
                                    q_emb = token_embeddings.mean(dim=1)
                                    q_emb = F.normalize(q_emb, p=2, dim=1)
                                # Method 3: Access through base_model.embed_tokens (PEFT/LoRA wrapped)
                                elif hasattr(model, 'base_model') and hasattr(model.base_model, 'embed_tokens'):
                                    token_embeddings = model.base_model.embed_tokens(inputs['input_ids'])
                                    q_emb = token_embeddings.mean(dim=1)
                                    q_emb = F.normalize(q_emb, p=2, dim=1)
                                # Method 4: Access through model.model.embed_tokens (nested structure)
                                elif hasattr(model, 'model') and hasattr(model.model, 'embed_tokens'):
                                    token_embeddings = model.model.embed_tokens(inputs['input_ids'])
                                    q_emb = token_embeddings.mean(dim=1)
                                    q_emb = F.normalize(q_emb, p=2, dim=1)
                                # Method 5: Minimal forward pass (last resort, slower but works)
                                else:
                                    outputs = model(**inputs, output_hidden_states=True, use_cache=False)
                                    if hasattr(outputs, 'hidden_states') and outputs.hidden_states:
                                        q_emb = outputs.hidden_states[-1].mean(dim=1)
                                        q_emb = F.normalize(q_emb, p=2, dim=1)
                                    elif hasattr(outputs, 'last_hidden_state'):
                                        q_emb = outputs.last_hidden_state.mean(dim=1)
                                        q_emb = F.normalize(q_emb, p=2, dim=1)
                                    else:
                                        raise ValueError("Could not extract embeddings from model output")
                                
                                # Move to memory graph device (if different)
                                if isinstance(q_emb, torch.Tensor) and hasattr(self, 'device'):
                                    try:
                                        target_device = self.device
                                        if q_emb.device != target_device:
                                            q_emb = q_emb.to(target_device)
                                    except:
                                        pass  # Keep on original device if move fails
                                
                            except Exception as e:
                                # Log error but don't print (too verbose)
                                import logging
                                logger = logging.getLogger(__name__)
                                logger.debug(f"Embedding extraction failed: {e}")
                                q_emb = None
                    else:
                        # Tokenizer or model not available
                        if tokenizer is None:
                            import logging
                            logger = logging.getLogger(__name__)
                            logger.debug("Tokenizer not available for shared model encoding")
                        if model is None:
                            import logging
                            logger = logging.getLogger(__name__)
                            logger.debug("Model not available for shared model encoding")
                        q_emb = None
            except Exception as e:
                # Log error but don't print (too verbose during training)
                import logging
                logger = logging.getLogger(__name__)
                logger.debug(f"Shared model encoding failed: {e}")
                q_emb = None
        
        if q_emb is None:
            if self.shared_adapter is not None:
                q_emb = self.shared_adapter.encode_text(query, max_length=512, use_hidden=False)
            else:
                return None
        
        if q_emb is None:
            return None

        # Retrieve the top_k similar memories from the MemoryGraph
        retrieved_experiences = self.retrieve_experience(q_emb, top_k)
        
        # Extract original texts to use as context for RAG
        scored_contexts = []
        for exp in retrieved_experiences:
            metadata = exp.get('metadata', {})
            raw_text = metadata.get('original_text', '')
            cleaned = self._sanitize_text(raw_text, max_chars=2000)
            if self._is_low_quality_chunk(cleaned):
                continue
            sim = metadata.get('similarity_score', 0.0)
            nid = metadata.get('node_id')
            importance = self.importance_scores.get(nid, 0.5) if nid is not None else 0.5
            last_accessed = self.last_accessed.get(nid, time.time()) if nid is not None else time.time()
            score = self._score_context(sim, importance, last_accessed)
            scored_contexts.append((score, cleaned))
        scored_contexts.sort(key=lambda x: x[0], reverse=True)
        context_texts = [text for _, text in scored_contexts[:top_k]]
        
        # Also retrieve context from the RAG's internal knowledge base (Wikitext)
        wikitext_context = [self._sanitize_text(t, max_chars=1500) for t in self.rag_generator.retrieve_documents(query, top_k=top_k)]
        wikitext_context = [t for t in wikitext_context if not self._is_low_quality_chunk(t)]
        
        # Combine contexts, prioritizing MemoryGraph's context
        # Use dict.fromkeys to preserve order and remove duplicates
        combined_context = list(dict.fromkeys(context_texts + wikitext_context))
        max_total_chars = 3500
        trimmed_context = []
        total_chars = 0
        for chunk in combined_context:
            if total_chars + len(chunk) > max_total_chars:
                remaining = max_total_chars - total_chars
                if remaining > 100:
                    trimmed_context.append(chunk[:remaining].rstrip() + "...")
                break
            trimmed_context.append(chunk)
            total_chars += len(chunk)
        combined_context = trimmed_context
        
        if not combined_context:
            print("[NOTICE] No context available from either MemoryGraph or Wikitext. Generating general response.")
            # If still no context, try to generate but catch CUDA errors
            try:
                return self.rag_generator.generate_response(query=query, retrieved_context=[], max_length=100)
            except (RuntimeError, torch.cuda.CudaError, AssertionError) as e:
                error_str = str(e).lower()
                if any(keyword in error_str for keyword in ["cuda", "assert", "device-side"]):
                    self._rag_cuda_error = True
                    # Return direct memory results instead
                    results = []
                    for exp in retrieved_experiences:
                        text = exp['metadata'].get('original_text', '')
                        if text:
                            results.append({'text': text, 'embedding': exp.get('embedding'), 'metadata': exp.get('metadata', {})})
                    return results if results else None
                else:
                    raise

        # Generate response using RAG - catch CUDA errors
        try:
            rag_response = self.rag_generator.generate_response(
                query=query,
                retrieved_context=combined_context
            )
            return rag_response
        except (RuntimeError, torch.cuda.CudaError, AssertionError) as e:
            error_str = str(e).lower()
            if any(keyword in error_str for keyword in ["cuda", "assert", "device-side"]):
                self._rag_cuda_error = True
                print(f"[WARN] RAG generation failed due to CUDA error, returning direct memory results")
                # Return direct memory results instead of RAG response
                results = []
                for exp in retrieved_experiences:
                    text = exp['metadata'].get('original_text', '')
                    if text:
                        results.append({'text': text, 'embedding': exp.get('embedding'), 'metadata': exp.get('metadata', {})})
                return results if results else None
            else:
                raise

    def propagate(self, input_embedding: torch.Tensor):
        # Safely unpack the tensor shape
        try:
            B, L, D = input_embedding.shape
        except ValueError as e:
            print(f"[INSPECT] Failed to unpack tensor shape {input_embedding.shape}: {e}")
            # Handle unexpected shapes
            if input_embedding.dim() == 2:
                # 2D tensor [batch, features] - add sequence dimension
                B, D = input_embedding.shape
                L = 1
                input_embedding = input_embedding.unsqueeze(1)  # [B, 1, D]
            elif input_embedding.dim() == 4:
                # 4D tensor - flatten to 3D
                try:
                    B, L1, L2, D = input_embedding.shape
                    L = L1 * L2
                    input_embedding = input_embedding.view(B, L, D)
                    print(f"[INSPECT] Reshaped 4D tensor {input_embedding.shape} to 3D: [{B}, {L}, {D}]")
                except ValueError as e2:
                    print(f"[INSPECT] Failed to unpack 4D tensor shape {input_embedding.shape}: {e2}")
                    # Fallback: flatten all dimensions except batch and last
                    B = input_embedding.shape[0]
                    D = input_embedding.shape[-1]
                    L = input_embedding.numel() // (B * D)
                    input_embedding = input_embedding.view(B, L, D)
                    print(f"[INSPECT] Fallback reshape to 3D: [{B}, {L}, {D}]")
            else:
                raise ValueError(f"Unsupported tensor shape for propagate: {input_embedding.shape}")
        
        output = torch.zeros_like(input_embedding)
        norm_in = input_embedding / (input_embedding.norm(dim=-1, keepdim=True) + 1e-8)
        
        # Iterate over a limited number of recent nodes to avoid excessive computation
        nodes_to_propagate = list(self.graph.nodes)[-min(100, self.graph.number_of_nodes()):]

        if not nodes_to_propagate:
            return output  # Return zero tensor if no nodes

        for nid in nodes_to_propagate:
            try:
                node_emb = self._decompress_embedding(self.graph.nodes[nid]['embedding'])
                
                # Pad/truncate node_emb to match D
                if node_emb.size(0) < D:
                    # Pad with zeros to match dimension D
                    node_emb = F.pad(node_emb, (0, D - node_emb.size(0)), 'constant', 0)
                elif node_emb.size(0) > D:
                    # Truncate to dimension D
                    node_emb = node_emb[:D]
                
                # Ensure node_emb is 1D for norm calculation if it somehow becomes 0D
                if node_emb.dim() == 0:
                    node_emb = node_emb.unsqueeze(0)

                node_norm = node_emb / (node_emb.norm() + 1e-8)
                
                # Ensure consistent dimensions for cosine_similarity
                # norm_in is (B, L, D) and node_norm is (D,)
                # Unsqueeze node_norm to (1, 1, D) for broadcasting
                sim = F.cosine_similarity(norm_in, node_norm.view(1,1,-1), dim=-1)
                
                # Output has shape (B, L) after similarity. Need to expand to (B, L, 1) for element-wise multiplication
                output += sim.unsqueeze(-1) * input_embedding
            except Exception as e:
                print(f"[ERROR] Error during propagation for node {nid}: {e}")
                continue
        
        # Avoid division by zero
        num_propagated_nodes = len(nodes_to_propagate)
        return output / max(num_propagated_nodes, 1)

    def consolidate(self, remove_isolated: bool = False):
        to_remove_edges = []
        for u, v, data in self.graph.edges(data=True):
            # Check if 'weight' key exists and if its value is None or less than 0.01
            if 'weight' not in data or data['weight'] is None or data['weight'] < 0.01:
                to_remove_edges.append((u, v))
        self.graph.remove_edges_from(to_remove_edges)
        
        if remove_isolated:
            # Use a copy of the graph to find isolated nodes, as graph can change during iteration
            isolated = list(nx.isolates(self.graph.copy()))
            self.graph.remove_nodes_from(isolated)
            # Remove from node_queue as well
            self.node_queue = deque([nid for nid in self.node_queue if nid not in set(isolated)])
            for nid in isolated:
                self.access_counts.pop(nid, None)
                self.last_accessed.pop(nid, None)
            print(f"[REMOVE] Removed {len(isolated)} isolated nodes.")

    def save_graph(self):
        try:
            # Prepare embedding data for JSON serialization
            serializable_nodes = []
            for nid, data in self.graph.nodes(data=True):
                node_copy = data.copy()
                # Ensure embedding is a dictionary with 'data' key for serialization
                current_embedding = node_copy.get('embedding')
                if isinstance(current_embedding, torch.Tensor):
                    node_copy['embedding'] = {'data': current_embedding.cpu().tolist(), 'compressed': False}
                elif isinstance(current_embedding, list):  # Handle old list format
                    node_copy['embedding'] = {'data': current_embedding, 'compressed': False}
                elif not isinstance(current_embedding, dict) or 'data' not in current_embedding:
                    # Fallback for unexpected format, ensure it's a dict
                    print(f"Warning: Node {nid} has unexpected embedding format during save. Attempting conversion.")
                    if isinstance(current_embedding, dict) and current_embedding:  # if non-empty dict but no 'data'
                        node_copy['embedding'] = {'data': list(current_embedding.values())[0], 'compressed': False}  # Best guess
                    else:  # empty dict or other type
                        node_copy['embedding'] = {'data': [], 'compressed': False}  # Fallback to empty

                # Make sure metadata is also serializable (e.g., if it contains numpy types)
                for key, value in node_copy.get('metadata', {}).items():
                    if isinstance(value, np.ndarray):
                        node_copy['metadata'][key] = value.tolist()
                    elif isinstance(value, np.generic):  # Catch single numpy types like np.float32
                        node_copy['metadata'][key] = value.item()  # Convert to standard Python type

                serializable_nodes.append((nid, node_copy))

            # Create a new graph for serialization to ensure correct format
            serializable_graph = nx.Graph()
            serializable_graph.add_nodes_from(serializable_nodes)
            serializable_graph.add_edges_from(self.graph.edges(data=True))

            data = json_graph.node_link_data(serializable_graph, edges="links")
            
            compression_info = {
                'original_dim': self.original_dim,
                'compressed_dim': getattr(self.pca, 'n_components_', None) if self.pca else None,
                'quantization_bits': self.quantization_bits,
                'pca_components': self.pca.components_.tolist() if self.pca and hasattr(self.pca, 'components_') else None,
                'pca_mean': self.pca.mean_.tolist() if self.pca and hasattr(self.pca, 'mean_') else None,
                'use_vae_compression': self.use_vae_compression,
                'use_combined_compression': self.use_combined_compression,
                'vae_latent_dim': self.vae_latent_dim if self.use_vae_compression else None,
                'max_nodes': self.max_nodes,  # Save current max_nodes setting
                # AGI Features: Save PCA on VAE state
                'pca_on_vae_components': self.pca_on_vae.components_.tolist() if self.pca_on_vae and hasattr(self.pca_on_vae, 'components_') else None,
                'pca_on_vae_mean': self.pca_on_vae.mean_.tolist() if self.pca_on_vae and hasattr(self.pca_on_vae, 'mean_') else None,
                'pca_on_vae_dim': getattr(self.pca_on_vae, 'n_components_', None) if self.pca_on_vae else None
            }
            
            # Save VAE weights to separate file (more efficient than JSON)
            if self.vae_compressor is not None:
                try:
                    vae_weights_file = self.graph_file.replace('.json.gz', '_vae_weights.pt')
                    self.vae_compressor.save_weights(vae_weights_file)
                    compression_info['vae_weights_file'] = vae_weights_file
                    compression_info['vae_enable_training'] = self.enable_vae_training
                    compression_info['vae_learning_rate'] = self.vae_learning_rate
                    print(f"[VAE] Saved VAE weights to {vae_weights_file}")
                except Exception as e:
                    print(f"[WARN] Could not save VAE weights: {e}")

            # OPTIMIZATION: Save as gzip-compressed JSON (much smaller file size)
            # Remove indent to save space (30-40% reduction)
            save_data = {
                'graph_data': data,
                'id_counter': self.id_counter,
                'node_queue': list(self.node_queue),
                'access_counts': self.access_counts,
                'last_accessed': self.last_accessed,
                'importance_scores': self.importance_scores,  # Save learnable importance scores
                'compression_info': compression_info,
                # AGI Features: Save human-like memory structures
                'node_types': self.node_types if self.enable_memory_types else {},
                'temporal_connections': {str(k): v for k, v in self.temporal_connections.items()} if self.enable_temporal_connections else {},
                'forgetting_curves': {str(k): v for k, v in self.forgetting_curves.items()} if self.enable_consolidation else {},
                'memory_hierarchies': {str(k): v for k, v in self.memory_hierarchies.items()} if self.enable_memory_types else {},
                'agi_features': {
                    'enable_memory_types': self.enable_memory_types,
                    'enable_temporal_connections': self.enable_temporal_connections,
                    'enable_consolidation': self.enable_consolidation
                }
            }
            
            # Save as gzip-compressed JSON for maximum compression (ONLY compressed version)
            json_str = json.dumps(save_data, separators=(',', ':'))  # No spaces, minimal JSON
            json_bytes = json_str.encode('utf-8')
            
            # Ensure graph_file always ends with .gz
            if not self.graph_file.endswith('.gz'):
                self.graph_file = self.graph_file + '.gz'
            
            # Save ONLY compressed version with maximum compression level
            with gzip.open(self.graph_file, 'wt', compresslevel=9) as f:  # Maximum compression
                f.write(json_str)
            
            # Clean up old uncompressed file if it exists (migration cleanup)
            old_uncompressed = self.graph_file.replace('.gz', '')
            if old_uncompressed.endswith('.json') and os.path.exists(old_uncompressed):
                try:
                    os.remove(old_uncompressed)
                    print(f"[CLEANUP] Removed old uncompressed file: {old_uncompressed}")
                except Exception as e:
                    print(f"[WARN] Could not remove old uncompressed file {old_uncompressed}: {e}")
            
            # Calculate size reduction
            uncompressed_size = len(json_bytes)
            compressed_size = os.path.getsize(self.graph_file) if os.path.exists(self.graph_file) else uncompressed_size
            compression_ratio = uncompressed_size / max(compressed_size, 1)
            print(f"[SAVE] MemoryGraph saved (compressed only). Size: {uncompressed_size/1024/1024:.2f}MB -> {compressed_size/1024/1024:.2f}MB ({compression_ratio:.1f}x compression)")
        except Exception as e:
            print(f"[FAIL] Save error: {e}")

    def load_graph(self):
        try:
            # Ensure we're looking for .gz file
            graph_file_gz = self.graph_file if self.graph_file.endswith('.gz') else self.graph_file + '.gz'
            
            # Load ONLY compressed version (fast and efficient)
            if os.path.exists(graph_file_gz):
                with gzip.open(graph_file_gz, 'rt') as f:
                    d = json.load(f)
                print(f"[LOAD] Loaded compressed MemoryGraph from {graph_file_gz}")
            else:
                # Check for old uncompressed version (one-time migration)
                old_file = graph_file_gz.replace('.gz', '')
                if old_file.endswith('.json') and os.path.exists(old_file):
                    print(f"[MIGRATE] Found old uncompressed file {old_file}, loading for migration...")
                    with open(old_file, 'r') as f:
                        d = json.load(f)
                    print(f"[LOAD] Loaded MemoryGraph from {old_file} (will be saved as compressed on next save)")
                    # Update graph_file to use .gz for future saves
                    self.graph_file = graph_file_gz
                else:
                    print(f"[WARN] MemoryGraph file not found: {graph_file_gz}")
                    return
            self.graph = json_graph.node_link_graph(d['graph_data'], directed=False, edges="links")
            self.id_counter = d.get('id_counter', 0)
            self.node_queue = deque(d.get('node_queue', []))
            self.access_counts = d.get('access_counts', {})
            self.last_accessed = d.get('last_accessed', {})
            # Load importance scores, initialize to 0.5 for nodes without scores (backward compatibility)
            loaded_scores = d.get('importance_scores', {})
            self.importance_scores = {}
            for nid in self.graph.nodes():
                self.importance_scores[nid] = loaded_scores.get(str(nid), loaded_scores.get(nid, 0.5))
            
            info = d.get('compression_info', {})
            self.original_dim = info.get('original_dim')
            comp_dim = info.get('compressed_dim')
            
            # Load AGI Features (backward compatible)
            agi_features = d.get('agi_features', {})
            self.enable_memory_types = agi_features.get('enable_memory_types', self.enable_memory_types)
            self.enable_temporal_connections = agi_features.get('enable_temporal_connections', self.enable_temporal_connections)
            self.enable_consolidation = agi_features.get('enable_consolidation', self.enable_consolidation)
            self.use_combined_compression = info.get('use_combined_compression', self.use_combined_compression)
            
            # Load memory types (backward compatible)
            if self.enable_memory_types:
                loaded_types = d.get('node_types', {})
                self.node_types = {int(k) if isinstance(k, str) else k: v for k, v in loaded_types.items()}
            else:
                self.node_types = {}
            
            # Load temporal connections (backward compatible)
            if self.enable_temporal_connections:
                loaded_temporal = d.get('temporal_connections', {})
                self.temporal_connections = {int(k) if isinstance(k, str) else k: v for k, v in loaded_temporal.items()}
            else:
                self.temporal_connections = {}
            
            # Load forgetting curves (backward compatible)
            if self.enable_consolidation:
                loaded_curves = d.get('forgetting_curves', {})
                self.forgetting_curves = {int(k) if isinstance(k, str) else k: v for k, v in loaded_curves.items()}
            else:
                self.forgetting_curves = {}
            
            # Load memory hierarchies (backward compatible)
            if self.enable_memory_types:
                loaded_hierarchies = d.get('memory_hierarchies', {})
                self.memory_hierarchies = {int(k) if isinstance(k, str) else k: v for k, v in loaded_hierarchies.items()}
            else:
                self.memory_hierarchies = {}
            
            # Load VAE compressor if it was used
            if info.get('use_vae_compression', False):
                self.use_vae_compression = True
                self.vae_latent_dim = info.get('vae_latent_dim', 128)
                
                # Load VAE training settings (backward compatible)
                self.enable_vae_training = info.get('vae_enable_training', self.enable_vae_training)
                self.vae_learning_rate = info.get('vae_learning_rate', self.vae_learning_rate)
                
                # Try to load VAE weights from separate file (new method)
                vae_weights_file = info.get('vae_weights_file')
                if vae_weights_file and os.path.exists(vae_weights_file):
                    try:
                        from neuro_fusion import VAECompressor
                        if self.original_dim:
                            hidden_dim = self.original_dim * 2
                            self.vae_compressor = VAECompressor(
                                input_dim=self.original_dim,
                                hidden_dim=hidden_dim,
                                latent_dim=self.vae_latent_dim,
                                shared_model=self.shared_model,
                                shared_tokenizer=self.shared_tokenizer,
                                enable_training=self.enable_vae_training,
                                learning_rate=self.vae_learning_rate,
                                device=self.device
                            )
                            if self.vae_compressor.load_weights(vae_weights_file, strict=False):
                                print(f"[CONFIG] VAE compressor loaded from {vae_weights_file}")
                                self.vae_weights_path = vae_weights_file
                            else:
                                print(f"[WARN] Could not load VAE weights, will re-initialize")
                                self.vae_compressor = None
                    except Exception as e:
                        print(f"[WARN] Error loading VAE compressor: {e}. Will re-initialize on first use.")
                        import traceback
                        traceback.print_exc()
                        self.vae_compressor = None
                else:
                    # Fallback: Try to load from old JSON format (backward compatibility)
                    vae_state = info.get('vae_state_dict')
                    if vae_state and self.original_dim:
                        try:
                            from neuro_fusion import VAECompressor
                            hidden_dim = self.original_dim * 2
                            self.vae_compressor = VAECompressor(
                                input_dim=self.original_dim,
                                hidden_dim=hidden_dim,
                                latent_dim=self.vae_latent_dim,
                                shared_model=self.shared_model,
                                shared_tokenizer=self.shared_tokenizer,
                                enable_training=self.enable_vae_training,
                                learning_rate=self.vae_learning_rate,
                                device=self.device
                            )
                            # Convert list back to tensors (old format)
                            vae_state_tensors = {}
                            for k, v in vae_state.items():
                                if isinstance(v, list):
                                    # Handle nested lists (2D tensors)
                                    if isinstance(v[0], list):
                                        vae_state_tensors[k] = torch.tensor(v, device=self.device)
                                    else:
                                        vae_state_tensors[k] = torch.tensor(v, device=self.device)
                                else:
                                    vae_state_tensors[k] = v
                            self.vae_compressor.load_state_dict(vae_state_tensors, strict=False)
                            if self.enable_vae_training:
                                self.vae_compressor.train()
                            else:
                                self.vae_compressor.eval()
                            print("[CONFIG] VAE compressor loaded from saved state (old format).")
                        except Exception as e:
                            print(f"[WARN] Error loading VAE compressor from old format: {e}. Will re-initialize on first use.")
                            self.vae_compressor = None
                
                # Load PCA on VAE if it was used (AGI feature)
                if self.use_combined_compression and info.get('pca_on_vae_dim'):
                    try:
                        pca_dim = info.get('pca_on_vae_dim')
                        pca_components = info.get('pca_on_vae_components')
                        pca_mean = info.get('pca_on_vae_mean')
                        if pca_components and pca_mean:
                            # Create a dummy dataset to properly fit PCA (required for sklearn)
                            dummy_data = np.random.randn(10, len(pca_mean))
                            self.pca_on_vae = PCA(n_components=pca_dim)
                            self.pca_on_vae.fit(dummy_data)
                            # Now set the actual components and mean
                            self.pca_on_vae.components_ = np.array(pca_components)
                            self.pca_on_vae.mean_ = np.array(pca_mean)
                            self.pca_on_vae.n_components_ = pca_dim
                            # Set other required attributes
                            if not hasattr(self.pca_on_vae, 'explained_variance_'):
                                self.pca_on_vae.explained_variance_ = np.var(self.pca_on_vae.components_, axis=0)
                            if not hasattr(self.pca_on_vae, 'explained_variance_ratio_'):
                                total_var = np.sum(self.pca_on_vae.explained_variance_)
                                self.pca_on_vae.explained_variance_ratio_ = self.pca_on_vae.explained_variance_ / total_var if total_var > 0 else np.ones(pca_dim) / pca_dim
                            print(f"[AGI] PCA on VAE loaded: {pca_dim} components")
                    except Exception as e:
                        print(f"[WARN] Error loading PCA on VAE: {e}")
                        import traceback
                        traceback.print_exc()
                        self.pca_on_vae = None
            
            # Load PCA if it was used (and VAE is not)
            if self.original_dim and comp_dim and not self.use_vae_compression:
                self.pca = PCA(n_components=comp_dim)
                if info.get('pca_components') is not None and info.get('pca_mean') is not None:
                    try:
                        self.pca.components_ = np.array(info['pca_components'])
                        self.pca.mean_ = np.array(info['pca_mean'])
                        print("[CONFIG] PCA components and mean loaded.")
                    except ValueError as ve:
                        print(f"[WARN] Error loading PCA components/mean (shape mismatch?): {ve}. PCA might need re-fitting.")
                        self.pca = None
                else:
                    print("[WARN] PCA components or mean not found in saved data. PCA might need re-fitting.")
                    self.pca = None  # Reset PCA if components not fully loaded

            self.quantization_bits = info.get('quantization_bits', self.quantization_bits)
            
            # CRITICAL: Don't restore max_nodes from saved file if current setting is None (infinite growth)
            # This ensures infinite growth is preserved even if old graph had a limit
            saved_max_nodes = info.get('max_nodes')
            if self.max_nodes is None:
                # If we're in infinite growth mode, never restore max_nodes from saved file
                if saved_max_nodes is not None:
                    print(f"[CONFIG] Preserving infinite growth mode (ignoring saved max_nodes={saved_max_nodes})")
                # Keep self.max_nodes as None (infinite growth)
            elif saved_max_nodes is not None and self.max_nodes is not None:
                # Only restore if both are set and we want to use saved value
                # But for now, prefer current setting
                pass
            
            print("[NETWORK] MemoryGraph loaded.")

            # Backwards compatibility: ensure all embeddings are in dict format {data:[], compressed: bool}
            for nid, node_data in self.graph.nodes(data=True):
                current_embedding = node_data.get('embedding')
                if isinstance(current_embedding, list):  # Old format: list of floats
                    self.graph.nodes[nid]['embedding'] = {'data': current_embedding, 'compressed': False}
                elif isinstance(current_embedding, dict) and 'data' not in current_embedding:
                    # Old dict format: might be {'embedding_key': [data]}
                    if current_embedding:  # if not empty dict
                        self.graph.nodes[nid]['embedding'] = {'data': list(current_embedding.values())[0], 'compressed': False}
                    else:  # empty dict
                        self.graph.nodes[nid]['embedding'] = {'data': [], 'compressed': False}
                elif not isinstance(current_embedding, dict):  # Other unexpected formats (e.g. None)
                    print(f"Warning: Node {nid} embedding format unexpected during load. Setting to empty dict.")
                    self.graph.nodes[nid]['embedding'] = {'data': [], 'compressed': False}

                # Ensure metadata items are standard Python types (not numpy types)
                metadata_to_clean = node_data.get('metadata', {})
                cleaned_metadata = {}
                for k, v in metadata_to_clean.items():
                    if isinstance(v, (np.ndarray, np.generic)):
                        cleaned_metadata[k] = v.tolist() if isinstance(v, np.ndarray) else v.item()
                    else:
                        cleaned_metadata[k] = v
                self.graph.nodes[nid]['metadata'] = cleaned_metadata

        except Exception as e:
            print(f"[FAIL] Load error: {e}")
            # If load fails, re-initialize an empty graph
            self.graph = nx.Graph()
            self.id_counter = 0
            self.node_queue = deque()
            self.pca = None
            self.original_dim = None

    def get_memory_summary(self):
        try:
            total_bytes = 0
            for _, n in self.graph.nodes(data=True):
                # Ensure embedding is in the expected dictionary format
                emb_data = n.get('embedding')
                if isinstance(emb_data, dict) and 'data' in emb_data:
                    # Approximate size based on data list length and assumed bit depth
                    item_size = emb_data.get('bits', 16) // 8 if emb_data.get('bits', 16) < 32 else 4  # float32 is 4 bytes
                    total_bytes += len(emb_data['data']) * item_size
                elif isinstance(emb_data, list):  # Handle old list format
                    total_bytes += len(emb_data) * 4  # Assume float32
                else:
                    print(f"Warning: Unexpected embedding format encountered for size calculation: {type(emb_data)}")

            return {
                'node_count': self.graph.number_of_nodes(),
                'edge_count': self.graph.number_of_edges(),
                'estimated_memory_mb': round(total_bytes / (1024**2), 2)
            }
        except Exception as e:
            return {'error': str(e)}

    def analyze_memory(self):
        return self.get_memory_summary()