# flake8: noqa
"""
memory.py
---------
Core memory system for HIEROMEM.

Now extended with integrated:
- EventLog (for structured event emission)
- Narrator (for natural-language cognitive reporting)
- GovernanceSafetyRegion (λ, τ, θ stability enforcement)
- SQLiteStore persistence
- GraphStore support (First-Class Graph Capability)

Implements HIEROMEM’s four core dynamics:
  1. Insertion      — encode & store new experiences
  2. Retrieval      — resonance-based recall (soft attention)
  3. Summarisation  — merge episodes to ensure bounded memory
  4. Forgetting     — trust decay & pruning (θ threshold)

"""

from __future__ import annotations

from collections import defaultdict
from dataclasses import dataclass
from contextlib import contextmanager, AbstractContextManager
from datetime import datetime
from typing import (
    Any,
    Iterable,
    Iterator,
    Literal,
    Mapping,
    Optional,
    Protocol,
    Sequence,
    Dict,
    List,
    TYPE_CHECKING,
)

if TYPE_CHECKING:
    from hieromem_core.replay.simulation import SimulationMemory

import numpy as np
import networkx as nx

from hieromem_core.episode import Episode
from hieromem_core.events import MemoryEvent, EventType, EventLog
from hieromem_core.compression import CompressionStrategy, NoOpCompressionStrategy
from hieromem_core.encoder import Encoder, EncoderConfig, MultiEncoderRegistry
from hieromem_core.index_manager import IndexManager
from hieromem_core.index_routing import IndexRouter
from hieromem_core.llm_extractor import LLMExtractor
from hieromem_core.graph_extractor import GraphExtractor, AutoGraphExtractor
from hieromem_core.persistence.memory_store import MemoryStore
from hieromem_core.persistence.graph_store import GraphStore
from hieromem_core.graph_extractor import GraphExtractor
from hieromem_core.governance import (
    GovernanceSafetyRegion,
    DEFAULT_FORBIDDEN_RELATIONS,
)
from hieromem_core.ingestion_queue import IngestionQueue
from hieromem_core.narrator import Narrator
from hieromem_core.realms import (
    MemoryRealm,
    RealmParticipant,
    check_permission,
    enforce_permission,
    PermissionError
)
from hieromem_core.multimodal import MultiModalEncoder
from hieromem_core.symbolic import GraphReasoner
from hieromem_core.retrieval import (
    DualRetriever,
    LearnedRetriever,
    DualRetrievalConfig,
    DualRetrievalInput,
    DualRetrievalResult,
)
from hieromem_core.indexing.entity_index import EntityIndex
from hieromem_core.graph_merge import hierarchical_merge
from hieromem_core.stability import GlobalStabilityMonitor
from hieromem_core.tracing import start_span
from hieromem_core.edge_types import (
    ConversationalEdgeBuilder,
    TURN_FOLLOWS, SPOKEN_BY, IN_SESSION,
)
from hieromem_core.policy.evaluator import MemoryPolicyEvaluator
from hieromem_core.policy.schema import MemoryPolicy


class VectorIndex(Protocol):
    @property
    def ids(self) -> Sequence[str]: ...

    def add_embeddings(
    self,
    ids: Sequence[str],
     matrix: np.ndarray) -> None: ...

    def add_many(self, ids: Sequence[str], matrix: np.ndarray) -> None: ...

    def add(self, episode_id: str, vec: np.ndarray) -> None: ...

    def remove(self, episode_id: str) -> None: ...

    def search(self,
    query: np.ndarray,
    k: int) -> Sequence[tuple[str,
     float]]: ...

    def clear(self) -> None: ...


# --------------------------------------------------------------------------- #
# Memory Class
# --------------------------------------------------------------------------- #
def _stringify_meta(value: Any) -> str:
    if isinstance(value, (str, int, float, bool)) or value is None:
        return repr(value)
    if isinstance(value, (list, tuple, set)):
        return "[" + ", ".join(_stringify_meta(v) for v in value) + "]"
    if isinstance(value, dict):
        items = ", ".join(
    f"{k}:{_stringify_meta(v)}" for k,
    v in sorted(
        value.items()))
        return "{" + items + "}"
    return repr(value)


@dataclass
class BatchSymbolicSearchInput:
    """Entry describing a single symbolic search in a batch."""

    path: Optional[Sequence[str]] = None
    pattern: Optional[nx.DiGraph] = None
    namespace: Optional[str] = None
    limit: Optional[int] = None


class Memory:
    def __init__(
        self,
        Nmax: int = 20,
        lam: float = 0.95,
        tau: float = 0.1,
        theta: float = 0.1,
        b: int = 2,
        store: Optional[MemoryStore] = None,
        debug: bool = False,
        event_log: Optional[EventLog] = None,
        narrator: Optional[Narrator] = None,
        encoder: Optional[Encoder] = None,
        encoder_registry: Optional[MultiEncoderRegistry] = None,
        multi_modal_encoder: Optional[MultiModalEncoder] = None,
        symbolic_reasoner: Optional[GraphReasoner] = None,
        stability_monitor: Optional[GlobalStabilityMonitor] = None,
        *,
        enable_multi_tenant: bool = False,
        default_namespace: str = "default",
        summary_distortion_bound: float | None = None,
        forbidden_relations: Optional[Iterable[str]] = None,
        index_router: Optional[IndexRouter] = None,
        index_manager: Optional[IndexManager] = None,
        ingestion_queue: Optional[IngestionQueue] = None,
        compression_strategy: Optional[CompressionStrategy] = None,
        async_ingestion: bool = False,
        tenant_configs: Optional[Dict[str, Any]] = None,
        strict_gsr: bool = True,
        auto_graph_extraction: bool = False,
        graph_extraction_backend: str = "regex",
        extractor: Optional[Any] = None,
        extractor_model: Optional[str] = None,
        graph_policy: str = "use_provided",
        policy: MemoryPolicy | None = None,
    ) -> None:
        """Initialize Memory with optional tenant configuration dict.

        ``tenant_configs`` is a mapping from tenant identifier (or ``default``) to a
        configuration dict containing ``quota`` and ``forbidden_vocab`` among other
        possible overrides. It is injected by ``deps.get_memory_engine`` after the
        TOML file is parsed.
        
        Graph Extraction Policy (graph_policy):
            - "use_provided": Use the graph passed to insert(), even if empty (default)
            - "auto_if_missing": Auto-extract from meta['text'] if graph is None
            - "disable": Never use graphs, ignore graph parameter
        """

        self.Nmax = Nmax
        self.lam = lam
        self.tau = tau
        self.theta = theta
        self.b = b
        self.debug = debug
        self.store = store

        # ============================================================= #
        # INVARIANT: graph_store is NEVER None
        # This ensures symbolic memory is always persistent and queryable
        # ============================================================= #
        from hieromem_core.persistence.graph_store import GraphStore
        from hieromem_core.persistence.in_memory_graph_store import InMemoryGraphStore

        if self.store is not None:
            is_graph_store = isinstance(self.store, GraphStore)

            if is_graph_store:
                # Store IS a GraphStore - use it as the single source of truth
                self._graph_store = self.store
                # Also set store.graph_store = store for internal consistency
                if getattr(self.store, "graph_store", None) is None:
                    try:
                        self.store.graph_store = self.store
                    except AttributeError:
                        pass  # __slots__ without graph_store - ok, we have _graph_store
            else:
                # Regular MemoryStore - check if it has graph_store attached
                gs = getattr(self.store, "graph_store", None)
                if gs is not None:
                    self._graph_store = gs
                else:
                    # No graph_store provided - create default InMemoryGraphStore
                    self._graph_store = InMemoryGraphStore()
                    # Attach to store for consistency if possible
                    try:
                        self.store.graph_store = self._graph_store
                    except AttributeError:
                        pass  # __slots__ without graph_store
        else:
            # No store at all - create default InMemoryGraphStore
            self._graph_store = InMemoryGraphStore()

        if self.debug:
            print(f"[Memory] Graph store: {type(self._graph_store).__name__} (always enabled)")


        self.summary_distortion_bound = summary_distortion_bound
        self.encoder = encoder or Encoder(EncoderConfig())
        self.encoder_registry = encoder_registry or MultiEncoderRegistry()
        if not self.encoder_registry.default():
            self.encoder_registry.register(
    self.encoder.config.model_name, self.encoder)
        self.metrics: dict[str, Any] = {}

        # Phase 4.2: GSR strict mode toggle
        self.strict_gsr = strict_gsr
        self._dropped_edges_count = 0  # Track dropped edges when strict_gsr=False

        # Tenant-aware state
        self._multi_tenant = bool(enable_multi_tenant)
        self.namespace = default_namespace or "default"
        if self._multi_tenant:
            self._episodes_by_namespace: dict[str,
                list[Episode]] = defaultdict(list)
            self._context_by_namespace: dict[str, Optional[np.ndarray]] = defaultdict(
                lambda: None)
        else:
            self._episodes_single: list[Episode] = []
            self._context_single: Optional[np.ndarray] = None
        
        # O(1) episode lookup index (ID -> Episode)
        self._episode_by_id: dict[str, Episode] = {}

        # Optional FAISS/approximate index for retrieval
        self._index: Optional[VectorIndex] = None
        self.index_router = index_router
        if self.index_router and self._index is not None:
            self.index_router.set_default(self._index)

        self.index_manager = index_manager
        if self.index_manager:
            self.index_manager.start()

        self.compression_strategy = compression_strategy or NoOpCompressionStrategy()
    
        # Graph extraction configuration
        self.graph_extraction_backend = graph_extraction_backend
        # Note: _llm_extractor, _graph_extractor, and _extraction_pipeline are 
        # initialized in the ExtractionPipeline block further below (lines ~280-320).

        self.ingestion_queue = ingestion_queue
        self._async_ingestion = bool(async_ingestion and ingestion_queue)
        if self._async_ingestion and self.ingestion_queue:
            self.ingestion_queue.start()

        # Event + narrator integration
        self.event_log = event_log or EventLog()
        self.narrator = narrator or Narrator()
        self.policy = policy
        self.policy_evaluator = (
            MemoryPolicyEvaluator(policy, event_log=self.event_log)
            if policy is not None
            else None
        )
        self.multi_modal_encoder = multi_modal_encoder or MultiModalEncoder()
        self.symbolic_reasoner = symbolic_reasoner or GraphReasoner()
        self.stability_monitor = stability_monitor or GlobalStabilityMonitor()
        if forbidden_relations is None:
            forbidden_source = DEFAULT_FORBIDDEN_RELATIONS
        else:
            forbidden_source = tuple(forbidden_relations)
        forbidden = tuple(
            {str(rel).lower() for rel in forbidden_source if str(rel).strip()}
        )
        self._forbidden_relations = forbidden
        self._forbidden_relations = forbidden
        self.entity_index = EntityIndex()
        self.dual_retriever = LearnedRetriever(
            self, reasoner=self.symbolic_reasoner)
        self._revision = 0

        # Conversational edge state tracking (TURN_FOLLOWS, SPOKEN_BY, IN_SESSION)
        self._previous_episode_id: Optional[str] = None
        self._current_session_id: Optional[str] = None
        self._turn_counter: int = 0
        self._conversation_state_by_namespace: dict[str, dict[str, Any]] = defaultdict(
            lambda: {
                "previous_episode_id": None,
                "current_session_id": None,
                "turn_counter": 0,
            }
        )
        # Lock for thread-safe conversation state updates
        import threading
        self._conversation_lock = threading.Lock()

        # Auto graph extraction for RAG workloads
        self._auto_graph_extraction = auto_graph_extraction
        self._graph_policy = graph_policy  # "use_provided", "auto_if_missing", "disable"
        self._extraction_pipeline = None
        # Legacy extractors (for backward compatibility)
        self._graph_extractor = None
        self._llm_extractor = None
        
        if extractor:
            # Use provided extractor in a new pipeline
            self._auto_graph_extraction = True
            from hieromem_core.extraction import ExtractionPipeline
            from hieromem_core.extraction.normalizer import EntityNormalizer
            from hieromem_core.extraction.linker import EntityLinker
            self._extraction_pipeline = ExtractionPipeline(
                extractors=[extractor],
                normalizer=EntityNormalizer(),
                linker=EntityLinker(),
            )
            if self.debug:
                 print(f"[Memory] Using injected extractor: {type(extractor).__name__}")
        elif auto_graph_extraction:
            # Use new ExtractionPipeline with always-on normalization
            from hieromem_core.extraction import (
                ExtractionPipeline,
                create_regex_pipeline,
                create_llm_pipeline,
                create_hybrid_pipeline,
            )
            
            if graph_extraction_backend in ("llm", "openai", "anthropic"):
                # Use LLM-powered extraction for higher quality
                self._extraction_pipeline = create_llm_pipeline(provider=graph_extraction_backend, api_key=None)
                if self.debug:
                    print(f"[Memory] Initialized LLM ExtractionPipeline ({graph_extraction_backend})")
            else:
                # Default to regex (fast, deterministic)
                self._extraction_pipeline = create_regex_pipeline()
                if self.debug:
                    print("[Memory] Initialized Regex ExtractionPipeline")

        # Phase 4: Realm Registry
        self.realms: dict[str, MemoryRealm] = {}


        # Connect narrator to event log automatically
        self.event_log.subscribe(lambda evt: self.narrator.observe(evt))

        # Governance & safety region (emit into event log)
        def _gsr_emitter(event_type: str, message: str,
                         details: dict[str, Any]) -> None:
            self.event_log.add(
    EventType.GOVERNANCE,
    message,
    details=details,
     source="GSR")

        gsr_forbidden = forbidden if forbidden else DEFAULT_FORBIDDEN_RELATIONS
        self.gsr = GovernanceSafetyRegion(
            lam=self.lam,
            tau=self.tau,
            theta=self.theta,
            Nmax=self.Nmax,
            b=self.b,
            clamp_on_validate=True,
            emitter=_gsr_emitter,
            forbidden_relations=gsr_forbidden,
        )
        self.gsr.validate()

        # Load episodes from persistent store
        if self.store is not None:
            self._load_from_store()

        self._emit(
            EventType.SYSTEM,
            "Memory initialized.",
            {
                "count": len(self.episodes),
                "namespace": self.namespace,
                "multi_tenant": self._multi_tenant,
                "graph_enabled": self.graph_store is not None
            },
        )

    def create_realm(
        self,
        name: str,
        realm_type: Literal["shared", "private", "hybrid"] = "private",
        participants: Optional[List[Dict[str, Any]]] = None,
        isolation_level: Literal["logical", "cryptographic"] = "logical"
    ) -> MemoryRealm:
        """Create a new memory realm.
        
        Args:
            name: Realm name (must be unique)
            realm_type: 'shared', 'private', or 'hybrid'
            participants: List of dicts with 'agent_id' and 'permissions'
            isolation_level: 'logical' (namespace) or 'cryptographic' (future)
        """
        if name in self.realms:
            raise ValueError(f"Realm '{name}' already exists")
            
        parts = []
        if participants:
            for p in participants:
                parts.append(RealmParticipant(
                    agent_id=p["agent_id"], 
                    permissions=p.get("permissions", ["read"])
                ))
                
        realm = MemoryRealm(
            name=name,
            type=realm_type,
            participants=parts,
            isolation_level=isolation_level
        )
        self.realms[name] = realm
        if self.debug:
            print(f"[Memory] Created realm '{name}' ({realm_type})")
        return realm

    # ------------------------------------------------------------------ #
    # Factory Methods
    # ------------------------------------------------------------------ #
    @classmethod
    def for_rag(
        cls,
        Nmax: int = 1000,
        encoder: Optional[Encoder] = None,
        store: Optional[MemoryStore] = None,
        *,
        extraction_backend: str = "regex",
        debug: bool = False,
        **kwargs: Any,
    ) -> "Memory":
        """
        Create a Memory instance optimized for RAG workloads.

        This factory method sets sensible defaults for retrieval-augmented generation:
        - Auto graph extraction enabled
        - graph_policy = "auto_if_missing" (extracts from text automatically)
        - Default encoder created if not provided

        Args:
            Nmax: Maximum episodes before summarization (default 1000)
            encoder: Encoder instance (created if not provided)
            store: Persistence store (optional)
            extraction_backend: "regex" (fast, offline) or "llm" (higher quality)
            debug: Enable debug logging
            **kwargs: Additional arguments passed to Memory.__init__

        Returns:
            Memory instance configured for RAG

        Example:
            >>> memory = Memory.for_rag()
            >>> memory.insert(vector, meta={'text': 'Alice knows Bob'})  # Auto-extracts!
            >>> results = memory.retrieve("Who does Alice know?")
        """
        if encoder is None:
            encoder = Encoder(EncoderConfig())

        return cls(
            Nmax=Nmax,
            encoder=encoder,
            store=store,
            debug=debug,
            auto_graph_extraction=True,
            graph_extraction_backend=extraction_backend,
            graph_policy="auto_if_missing",
            **kwargs,
        )

    # ------------------------------------------------------------------ #
    # Tenant Management
    # ------------------------------------------------------------------ #
    @property
    def episodes(self) -> list[Episode]:
        if self._multi_tenant:
            return self._episodes_by_namespace[self.namespace]
        return self._episodes_single

    @episodes.setter
    def episodes(self, value: list[Episode]) -> None:
        if self._multi_tenant:
            self._episodes_by_namespace[self.namespace] = value
        else:
            self._episodes_single = value

    @property
    def context(self) -> Optional[np.ndarray]:
        if self._multi_tenant:
            return self._context_by_namespace[self.namespace]
        return self._context_single

    @context.setter
    def context(self, value: Optional[np.ndarray]) -> None:
        if self._multi_tenant:
            self._context_by_namespace[self.namespace] = value
        else:
            self._context_single = value

    @property
    def index(self) -> Optional[VectorIndex]:
        if self.index_router:
            return getattr(self.index_router, 'get_default', lambda: None)()
        return self._index

    @property
    def graph_store(self) -> "GraphStore":
        """
        Return the graph store - NEVER None (Phase 1 invariant).

        The graph store is the single source of truth for symbolic memory.
        """
        return self._graph_store  # type: ignore[return-value]

    @graph_store.setter
    def graph_store(self, gs: "GraphStore") -> None:
        """
        Set graph store. Warns if setting to None (deprecated behavior).
        """
        if gs is None:
            import warnings
            warnings.warn(
                "Setting graph_store to None is deprecated. Graph memory is now always enabled.",
                DeprecationWarning,
                stacklevel=2
            )
            # Don't actually set to None - keep existing
            return
        self._graph_store = gs
        # Also update store.graph_store for consistency if possible
        if self.store is not None:
            try:
                self.store.graph_store = gs
            except AttributeError:
                pass  # Store uses __slots__

    def set_namespace(self, namespace: str) -> None:
        if not self._multi_tenant:
            return
        self.namespace = namespace

    @contextmanager
    def use_namespace(self, namespace: Optional[str]) -> Iterator[None]:
        if not self._multi_tenant or namespace is None:
            yield
            return
        previous = self.namespace
        self.set_namespace(namespace)
        try:
            yield
        finally:
            self.set_namespace(previous)

    def _bump_revision(self) -> None:
        self._revision += 1

    @property
    def revision(self) -> int:
        """Expose the monotonic revision counter for cache invalidation."""

        return self._revision

    def _default_encoder_id(self) -> str:
        """Return the default encoder identifier for this memory instance."""
        if hasattr(self, 'encoder') and hasattr(self.encoder, 'config'):
            return getattr(self.encoder.config, 'model_name', 'default')
        return 'default'
    # ------------------------------------------------------------------ #
    # Core Operations
    # ------------------------------------------------------------------ #
    def _validate_graph_safety(self, G: nx.DiGraph) -> nx.DiGraph:
        """
        Filter graph edges against GSR, dropping forbidden edges.

        Phase 3 Invariant: GSR filters graphs, it never prevents memory formation.
        This ensures that once something is accepted as memory, it persists.

        Legacy behavior (strict_gsr=True raises) is deprecated and will be
        removed in v2.0. For now, strict_gsr=True still raises for backwards
        compatibility, but strict_gsr=False is strongly recommended.
        """
        if not hasattr(self, 'gsr') or self.gsr is None or G is None:
            return G

        # Use the new filter_graph method for consistent filtering
        filtered_graph, dropped_edges = self.gsr.filter_graph(G)

        if dropped_edges:
            self._dropped_edges_count += len(dropped_edges)

            if self.strict_gsr:
                # Legacy behavior: raise on forbidden edges
                # DEPRECATED: will be removed in v2.0
                first_dropped = dropped_edges[0]
                raise ValueError(
                    f"Relation '{first_dropped['relation']}' between "
                    f"'{first_dropped['source']}' and '{first_dropped['target']}' "
                    f"is forbidden by GSR."
                )

            # Preferred behavior: filter and emit event
            self._emit(
                EventType.GOVERNANCE,
                f"Filtered {len(dropped_edges)} forbidden edges from graph.",
                {"dropped": dropped_edges, "namespace": self.namespace},
            )

        return filtered_graph

    def insert(
        self,
        vector: np.ndarray,
        graph: Optional[nx.DiGraph] = None,
        meta: dict[str, Any] | None = None,
        *,
        namespace: Optional[str] = None,
        realm: Optional[str] = None,
        agent: Optional[str] = None,
    ) -> Episode:
        """Insert a new episode into memory.
        
        Graph behavior depends on graph_policy (set in __init__):
        - "use_provided": Use graph as passed (default, backward compatible)
        - "auto_if_missing": If graph is None, auto-extract from meta['text']
        - "disable": Ignore graphs entirely
        """
        # Phase 4: Realm enforcement
        if realm:
            if realm not in self.realms:
                raise ValueError(f"Realm '{realm}' does not exist")
            if agent:
                enforce_permission(self.realms[realm], agent, "write")
            namespace = realm

        with start_span("memory.insert", attributes={"namespace": namespace or self.namespace}):
            with self.use_namespace(namespace):
                # ================================================================ #
                # Graph Policy: Handle auto-extraction based on policy
                # ================================================================ #
                if self._graph_policy == "disable":
                    # Never use graphs
                    graph = None
                elif self._graph_policy == "auto_if_missing" and graph is None:
                    # Auto-extract from text if available
                    text = None
                    if meta:
                        text = meta.get('text') or meta.get('content') or meta.get('content_preview')
                    
                    if text and self._extraction_pipeline:
                        try:
                            result = self._extraction_pipeline.extract(str(text), link_entities=False)
                            graph = result.graph if result.graph else nx.DiGraph()
                            if self.debug and graph.number_of_nodes() > 0:
                                print(f"[Memory] Auto-extracted graph: {graph.number_of_nodes()} nodes")
                        except Exception as e:
                            if self.debug:
                                print(f"[Memory] Auto-extraction failed: {e}")
                            graph = nx.DiGraph()
                    else:
                        graph = nx.DiGraph()
                # else: "use_provided" - use graph as passed
                
                # Validate and potentially filter graph edges (returns modified graph)
                if graph is not None:
                    graph = self._validate_graph_safety(graph)

                if self.policy_evaluator is not None:
                    decision = self.policy_evaluator.evaluate_insert(
                        graph=graph,
                        total_episodes=len(self.episodes),
                        namespace=self.namespace,
                    )
                    if not decision.allowed:
                        raise ValueError(
                            "Insert blocked by policy: "
                            + "; ".join(decision.reasons)
                        )

                # Phase 4: Stamp namespace into metadata for portability
                effective_meta = (meta or {}).copy()
                if (namespace or self.namespace) and "namespace" not in effective_meta:
                    effective_meta["namespace"] = namespace or self.namespace

                ep = Episode(
                    e=np.asarray(vector, dtype=np.float32),
                    G=graph,
                    tocc=datetime.now(),
                    meta=effective_meta,
                )
                self.episodes.append(ep)
                self._episode_by_id[ep.id] = ep  # O(1) lookup index

                episodes_to_persist = [ep]
                # Pass explicit namespace to ensure graph edges are correctly namespaced
                effective_ns = namespace or self.namespace
                if self._async_ingestion and self.ingestion_queue:
                    self.ingestion_queue.enqueue(ep)
                else:
                    self._persist_batch(episodes_to_persist, namespace=effective_ns)

                # ================================================================ #
                # Create Conversational Edges (TURN_FOLLOWS, SPOKEN_BY, IN_SESSION)
                # ================================================================ #
                self._create_conversational_edges(ep, meta or {}, namespace=effective_ns)

                self._emit(
                    EventType.INSERT,
                    f"Inserted episode {ep.id}",
                    {"meta": meta, "namespace": self.namespace},
                )

                # Trigger summarisation if policy or capacity requires it
                if (
                    self.policy_evaluator is not None
                    and self.policy_evaluator.should_summarise(self.episodes)
                ):
                    self.summarise()
                if len(self.episodes) > self.Nmax:
                    self.summarise()

                # Bug 5 fix: use explicit None check to avoid boolean trap with empty graphs
                if graph is not None:
                    self.entity_index.add(ep)
                    # Note: graph_store.save_episode is handled by _persist_batch

                self._bump_revision()

                return ep

    def _create_conversational_edges(
        self,
        episode: Episode,
        meta: dict[str, Any],
        *,
        namespace: Optional[str] = None
    ) -> None:
        """
        Create structural edges for conversational memory.
        
        Creates:
        - TURN_FOLLOWS: Links to previous episode (temporal ordering)
        - SPOKEN_BY: Links to speaker node (speaker attribution)
        - IN_SESSION: Links to session node (session membership)
        
        All edges include weights for retrieval ranking boost.
        Thread-safe: uses _conversation_lock for state updates.
        """
        gs = self.graph_store
        if gs is None:
            return
        
        effective_ns = namespace or self.namespace
        
        # Thread-safe access to conversation state
        with self._conversation_lock:
            state = self._conversation_state_by_namespace[effective_ns]
            # 1. TURN_FOLLOWS edge (temporal ordering)
            if state["previous_episode_id"] is not None:
                state["turn_counter"] += 1
                edge = ConversationalEdgeBuilder.build_turn_follows_edge(
                    prev_episode_id=state["previous_episode_id"],
                    curr_episode_id=episode.id,
                    temporal_distance=1,  # Consecutive turns
                    weight_decay=0.9,
                    namespace=effective_ns,
                )
                gs.add_structural_edge(edge[0], edge[1], edge[2], namespace=effective_ns)
                
                if self.debug:
                    print(f"[Memory] Created TURN_FOLLOWS: {edge[0][:8]} -> {edge[1][:8]}")
            
            # Update previous episode tracker
            state["previous_episode_id"] = episode.id

            # 3. IN_SESSION edge (session membership) - needs state tracking
            session_id = meta.get("session_id") or meta.get("_session")
            if session_id:
                session_id = str(session_id).strip()

                # Track current session
                if state["current_session_id"] != session_id:
                    state["current_session_id"] = session_id
                    state["turn_counter"] = 0  # Reset turn counter for new session

                edge = ConversationalEdgeBuilder.build_in_session_edge(
                    episode_id=episode.id,
                    session_id=session_id,
                    turn_index=meta.get("turn_idx", state["turn_counter"]),
                    weight=1.0,
                    session_timestamp=meta.get("session_timestamp"),
                    namespace=effective_ns,
                )
                gs.add_structural_edge(edge[0], edge[1], edge[2], namespace=effective_ns)
        
        # 2. SPOKEN_BY edge (speaker attribution) - no state needed, outside lock
        speaker = meta.get("speaker") or meta.get("speakers")
        if speaker:
            # Handle both single speaker and list of speakers
            speakers = speaker if isinstance(speaker, (list, tuple)) else [speaker]
            for spk in speakers:
                if spk and str(spk).strip():
                    edge = ConversationalEdgeBuilder.build_spoken_by_edge(
                        episode_id=episode.id,
                        speaker=str(spk).strip(),
                        weight=1.0,
                        namespace=effective_ns,
                    )
                    gs.add_structural_edge(edge[0], edge[1], edge[2], namespace=effective_ns)
                    
                    if self.debug:
                        print(f"[Memory] Created SPOKEN_BY: {episode.id[:8]} -> SPEAKER:{spk}")

    # ------------------------------------------------------------------ #
    # Multi-Modal Episode Methods
    # ------------------------------------------------------------------ #

    def add_text_episode(
        self,
        content: str,
        G: nx.DiGraph | None = None,
        meta: dict[str, Any] | None = None,
        *,
        namespace: Optional[str] = None,
        doc_id: Optional[str] = None,
        auto_extract: Optional[bool] = None,
    ) -> Episode:
        """Add a text episode to memory.

        Args:
            content: Text content to encode
            G: Optional knowledge graph. If None and auto_graph_extraction is enabled,
               automatically extracts entities and relations from the text.
            meta: Optional metadata dictionary
            namespace: Optional namespace for multi-tenant scenarios
            doc_id: Optional document ID for cross-document linking
            auto_extract: Override auto_graph_extraction setting for this call

        Returns:
            The created Episode
        """
        from hieromem_core.encoders import EncoderFactory

        encoder = EncoderFactory.get_encoder('text')
        embedding = encoder.encode(content)

        # Auto-extract graph if enabled and no graph provided
        should_extract = auto_extract if auto_extract is not None else self._auto_graph_extraction
        if G is None and should_extract:
            # Use new ExtractionPipeline (always-on normalization)
            if self._extraction_pipeline is not None:
                result = self._extraction_pipeline.extract(
                    content,
                    episode_id=doc_id,  # For cross-doc linking
                    link_entities=True,
                )
                G = result.graph
                if self.debug:
                    print(f"[Memory] Extracted graph: {G.number_of_nodes()} nodes, "
                          f"{G.number_of_edges()} edges, "
                          f"entities={[e.canonical for e in result.entities[:5]]}")
            # Legacy fallback for backward compatibility
            elif self._llm_extractor is not None:
                legacy_result = self._llm_extractor.extract(content)
                G = legacy_result.graph
            elif self._graph_extractor is not None:
                G = self._graph_extractor.extract(content, doc_id=doc_id)
            else:
                G = nx.DiGraph()
        elif G is None:
            G = nx.DiGraph()

        if meta is None:
            meta = {}
        meta['media_type'] = 'text'
        meta['content_preview'] = content[:100] if len(
            content) > 100 else content
        if doc_id:
            meta['doc_id'] = doc_id

        return self.insert(embedding, G, meta, namespace=namespace)

    def add_reasoning_chain(self, chain: Any, namespace: Optional[str] = None) -> Episode:
        """
        Store a symbolic reasoning chain as a derived episode.
        
        This enables 'meta-reasoning' where past reasoning chains can be
        retrieved and used to answer future queries.
        """
        text = chain.to_text()
        # We merge the chain's metadata with standard tracking
        meta = chain.to_episode().meta
        
        # Use add_text_episode to handle embedding and ID assignment
        # We disable auto-extraction because the chain IS the structure
        return self.add_text_episode(
            content=text, 
            meta=meta, 
            auto_extract=False,
            namespace=namespace
        )

    def add_image_episode(
        self,
        image_path: str,
        G: nx.DiGraph | None = None,
        meta: dict[str, Any] | None = None,
        *,
        namespace: Optional[str] = None,
    ) -> Episode:
        """Add an image episode to memory.

        Args:
            image_path: Path to image file or PIL Image object
            G: Optional knowledge graph. If None, creates empty graph
            meta: Optional metadata dictionary
            namespace: Optional namespace for multi-tenant scenarios

        Returns:
            The created Episode
        """
        from hieromem_core.encoders import EncoderFactory

        encoder = EncoderFactory.get_encoder('image')
        embedding = encoder.encode(image_path)

        if G is None:
            G = nx.DiGraph()

        if meta is None:
            meta = {}
        meta['media_type'] = 'image'
        meta['image_path'] = str(image_path)

        return self.insert(embedding, G, meta, namespace=namespace)

    def add_audio_episode(
        self,
        audio_path: str,
        G: nx.DiGraph | None = None,
        meta: dict[str, Any] | None = None,
        *,
        namespace: Optional[str] = None,
    ) -> Episode:
        """Add an audio episode to memory.

        Args:
            audio_path: Path to audio file or numpy array of audio samples
            G: Optional knowledge graph. If None, creates empty graph
            meta: Optional metadata dictionary
            namespace: Optional namespace for multi-tenant scenarios

        Returns:
            The created Episode
        """
        from hieromem_core.encoders import EncoderFactory

        encoder = EncoderFactory.get_encoder('audio')
        embedding = encoder.encode(audio_path)

        if G is None:
            G = nx.DiGraph()

        if meta is None:
            meta = {}
        meta['media_type'] = 'audio'
        meta['audio_path'] = str(audio_path)

        return self.insert(embedding, G, meta, namespace=namespace)

    # ------------------------------------------------------------------ #
    def update_trust(
        self,
        used_ids: list[str] | None=None,
        *,
        namespace: Optional[str]=None,
    ) -> None:
        """Apply trust decay and prune low-trust episodes."""
        with self.use_namespace(namespace):
            used_set: set[str] = set(used_ids or [])
            pruned_ids: list[str] = []

            # 1. Update trust and identify pruned
            for ep in self.episodes:
                used = ep.id in used_set
                ep.update_trust(used=used, lam=self.lam)
                if ep.trust < self.theta:
                    pruned_ids.append(ep.id)

            # 2. Batch remove from store
            if pruned_ids and self.store is not None:
                try:
                    if hasattr(self.store, "remove_many"):
                        self.store.remove_many(pruned_ids)
                    else:
                        for eid in pruned_ids:
                            self.store.remove(eid)
                except Exception as e:
                    self._emit(EventType.SYSTEM, "Failed to batch remove episodes from store", {"error": str(e)}, level="warn")

            for eid in pruned_ids:
                self._remove_graph_episode(eid, namespace=self.namespace)

            # 3. Remove from active memory and index
            if pruned_ids:
                pruned_set = set(pruned_ids)
                self.episodes = [ep for ep in self.episodes if ep.id not in pruned_set]

                for eid in pruned_ids:
                    self._episode_by_id.pop(eid, None)
                    self._index_remove_episode(eid)
                    self.entity_index.remove(eid)

            for ep in self.episodes:
                self._persist_trust_update(ep.id, ep.trust, namespace=self.namespace)

            details = {
                "used": list(used_set),
                "pruned": pruned_ids,
                "namespace": self.namespace
            }
            if pruned_ids:
                 self._emit(
                    EventType.FORGET,
                    f"Trust updated and {len(pruned_ids)} pruned.",
                    details
                )
            self.stability_monitor.observe_trust(
                [ep.trust for ep in self.episodes])
            self._bump_revision()

    # ------------------------------------------------------------------ #
    def remove_episode_ids(
        self, episode_ids: Iterable[str], *, namespace: Optional[str]=None
    ) -> list[str]:
        """Remove specific episodes across persistence + index layers."""

        ids = [str(eid) for eid in episode_ids]
        if not ids:
            return []

        removed: list[str] = []
        with self.use_namespace(namespace):
            if self.policy_evaluator is not None:
                self.policy_evaluator.evaluate_forget(
                    episode_ids=ids,
                    namespace=self.namespace,
                    reason="remove_by_id",
                )
            for episode_id in ids:
                if self.store is not None:
                    try:
                        self.store.remove(episode_id)
                    except Exception as e:
                        self._emit(EventType.SYSTEM, f"Failed to remove episode {episode_id} from store", {"error": str(e)}, level="warn")
                self._remove_graph_episode(episode_id, namespace=self.namespace)
                for episode in list(self.episodes):
                    if episode.id == episode_id:
                        self.episodes.remove(episode)
                        removed.append(episode_id)
                        break
                self._episode_by_id.pop(episode_id, None)
                self._index_remove_episode(episode_id)
                self.entity_index.remove(episode_id)

            if removed:
                self._emit(
                    EventType.FORGET,
                    f"Removed {len(removed)} episodes by id.",
                    {"removed": removed, "namespace": self.namespace},
                )
                self._bump_revision()

        return removed

    # ------------------------------------------------------------------ #
    def summarise(self) -> None:
        """
        Merge episodes to maintain memory capacity (Nmax).

        Uses hierarchical_merge to combine the two oldest episodes or
        episodes selected by a strategy.
        """
        if len(self.episodes) <= self.Nmax:
            return
        if self.policy_evaluator is not None:
            self.policy_evaluator.evaluate_summarise(self.episodes, namespace=self.namespace)

        # Simple strategy: Merge the two oldest episodes
        # Sort by creation time (tocc)
        sorted_episodes = sorted(self.episodes, key=lambda e: e.tocc)

        # Take top 'b' episodes (branching factor) or just 2
        b = int(self.b) if hasattr(self, 'b') else 2
        to_merge = sorted_episodes[:b]

        if len(to_merge) < 2:
            return

        # Merge graphs
        graphs = [ep.G for ep in to_merge]
        weights = [ep.trust for ep in to_merge]
        merged_graph = hierarchical_merge(graphs, weights=weights)

        # Merge vectors (weighted average)
        vectors = np.array([ep.e for ep in to_merge])
        total_weight = sum(weights)
        if total_weight > 0:
            merged_vector = np.average(vectors, axis=0, weights=weights)
        else:
            merged_vector = np.mean(vectors, axis=0)

        # Determine parent (inherit from first child if present)
        # This aligns with test expectations where children share a parent
        parent_id = to_merge[0].parent

        # Create new episode
        summary_ep = Episode(
            e=merged_vector,
            G=merged_graph,
            tocc=datetime.now(),
            meta={"type": "summary", "summary_of": [ep.id for ep in to_merge]},
            parent=parent_id,
            trust=1.0
        )
        
        # Update children to point to summary as parent
        for ep in to_merge:
            ep.parent = summary_ep.id

        # Batch update children in store
        if self.store is not None:
            try:
                # Update in store - use consistent API (save_many or save_episode)
                if hasattr(self.store, "save_many"):
                    self.store.save_many(to_merge)
                elif hasattr(self.store, "save_episode"):
                    for ep in to_merge:
                        self.store.save_episode(ep)
            except Exception as e:
                self._emit(EventType.SYSTEM, "Failed to update parents for merged episodes", {"error": str(e)}, level="warn")

        # Remove old episodes from active memory and index
        # But NOT from store (so they persist as history)
        ids_to_remove = [ep.id for ep in to_merge]
        with self.use_namespace(self.namespace):
            for episode_id in ids_to_remove:
                # Remove from active list
                for episode in list(self.episodes):
                    if episode.id == episode_id:
                        self.episodes.remove(episode)
                        break
                self._episode_by_id.pop(episode_id, None)
                # Remove from index
                self._index_remove_episode(episode_id)

        # Insert new episode
        self.episodes.append(summary_ep)
        self._episode_by_id[summary_ep.id] = summary_ep
        self._persist_batch([summary_ep])

        # Explicit graph store insertion for summary episode (fix: summarise graph consistency)
        # Skip if store IS the graph_store (same object) to avoid duplicate edges
        if summary_ep.G is not None:
            self._persist_graph_episode(summary_ep, namespace=self.namespace)

        if self._index:
            self._index.add(summary_ep.id, summary_ep.e)

        self._emit(
            EventType.SYSTEM,
            f"Summarised {len(to_merge)} episodes into {summary_ep.id}",
            {"source_ids": ids_to_remove, "summary_id": summary_ep.id}
        )
        self._bump_revision()

    # ------------------------------------------------------------------ #
    # Conversational Edge Query Methods
    # ------------------------------------------------------------------ #
    def get_following_turns(
        self,
        episode_id: str,
        k: int = 5,
        *,
        namespace: Optional[str] = None
    ) -> list[Episode]:
        """
        Get episodes that temporally follow the given one.
        
        Uses TURN_FOLLOWS edges for traversal.
        
        Args:
            episode_id: Starting episode ID
            k: Maximum number of following episodes to return
            namespace: Optional namespace filter
            
        Returns:
            List of Episode objects in temporal order (closest first)
        """
        gs = self.graph_store
        if gs is None:
            return []
        
        effective_ns = namespace or self.namespace
        following_ids = gs.query_turns_after(episode_id, limit=k, namespace=effective_ns)
        
        return self._resolve_episode_ids(following_ids, namespace=effective_ns)

    def get_by_speaker(
        self,
        speaker: str,
        k: int = 20,
        *,
        namespace: Optional[str] = None
    ) -> list[Episode]:
        """
        Get episodes from a specific speaker.
        
        Uses SPOKEN_BY edges for lookup.
        
        Args:
            speaker: Speaker name to search for
            k: Maximum number of episodes to return
            namespace: Optional namespace filter
            
        Returns:
            List of Episode objects (ordered by weight, highest first)
        """
        gs = self.graph_store
        if gs is None:
            return []
        
        effective_ns = namespace or self.namespace
        episode_ids = gs.query_by_speaker(speaker, limit=k, namespace=effective_ns)
        
        return self._resolve_episode_ids(episode_ids, namespace=effective_ns)

    def get_in_session(
        self,
        session_id: str,
        *,
        namespace: Optional[str] = None
    ) -> list[Episode]:
        """
        Get all episodes in a specific session.
        
        Uses IN_SESSION edges for lookup.
        
        Args:
            session_id: Session identifier to search for
            namespace: Optional namespace filter
            
        Returns:
            List of Episode objects (ordered by turn_index)
        """
        gs = self.graph_store
        if gs is None:
            return []
        
        effective_ns = namespace or self.namespace
        episode_ids = gs.query_by_session(session_id, namespace=effective_ns)
        
        return self._resolve_episode_ids(episode_ids, namespace=effective_ns)

    def _resolve_episode_ids(
        self,
        episode_ids: list[str],
        *,
        namespace: Optional[str] = None
    ) -> list[Episode]:
        """
        Convert episode IDs to Episode objects with O(1) lookup.
        
        Priority order:
        1. Local index (_episode_by_id) - O(1)
        2. GraphStore.load_by_id() - O(1)  
        3. Log warning if not found
        
        Args:
            episode_ids: List of episode IDs to resolve
            namespace: Optional namespace filter
            
        Returns:
            List of Episode objects (missing episodes are logged and skipped)
        """
        gs = self.graph_store
        result = []
        
        for ep_id in episode_ids:
            # Try O(1) local index first
            ep = self._episode_by_id.get(ep_id)
            
            if ep is None and gs is not None and hasattr(gs, 'load_by_id'):
                # Try loading from graph store
                load_by_id = getattr(gs, "load_by_id", None)
                if callable(load_by_id):
                    try:
                        ep = load_by_id(ep_id, namespace=namespace)
                    except TypeError:
                        ep = load_by_id(ep_id)
            
            if ep is not None:
                result.append(ep)
            elif self.debug:
                print(f"[Memory] Warning: Episode {ep_id} not found in index or store")
        
        return result

    def reset_conversation(self, namespace: Optional[str] = None) -> None:
        """
        Reset conversation state for TURN_FOLLOWS edge tracking.
        
        Call this when starting a new conversation to prevent
        TURN_FOLLOWS edges from connecting unrelated episodes.
        
        Thread-safe: uses _conversation_lock for state updates.
        """
        with self._conversation_lock:
            if namespace:
                state = self._conversation_state_by_namespace[namespace]
                state["previous_episode_id"] = None
                state["current_session_id"] = None
                state["turn_counter"] = 0
            else:
                self._conversation_state_by_namespace.clear()
                self._previous_episode_id = None
                self._current_session_id = None
                self._turn_counter = 0
        
        if self.debug:
            print("[Memory] Conversation state reset")

    def symbolic_inconsistency_metrics(self, namespace: str | None = None) -> dict[str, Any]:
        """
        Detect symbolic contradictions in the current memory state.
        """
        ns = namespace or self.namespace

        if self._multi_tenant:
             episodes = self._episodes_by_namespace[ns]
        else:
             episodes = self._episodes_single

        graph_entries = []
        for ep in episodes:
            if ep.G:
                graph_entries.append((ep.G, {"episode_id": ep.id, "trust": ep.trust}))

        inconsistencies = self.symbolic_reasoner.inconsistencies(graph_entries)
        result = inconsistencies.to_dict()
        result["namespace"] = ns
        return result

    def build_query_graph(self, query_text: str) -> Optional[nx.DiGraph]:
        """
        Build a query graph from query text for retrieval.

        Uses the configured graph extractor to extract entities from
        the query, enabling graph-enhanced retrieval without oracle knowledge.

        When LLM extraction is enabled, also expands the query to find
        implicit entities (e.g., "Who founded Facebook?" -> ["Facebook", "founder", "Mark Zuckerberg"])

        Args:
            query_text: The query/question text

        Returns:
            A DiGraph containing extracted entities, or None if extraction disabled
        """
        # 1. Use new ExtractionPipeline (Preferred)
        if self._extraction_pipeline is not None:
             # Just use standard extract - LLM extractors usually have query expansion logic built-in
             # or we rely on extracting entities from the query itself.
             # For now, let's treat the query as text to extract from.
             result = self._extraction_pipeline.extract(query_text, link_entities=False)
             if result.graph and result.graph.number_of_nodes() > 0:
                 return result.graph
             elif result.entities:
                 # Fallback: make a graph from entities if no edges
                 G = nx.DiGraph()
                 for ent in result.entities:
                     G.add_node(ent.canonical)
                 return G

        # 2. Legacy LLM Extractor
        if self._llm_extractor is not None:
            # Note: expand_query returns a dict, not a graph object directly unless converted
            # This legacy path might be broken, let's just use extract()
            result = self._llm_extractor.extract(query_text)
            if result.graph:
                return result.graph
        
        # 3. Legacy Basic/Regex Extractor
        if self._graph_extractor is None:
            return None

        entities = self._graph_extractor.extract_entities(query_text)
        if not entities:
            return None

        G = nx.DiGraph()
        for ent in entities:
             G.add_node(ent)

        # Add edges between entities
        from itertools import combinations
        for e1, e2 in combinations(entities, 2):
            G.add_edge(e1, e2, relation="co_mentioned")

        return G

    def symbolic_inconsistency_by_namespace(self) -> dict[str, dict[str, Any]]:
        """
        Return symbolic inconsistencies grouped by namespace.
        """
        if not self._multi_tenant:
            return {self.namespace: self.symbolic_inconsistency_metrics(self.namespace)}

        results = {}
        for ns in self._episodes_by_namespace.keys():
            results[ns] = self.symbolic_inconsistency_metrics(ns)
        return results

    def prune_below_trust(
        self, threshold: float, *, namespace: Optional[str]=None
    ) -> list[str]:
        """Prune all episodes with trust strictly below the provided threshold."""

        removed: list[str] = []
        with self.use_namespace(namespace):
            for episode in list(self.episodes):
                if episode.trust < threshold:
                    self.episodes.remove(episode)
                    removed.append(episode.id)
                    if self.store is not None:
                        try:
                            self.store.remove(episode.id)
                        except Exception as e:
                            self._emit(EventType.SYSTEM, f"Failed to remove episode {episode.id} from store", {"error": str(e)}, level="warn")
                    self._remove_graph_episode(episode.id, namespace=self.namespace)
                    self._episode_by_id.pop(episode.id, None)
                    self._index_remove_episode(episode.id)
                    self.entity_index.remove(episode.id)

            if removed:
                details = {
                    "threshold": float(threshold),
                    "removed": removed,
                    "namespace": self.namespace,
                }
                self._emit(
                    EventType.FORGET,
                    f"Pruned {len(removed)} episodes below trust threshold.",
                    details,
                )
                self._bump_revision()

        return removed

    # ------------------------------------------------------------------ #
    def clear(self) -> None:
        """Wipe memory clean."""
        removed_ids: list[str] = []
        if self._multi_tenant:
            # Capture IDs across all namespaces before clearing
            for ns, episodes in list(self._episodes_by_namespace.items()):
                removed_ids.extend(ep.id for ep in episodes)
            
            self._episodes_by_namespace = defaultdict(list)
            self._context_by_namespace = defaultdict(lambda: None)
            # Re-establish a default namespace so callers can immediately insert
            self.set_namespace("default")
        else:
            removed_ids = [ep.id for ep in self.episodes]
            self.episodes.clear()
        self._episode_by_id.clear()

        if self.store is not None:
            try:
                self.store.clear()
            except Exception as e:
                self._emit(EventType.SYSTEM, "Failed to clear store", {"error": str(e)}, level="error")

        gs = self.graph_store
        if gs is not None and gs is not self.store:
            try:
                gs.clear()
            except Exception as e:
                self._emit(EventType.SYSTEM, "Failed to clear graph store", {"error": str(e)}, level="warn")

        index = self.index
        if index is not None:
            if hasattr(index, "clear"):
                index.clear()
            else:
                for episode_id in removed_ids:
                    self._index_remove_episode(episode_id)
        
        self.entity_index.clear()

        self._emit(EventType.CLEAR, "All episodes cleared.")
        self._bump_revision()

    # ------------------------------------------------------------------ #
    # Stats & Monitoring
    # ------------------------------------------------------------------ #
    def stats(self) -> dict[str, Any]:
        """Return snapshot of memory and governance status."""
        trusts = np.array([ep.trust for ep in self.episodes]
                          ) if self.episodes else np.array([0])
        stability = self.stability_monitor.stability()
        
        # Base stats
        result = {
            "count": len(self.episodes),
            "avg_trust": float(np.mean(trusts)),
            "min_trust": float(np.min(trusts)),
            "max_trust": float(np.max(trusts)),
            "lambda": self.lam,
            "tau": self.tau,
            "theta": self.theta,
            "governance_safe": self.gsr.is_valid(),
            "stability_score": self.gsr.stability_score(),
            "risk_level": self.gsr.risk_level(),
            "retrieval_entropy": stability.retrieval_entropy,
            "summary_fidelity": stability.summary_fidelity,
            "summary_distortion": stability.summary_distortion,
            "summary_distortion_bound": stability.summary_distortion_bound,
            "trust_floor": stability.trust_floor,
            "memory_utilisation": stability.memory_utilisation,
            "contraction_factor": stability.contraction_factor,
            "lyapunov_margin": stability.lyapunov_margin,
            "lyapunov_warning": stability.lyapunov_warning,
            "global_stability": stability.stability_score,
            "namespace": self.namespace,
            "multi_tenant": self._multi_tenant,
        }
        
        # Add graph stats if available (Phase 0.3: consistent schema)
        if self.graph_store is not None:
            graph_stats = self.graph_store.get_graph_stats(namespace=self.namespace)
            # Ensure enabled flag is always present
            graph_stats["enabled"] = True
            result["graph_store"] = graph_stats
        else:
            result["graph_store"] = {
                "enabled": False,
                "node_count": 0,
                "edge_count": 0,
                "backend": None
            }
        
        return result

    # ------------------------------------------------------------------ #
    def graph_diagnostics(self) -> dict[str, Any]:
        """
        Return detailed graph diagnostics for observability (Phase 4.1).

        Returns:
            Dict containing:
            - enabled: bool
            - backend: str or None
            - node_count: int
            - edge_count: int
            - top_relations: list of (relation_type, count) tuples
            - namespaces: list of namespace strings present in graph
        """
        if self.graph_store is None:
            return {
                "enabled": False,
                "backend": None,
                "node_count": 0,
                "edge_count": 0,
                "top_relations": [],
                "namespaces": []
            }

        stats = self.graph_store.get_graph_stats(namespace=self.namespace)

        # Collect relation type distribution if possible
        top_relations: list[tuple[str, int]] = []
        namespaces: set[str] = set()

        if hasattr(self.graph_store, "_G"):
            relation_counts: dict[str, int] = {}
            for _, _, data in self.graph_store._G.edges(data=True):
                rel = data.get("relation_type", "UNKNOWN")
                relation_counts[rel] = relation_counts.get(rel, 0) + 1
                ns = data.get("namespace")
                if ns:
                    namespaces.add(ns)
            # Sort by count descending, take top 10
            top_relations = sorted(relation_counts.items(), key=lambda x: -x[1])[:10]

        return {
            "enabled": True,
            "backend": stats.get("backend", "unknown"),
            "type": stats.get("type", type(self.graph_store).__name__),
            "node_count": stats.get("node_count", 0),
            "edge_count": stats.get("edge_count", 0),
            "top_relations": top_relations,
            "namespaces": sorted(namespaces)
        }

    # ------------------------------------------------------------------ #
    # Evaluation Methods
    # ------------------------------------------------------------------ #

    def evaluate(
        self,
        examples: List[Any],
        k: int = 10,
        use_graph: bool = True,
        alpha_dense: float = 1.0,
        beta_symbolic: float = 50.0,
    ) -> Any:
        """
        Evaluate retrieval performance on a test set.

        Args:
            examples: List of EvaluationExample objects
            k: Number of results to retrieve
            use_graph: Whether to use graph-enhanced retrieval
            alpha_dense: Dense weight for hybrid retrieval
            beta_symbolic: Symbolic weight for hybrid retrieval

        Returns:
            EvaluationResult with metrics (recall, NDCG, etc.)

        Example:
            >>> from hieromem_core.evaluation import create_example
            >>> examples = [create_example("query", ["doc1", "doc2"])]
            >>> result = memory.evaluate(examples, k=5)
            >>> print(f"Recall@5: {result.recall_at_k:.3f}")
        """
        from hieromem_core.evaluation import Evaluator
        evaluator = Evaluator(self)  # type: ignore[no-untyped-call]
        return evaluator.evaluate(
            examples, k=k, use_graph=use_graph,
            alpha_dense=alpha_dense, beta_symbolic=beta_symbolic
        )

    def compare(
        self,
        examples: List[Any],
        k: int = 10,
        methods: Optional[List[str]] = None,
    ) -> Any:
        """
        Compare multiple retrieval methods on a test set.

        Args:
            examples: List of EvaluationExample objects
            k: Number of results to retrieve
            methods: Methods to compare ('graph', 'dense')

        Returns:
            ComparisonResult with per-method results and winner

        Example:
            >>> result = memory.compare(examples, methods=['graph', 'dense'])
            >>> print(result.summary())
        """
        from hieromem_core.evaluation import Evaluator
        evaluator = Evaluator(self)  # type: ignore[no-untyped-call]
        return evaluator.compare(examples, k=k, methods=methods)

    # ------------------------------------------------------------------ #
    # Learning
    # ------------------------------------------------------------------ #

    def train_retriever(self, examples: Iterable[Any]) -> None:
        """
        Train the Learned Retriever using evaluation examples.

        Args:
            examples: List of EvaluationExample objects with ground truth.
        """
        if hasattr(self.dual_retriever, "train"):
            # Ensure list
            ex_list = list(examples)
            self.dual_retriever.train(ex_list)
        else:
            print("[Memory] Retriever is not trainable.")

    # ------------------------------------------------------------------ #
    # Global Knowledge Graph
    # ------------------------------------------------------------------ #

    def get_unified_graph(self, namespace: Optional[str] = None) -> nx.DiGraph:
        """
        Get the Unified Knowledge Graph.

        This returns a consolidated view of the memory, merging duplicate edges
        and resolving conflicts using consensus logic.

        Args:
            namespace: Optional namespace to filter (None = global).

        Returns:
            nx.DiGraph: High-quality directed graph.
        """
        from hieromem_core.graph import UnifiedGraphBuilder
        if self.graph_store:
            builder = UnifiedGraphBuilder(self.graph_store)  # type: ignore[no-untyped-call]
            return builder.build(namespace=namespace)
        return nx.DiGraph()

    # ------------------------------------------------------------------ #
    def multi_modal_insert(
        self,
        payload: dict[str, Any],
        *,
        graph: Optional[nx.DiGraph]=None,
        meta: Optional[dict[str, Any]]=None,
        namespace: Optional[str]=None,
    ) -> Episode:
        payload = dict(payload)
        embedding_override = payload.pop("embedding", None)
        output = None
        vector: Optional[np.ndarray] = None
        if embedding_override is not None:
            vector = np.asarray(embedding_override, dtype=np.float32)
        combined_meta = dict(meta or {})
        if vector is None and "text" in payload and hasattr(self, "encoder"):
            text_values = list(self._ensure_iterable(payload["text"]))
            if text_values:
                encodings = np.stack(
                    [self.encoder.encode_single(text) for text in text_values])
                vector = np.mean(encodings, axis=0)
                norm = np.linalg.norm(vector)
                if norm > 0:
                    vector = vector / norm
                vector = vector.astype(np.float32)
                combined_meta.setdefault("modalities", {"text": 1.0})

        if vector is None:
            output = self.multi_modal_encoder.encode(payload)
            vector = output.vector
        if output is not None:
            combined_meta.setdefault(
    "modalities", dict(
        output.modality_weights))
        else:
            combined_meta.setdefault("modalities", {"embedding": 1.0})
        combined_meta.setdefault("payload_keys", sorted(payload.keys()))
        combined_meta.setdefault("encoder_id", self._default_encoder_id())
        auto_graph = graph or self._build_scene_graph_from_payload(payload)
        return self.insert(
    vector,
    auto_graph,
    meta=combined_meta,
     namespace=namespace)
     
    # ------------------------------------------------------------------ #
    # Explicit Graph Retrieval (Zero-Hop and Multi-Hop)
    # ------------------------------------------------------------------ #
    def graph_search(
        self,
        episode_id: str,
        depth: int = 2,
        limit: int = 50,
        *,
        namespace: Optional[str] = None
    ) -> list[dict[str, Any]]:
        """
        Perform a graph expansion search from a given episode.
        
        Delegates to the configured GraphStore.
        
        Args:
            episode_id: The ID of the episode to use as a seed.
            depth: Traversal depth (hops).
            limit: Maximum number of related entities/paths to return.
            namespace: Optional namespace.
            
        Returns:
            List of related entities/paths found in the graph.
        
        Raises:
            RuntimeError: If no GraphStore is configured.
        """
        if self.graph_store is None:
            raise RuntimeError("Graph search requested but no GraphStore is configured.")
            
        with self.use_namespace(namespace):
            return self.graph_store.search_related(
                episode_id,
                depth=depth,
                limit=limit,
                namespace=namespace or self.namespace
            )

    # ------------------------------------------------------------------ #
    def symbolic_search(
        self,
        *,
        path: Optional[Sequence[str]]=None,
        pattern: Optional[nx.DiGraph]=None,
        namespace: Optional[str]=None,
        limit: int=5,
    ) -> list[tuple[Episode, dict[str, Any]]]:
        """Return episodes whose scene graphs satisfy symbolic predicates."""

        if path is None and pattern is None:
            raise ValueError("symbolic_search requires a path or pattern")

        results: list[tuple[Episode, dict[str, Any]]] = []
        with self.use_namespace(namespace):
            for ep in self.episodes:
                matches: dict[str, Any] = {}
                graph = getattr(ep, "G", None)
                if path is not None and self.symbolic_reasoner.has_path(
                    graph, path):
                    matches["path"] = list(path)
                    matches["paths_found"] = self.symbolic_reasoner.find_paths(
                        graph, path[0], path[-1], limit=limit)
                if pattern is not None:
                    subgraphs = self.symbolic_reasoner.find_subgraphs(
                        graph, pattern, limit=limit)
                    if subgraphs:
                        matches.setdefault("subgraphs", subgraphs)
                if matches:
                    results.append((ep, matches))
        return results

    # ------------------------------------------------------------------ #
    def batch_symbolic_search(
        self,
        queries: Sequence[BatchSymbolicSearchInput],
    ) -> list[list[tuple[Episode, dict[str, Any]]]]:
        """Execute multiple symbolic graph searches sequentially."""

        batch_results: list[list[tuple[Episode, dict[str, Any]]]] = []
        for entry in queries:
            kwargs: dict[str, Any] = {}
            if entry.path is not None:
                kwargs["path"] = entry.path
            if entry.pattern is not None:
                kwargs["pattern"] = entry.pattern
            if entry.limit is not None:
                kwargs["limit"] = entry.limit
            if "path" not in kwargs and "pattern" not in kwargs:
                raise ValueError(
                    "Each batch symbolic search entry requires a path or pattern")
            batch_results.append(
                self.symbolic_search(
                    namespace=entry.namespace,
                    **kwargs,
                )
            )
        return batch_results

    # ------------------------------------------------------------------ #
    def _build_scene_graph_from_payload(
        self, payload: Mapping[str, Any]) -> nx.DiGraph:
        graph = nx.DiGraph()
        root_id = "experience"
        graph.add_node(root_id, type="experience", namespace=self.namespace)

        for idx, text in enumerate(self._ensure_iterable(payload.get("text"))):
            node_id = f"text:{idx}"
            graph.add_node(node_id, type="text", value=str(text))
            graph.add_edge(root_id, node_id, relation="describes")

        for idx, image in enumerate(
    self._ensure_iterable(
        payload.get("image"))):
            node_id = f"image:{idx}"
            graph.add_node(node_id, type="image", ref=str(image))
            graph.add_edge(root_id, node_id, relation="illustrates")

        for idx, audio in enumerate(
    self._ensure_iterable(
        payload.get("audio"))):
            node_id = f"audio:{idx}"
            graph.add_node(node_id, type="audio", ref=str(audio))
            graph.add_edge(root_id, node_id, relation="narrates")

        for key, value in payload.items():
            if key in {"text", "image", "audio"}:
                continue
            node_id = f"meta:{key}"
            graph.add_node(
    node_id,
    type="metadata",
    value=_stringify_meta(value))
            graph.add_edge(root_id, node_id, relation="has_meta")

        return graph

    @ staticmethod
    def _ensure_iterable(value: Any) -> Iterable[Any]:
        if value is None:
            return []
        if isinstance(value, (list, tuple, set)):
            return value
        return [value]

    # ------------------------------------------------------------------ #
    # Unified Retrieval API (Canonical Method)
    # ------------------------------------------------------------------ #
    def retrieve(
        self,
        query: str,
        k: int = 5,
        *,
        mode: str = "graph_hybrid",
        alpha_dense: float = 1.0,
        beta_symbolic: float = 50.0,
        min_confidence: float = 0.0,
        namespace: Optional[str] = None,
        realm: Optional[str] = None,
        agent: Optional[str] = None,
    ) -> list[Episode]:
        """
        Unified retrieval method - the recommended way to query HIEROMEM.

        This method automatically:
        1. Encodes the query text to a vector
        2. Extracts entities and builds a query graph
        3. Uses graph-enhanced dual retrieval

        Args:
            query: The query text (question or search query)
            k: Number of episodes to retrieve
            mode: Retrieval mode:
                - "graph_hybrid": Use both dense and graph (default)
                - "dense_only": Pure vector similarity
                - "graph_only": Pure symbolic matching (requires query graph)
            alpha_dense: Weight for dense scores (default 1.0)
            beta_symbolic: Weight for symbolic/graph scores (default 50.0)
            beta_symbolic: Weight for symbolic/graph scores (default 50.0)
            min_confidence: Minimum confidence score (0.0 to 1.0) to return (default 0.0)
            namespace: Optional namespace for multi-tenant scenarios
            realm: Optional realm (overrides namespace)
            agent: Agent ID for permission check

        Returns:
            List of top-k most relevant episodes

        Example:
            >>> episodes = memory.retrieve("What did Alice tell Bob?", k=5)
            >>> for ep in episodes:
            ...     print(ep.meta.get('text', '')[:100])
        """
        # Phase 4: Realm enforcement
        if realm:
            if realm not in self.realms:
                raise ValueError(f"Realm '{realm}' does not exist")
            if agent:
                enforce_permission(self.realms[realm], agent, "read")
            namespace = realm

        # Step 1: Encode query to vector
        if self.encoder is None:
            raise ValueError("Memory.retrieve requires an encoder. Pass encoder= to __init__.")
        
        query_vec = self.encoder.encode(query)
        # Handle 2D output from some encoders - squeeze to remove extra dimensions
        if hasattr(query_vec, 'squeeze'):
            query_vec = query_vec.squeeze()
        query_vec = np.asarray(query_vec, dtype=np.float32).flatten()

        # Step 2: Build query graph (if using graph-enhanced modes)
        query_graph = None
        if mode in ("graph_hybrid", "graph_only"):
            query_graph = self.build_query_graph(query)

        # Step 3: Dispatch to appropriate retrieval method
        if mode == "dense_only":
            # Pure vector similarity
            return self.resonant_retrieve(query_vec, k=k, namespace=namespace)
        
        elif mode in ("graph_hybrid", "graph_only"):
            # Use dual retrieval with graph
            result = self.dual_retrieve(
                query_vec,
                query_graph=query_graph,
                k=k,
                alpha_dense=alpha_dense if mode == "graph_hybrid" else 0.0,
                beta_symbolic=beta_symbolic,
                namespace=namespace,
            )
            episodes = result.episodes
        
        else:
            raise ValueError(f"Unknown mode: {mode}. Use 'graph_hybrid', 'dense_only', or 'graph_only'.")

        # Step 4: Filter by confidence
        if min_confidence > 0:
            filtered = []
            for ep in episodes:
                # meta['score'] is populated by dual_retrieve / resonant_retrieve
                # If score is missing, we assume 0.0 (or we could assume 1.0 if we trust retrieval implicitly?)
                # Safety decision: score must be present and >= min_confidence
                score = ep.meta.get("score", 0.0)
                if score >= min_confidence:
                    filtered.append(ep)
            return filtered
            
        return episodes

    def explain_retrieval(self, query: str, result: Episode) -> str:
        """
        Generate an explanation for why this result was retrieved.
        
        Handles both atomic episodes and derived reasoning chains.
        """
        try:
            # Case 1: Reasoning Chain (Detailed symbolic explanation)
            if result.meta.get("type") == "reasoning_chain":
                # Reconstruct chain from meta if needed, or if it's already a chain object
                # For now, we assume we might need to rely on the meta
                
                # Check for low trust warning
                score = result.meta.get("score", 0.0)
                warning = ""
                if score < 0.4:
                    warning = "\n[WARNING: Low-trust reasoning chain. Verify sources.]\n"

                # If we have the components, we can reconstruct a reasonable explanation text
                path_str = result.meta.get("path", [])
                explanation = (
                    f"Reasoning Chain Result\n"
                    f"======================\n"
                    f"Query: {query}\n"
                    f"{warning}\n"
                    f"Path: {' -> '.join(path_str) if isinstance(path_str, list) else path_str}\n"
                    f"Score: {score:.2f}\n"
                )
                if "components" in result.meta:
                    comps = result.meta["components"]
                    explanation += (
                        f"  - Trust: {comps.get('trust', 0):.2f}\n"
                        f"  - Semantic: {comps.get('semantic', 0):.2f}\n"
                    )
                return explanation

            # Case 2: Atomic Episode (Standard retrieval metrics)
            else:
                score = result.meta.get("score", "N/A")
                return (
                    f"Standard Retrieval Result\n"
                    f"-------------------------\n"
                    f"Score: {score}\n"
                    f"Source: Atomic Episode (Dense/Symbolic match)\n"
                    f"Content: {result.meta.get('content_preview', 'No preview')[:100]}..."
                )
        except Exception as e:
            return f"Could not explain retrieval: {e}"

    # ------------------------------------------------------------------ #
    # Phase 3: Reasoning Facade Methods
    # ------------------------------------------------------------------ #
    
    def recall_before(
        self,
        event: str | Episode,
        limit: int = 10,
        namespace: Optional[str] = None
    ) -> list[Episode]:
        """
        Recall episodes that occurred strictly before the given event.
        
        Delegates to reasoning.temporal.episodes_before.
        
        Args:
            event: Episode ID or Episode object acting as the reference point.
            limit: Max results to return.
            namespace: Namespace filter.
            
        Returns:
            List of episodes sorted by time (newest first).
        """
        from hieromem_core.reasoning.temporal import episodes_before
        
        ref_time: datetime | None = None
        
        if isinstance(event, str):
            ep = self._get_episode_by_id(event)
            if ep:
                ref_time = ep.tocc
        else:
            ref_time = event.tocc
            
        if ref_time is None:
            # If we can't find the event or its time, we can't search before it
            if self.debug:
                print(f"[Memory] recall_before: Reference event has no timestamp or not found.")
            return []
            
        with self.use_namespace(namespace):
            # We filter self.episodes which is the active working set
            # For persistent store queries, one would use the TemporalIndex (Phase 3.2 item)
            # but here we use the loaded episodes or would need to query the store if index supported it.
            # Assuming self.episodes contains the relevant working set for now as per v1 architecture.
            return episodes_before(self.episodes, ref_time, limit=limit)

    def recall_temporal(
        self,
        query: Optional[str] = None,
        start: Optional[datetime | str] = None,
        end: Optional[datetime | str] = None,
        namespace: Optional[str] = None,
        limit: int = 50,
    ) -> list[Episode]:
        """
        Recall episodes within a specific time range, optionally filtering by semantic query.
        
        Delegates to reasoning.temporal.episodes_in_range.
        
        Args:
            query: Optional text query to filter by relevance first.
            start: Start datetime (inclusive) or ISO string.
            end: End datetime (inclusive) or ISO string.
            namespace: Namespace filter.
            limit: Max results.
            
        Returns:
            List of episodes sorted chronologically (oldest first).
        """
        from hieromem_core.reasoning.temporal import episodes_in_range
        
        # Parse strings to datetime if needed
        start_dt = datetime.fromisoformat(start) if isinstance(start, str) else start
        end_dt = datetime.fromisoformat(end) if isinstance(end, str) else end
        
        with self.use_namespace(namespace):
            # If query provided, retrieve first then filter
            candidates = self.episodes
            if query:
                # We retrieve more than limit to allow for temporal filtering
                candidates = self.retrieve(query, k=limit * 2, namespace=self.namespace)
                
            filtered = episodes_in_range(
                candidates,
                start=start_dt,
                end=end_dt,
                limit=limit,
                ascending=True # Chronological
            )
            return filtered

    def reconstruct_sequence(
        self,
        start_event: Optional[str] = None,
        end_event: Optional[str] = None,
        limit: int = 100,
        namespace: Optional[str] = None
    ) -> list[Episode]:
        """
        Reconstruct a sequence of episodes (causal or temporal).
        
        Delegates to reasoning.sequence.reconstruct_sequence.
        
        Args:
            start_event: Start episode ID.
            end_event: End episode ID.
            limit: Max sequence length.
            namespace: Namespace filter.
            
        Returns:
            Ordered list of Episode objects forming the sequence.
        """
        from hieromem_core.reasoning.sequence import reconstruct_sequence as reconstruct_helper
        
        # reconstruct_sequence helper in sequence.py takes (memory, ...)
        # We pass 'self' as the memory interface
        return reconstruct_helper(
            self,
            start_episode_id=start_event,
            end_episode_id=end_event,
            limit=limit,
            namespace=namespace
        )

    def find_contradictions(
        self,
        query: Optional[str] = None,
        namespace: Optional[str] = None
    ) -> list[Any]:
        """
        Find contradictions in memory, optionally focused on a specific topic.
        
        Delegates to reasoning.contradiction.ContradictionDetector.
        
        Args:
            query: Optional query to focus the search (e.g., "Alice's job").
            namespace: Namespace filter.
            
        Returns:
            List of Contradiction objects.
        """
        from hieromem_core.reasoning.contradiction import ContradictionDetector
        
        with self.use_namespace(namespace):
            detector = ContradictionDetector()
            
            candidates = self.episodes
            if query:
                # If query is focused, we only look at relevant episodes
                # but we need enough context to find conflicts.
                # Heuristic: retrieve top 20 relevant items
                candidates = self.retrieve(query, k=20, namespace=self.namespace)
                
            return detector.find_contradictions(candidates)

    # ------------------------------------------------------------------
    def dual_retrieve(
        self,
        query_vector: np.ndarray,
        *,
        query_graph: Optional[nx.DiGraph]=None,
        k: int=3,
        alpha_dense: Optional[float]=None,
        beta_symbolic: Optional[float]=None,
        temperature: Optional[float]=None,
        max_iterations: Optional[int]=None,
        tolerance: Optional[float]=None,
        namespace: Optional[str]=None,
    ) -> DualRetrievalResult:
        with start_span("memory.dual_retrieve", attributes={"namespace": namespace or self.namespace}):
            with self.use_namespace(namespace):
                base_config = DualRetrievalConfig()
                config = DualRetrievalConfig(
                    alpha=alpha_dense if alpha_dense is not None else base_config.alpha,
                    beta=beta_symbolic if beta_symbolic is not None else base_config.beta,
                    temperature=temperature if temperature is not None else base_config.temperature,
                    max_iter=max_iterations if max_iterations is not None else base_config.max_iter,
                    tol=tolerance if tolerance is not None else base_config.tol,
                    k=k if k is not None else base_config.k,
                )
                result = self.dual_retriever.retrieve(
                    query_vector,
                    query_graph=query_graph,
                    config=config,
                    namespace=self.namespace,
                )
                if result.weights.size:
                    self.stability_monitor.observe_retrieval(result.weights)
                self._emit(
                    EventType.RETRIEVE,
                    f"Dual retrieval returned {len(result.episodes)} episodes.",
                    {
                        "ids": [ep.id for ep in result.episodes],
                        "weights": result.weights.tolist(),
                        "dense": dict(result.dense_scores),
                        "symbolic": dict(result.symbolic_scores),
                        "iterations": result.iterations,
                        "namespace": self.namespace,
                    },
                )
                return result

    # ------------------------------------------------------------------ #
    def batch_dual_retrieve(
        self,
        queries: Sequence[DualRetrievalInput],
    ) -> list[DualRetrievalResult]:
        """Execute multiple dual retrieval queries sequentially."""

        results: list[DualRetrievalResult] = []
        for entry in queries:
            kwargs: dict[str, Any] = {}
            if entry.k is not None:
                kwargs["k"] = entry.k
            if entry.alpha is not None:
                kwargs["alpha_dense"] = entry.alpha
            if entry.beta is not None:
                kwargs["beta_symbolic"] = entry.beta
            if entry.temperature is not None:
                kwargs["temperature"] = entry.temperature
            if entry.max_iterations is not None:
                kwargs["max_iterations"] = entry.max_iterations
            if entry.tolerance is not None:
                kwargs["tolerance"] = entry.tolerance

            results.append(
                self.dual_retrieve(
                    entry.vector,
                    query_graph=entry.query_graph,
                    namespace=entry.namespace,
                    **kwargs,
                )
            )
        return results

    def __len__(self) -> int:
        return len(self.episodes)

    # ------------------------------------------------------------------ #
    # Persistence
    # ------------------------------------------------------------------ #
    def _load_from_store(self) -> None:
        """Load persisted episodes at startup."""
        store = self.store
        if store is None:
            return
        try:
            eps = store.load_all()
        except Exception as e:
            self._emit(EventType.SYSTEM, "Failed to load episodes from store", {"error": str(e)}, level="error")
            eps = []
        self._episode_by_id.clear()
        if self._multi_tenant:
            current = self.namespace
            for ep in eps:
                ns = str(
    ep.meta.get(
        "namespace",
        current)) if isinstance(
            ep.meta,
             dict) else current
                self.set_namespace(ns)
                self.episodes.append(ep)
                self._episode_by_id[ep.id] = ep
                self._persist_graph_episode(ep, namespace=ns)
            self.set_namespace(current)
        else:
            self.episodes.extend(eps)
            for ep in eps:
                self._episode_by_id[ep.id] = ep
                ns = self.namespace
                if isinstance(ep.meta, dict):
                    ns = str(ep.meta.get("namespace", ns))
                self._persist_graph_episode(ep, namespace=ns)
        self._emit(
    EventType.LOAD, "Loaded episodes from store.", {
        "count": len(eps)})
        self._reconcile_index()
        if eps:
            self._bump_revision()
            # Index loaded episodes
            for ep in self.episodes:
                 self.entity_index.add(ep)

    # ------------------------------------------------------------------ #
    # Event emission (to log + narrator)
    # ------------------------------------------------------------------ #
    def _emit(
        self,
        etype: EventType,
        message: str,
        details: dict[str, Any] | None=None,
        *,
        level: Literal["info", "warn", "error"]="info",
    ) -> None:
        payload = details or {}
        evt = MemoryEvent(
    event_type=etype,
    message=message,
    details=payload,
     level=level)
        self.event_log.add(
    etype,
    message,
    details=payload,
    source="Memory",
     level=level)
        # Narrator will automatically observe via subscription
        if self.debug:
            print(evt)

    # ------------------------------------------------------------------ #
    def get_recent_events(self, n: int=10) -> list[MemoryEvent]:
        return self.event_log.recent(n)

    # ------------------------------------------------------------------ #
    # Replay & Simulation (Phase 4)
    # ------------------------------------------------------------------ #
    def link_episodes(
        self,
        source: tuple[str, str],
        target: tuple[str, str],
        relation: str
    ) -> None:
        """
        Create a cross-reference between two episodes, potentially across realms.
        
        Args:
            source: (namespace, episode_id) of the source episode.
            target: (namespace, episode_id) of the target episode.
            relation: Description of the relationship (e.g., "derived_from").
        """
        src_ns, src_id = source
        tgt_ns, tgt_id = target
        
        # Verify source exists
        with self.use_namespace(src_ns):
            src_ep = self._get_episode_by_id(src_id)
            if not src_ep and self.store:
                src_ep = self.store.load_by_id(src_id)
                
        if not src_ep:
             raise ValueError(f"Source episode {src_id} in {src_ns} not found")

        # Verify target exists
        with self.use_namespace(tgt_ns):
            tgt_ep = self._get_episode_by_id(tgt_id)
            if not tgt_ep and self.store:
                tgt_ep = self.store.load_by_id(tgt_id)
        
        if not tgt_ep:
             raise ValueError(f"Target episode {tgt_id} in {tgt_ns} not found")
             
        # Add reference to source metadata
        # Initialize references list if not present
        if "references" not in src_ep.meta:
            src_ep.meta["references"] = []
            
        # Avoid duplicates
        ref_entry = {"namespace": tgt_ns, "episode_id": tgt_id, "relation": relation}
        if ref_entry not in src_ep.meta["references"]:
            src_ep.meta["references"].append(ref_entry)
            
            # Persist update
            with self.use_namespace(src_ns):
                if self.store:
                    self.store.save_episode(src_ep)
                # Update in-memory if loaded
                # We need to ensure we update the correct list if it's cached
                # But _get_episode_by_id returns from self.episodes which is usually the live list
                pass

    def export_training_data(
        self,
        namespace: Optional[str] = None,
        format: str = "jsonl",
        include_graph: bool = True,
        start: Optional[datetime] = None,
        end: Optional[datetime] = None,
    ) -> Any:
        """Export memory episodes for fine-tuning or analysis."""
        from hieromem_core.replay.export import export_training_data as _export
        return _export(self, namespace, format, include_graph, start, end)

    def replay(
        self, 
        start: Optional[datetime] = None, 
        end: Optional[datetime] = None,
        namespace: Optional[str] = None
    ) -> Iterator[Episode]:
        """
        Replay episodes in chronological order.
        
        Args:
            start: Start datetime (inclusive).
            end: End datetime (inclusive).
            namespace: Optional namespace to filter by.
        """
        from hieromem_core.replay import replay as _replay
        return _replay(self, start, end, namespace)

    def simulate(self) -> "AbstractContextManager[SimulationMemory]":
        """
        Create a simulation context (transactional sandbox).
        
        Usage:
            with memory.simulate() as sim:
                sim.insert(...) 
                # Changes are discarded on exit
        """
        from hieromem_core.replay import simulate as _simulate
        return _simulate(self)

    # ------------------------------------------------------------------ #
    # Private / Internal helpers for index + retrieval logic
    # ------------------------------------------------------------------ #
    def _get_episode_by_id(self, episode_id: str) -> Optional[Episode]:
        for ep in self.episodes:
            if ep.id == episode_id:
                return ep
        store = self.store
        if store is not None:
            try:
                candidate = store.load_by_id(episode_id)
                if candidate is not None:
                    self.episodes.append(candidate)
                    self._episode_by_id[candidate.id] = candidate
                    ns = self.namespace
                    if isinstance(candidate.meta, dict):
                        ns = str(candidate.meta.get("namespace", ns))
                    self._persist_graph_episode(candidate, namespace=ns)
                    return candidate
            except Exception as e:
                self._emit(EventType.SYSTEM, f"Failed to load episode {episode_id} from store", {"error": str(e)}, level="warn")
        return None

    def _normalise_weights(self, scores: np.ndarray) -> np.ndarray:
        if scores.size == 0:
            return np.zeros(0, dtype=float)
        tau = float(self.tau)
        if tau <= 0:
            return np.ones_like(scores) / len(scores)
        scaled = scores / tau
        scaled -= np.max(scaled)
        exps = np.exp(scaled)
        total = np.sum(exps)
        if total == 0:
            return np.ones_like(scores) / len(scores)
        normalised = exps / total
        return np.asarray(normalised, dtype=float)

    def _brute_force_retrieve(
        self, query: np.ndarray, k: int
    ) -> tuple[list[Episode], np.ndarray, np.ndarray]:
        E = np.stack([ep.e for ep in self.episodes])
        sims = E @ query / (np.linalg.norm(E, axis=1) *
                            np.linalg.norm(query) + 1e-9)
        weights = self._normalise_weights(sims)
        cache_key = (
    "brute",
    self.namespace,
    hash(
        np.asarray(
            query,
             dtype=float).tobytes()))
        weights = self.stability_monitor.stabilise_weights(
            weights, key=cache_key)
        idx_sorted = np.argsort(-weights)
        top_indices = idx_sorted[:k]
        retrieved = [self.episodes[i] for i in top_indices]
        self.context = np.average(E, axis=0, weights=weights)
        return retrieved, weights, top_indices

    def resonant_retrieve(
        self, query: np.ndarray, k: int = 3, *, namespace: Optional[str] = None
    ) -> list[Episode]:
        """Retrieve top-k episodes using resonance-based similarity.

        Args:
            query: Query vector for retrieval
            k: Number of episodes to retrieve
            namespace: Optional namespace for multi-tenant scenarios

        Returns:
            List of top-k most similar episodes
        """
        with self.use_namespace(namespace):
            if np.allclose(query, 0):
                raise ValueError("Query vector cannot be zero")
            if not self.episodes:
                return []

            retrieved: list[Episode] = []
            weights: np.ndarray = np.array([])
            indices: list[int] | np.ndarray = []

            # Use index if available
            if self._index:
                # Search returns list of (id, score) tuples
                results = self._index.search(query, k)

                # Map IDs to episodes
                id_map = {ep.id: ep for ep in self.episodes}
                retrieved = []
                scores = []

                for eid, score in results:
                    if eid in id_map:
                        retrieved.append(id_map[eid])
                        scores.append(score)

                weights = self._normalise_weights(np.array(scores))

                # Update context based on retrieved episodes
                if retrieved:
                    vectors = np.stack([ep.e for ep in retrieved])
                    self.context = np.average(vectors, axis=0, weights=weights)

                if retrieved:
                    # Emit retrieval event for index hits
                    self._emit(
                        EventType.RETRIEVE,
                        f"Retrieved {len(retrieved)} episodes via index.",
                        {
                            "indices": [ep.id for ep in retrieved],
                            "weights": weights.tolist(),
                            "backend": self._index.index_type if hasattr(self._index, "index_type") else "faiss",
                            "namespace": self.namespace,
                        },
                    )

            # Fallback to brute force when index is unavailable or returns no hits
            if not retrieved:
                retrieved, weights, indices = self._brute_force_retrieve(query, k)

                # Emit retrieval event (brute force)
                self._emit(
                    EventType.RETRIEVE,
                    f"Retrieved {len(retrieved)} episodes via brute force.",
                    {
                        "indices": [int(i) for i in indices],
                        "weights": weights.tolist(),
                        "backend": "numpy",
                        "namespace": self.namespace,
                    },
                )

            # Clean up stale IDs from index
            if self._index and hasattr(self._index, 'ids'):
                valid_ids = {ep.id for ep in self.episodes}
                for episode_id in list(self._index.ids):
                    if episode_id not in valid_ids:
                        self._index.remove(episode_id)

            return retrieved

    def _persist_graph_episode(self, ep: Episode, *, namespace: str) -> None:
        """Persist an episode into the graph store when it is separate."""
        gs = self.graph_store
        if gs is None or gs is self.store:
            return
        save_episode = getattr(gs, "save_episode", None)
        if not callable(save_episode):
            return
        try:
            save_episode(ep, namespace=namespace)
        except TypeError:
            save_episode(ep)
        except Exception as e:
            self._emit(
                EventType.SYSTEM,
                "Failed to persist episode to graph store",
                {"error": str(e), "episode_id": ep.id},
                level="error",
            )

    def _remove_graph_episode(self, episode_id: str, *, namespace: Optional[str]) -> None:
        """Remove an episode from the graph store when it is separate."""
        gs = self.graph_store
        if gs is None or gs is self.store:
            return
        remove_fn = getattr(gs, "remove", None)
        if not callable(remove_fn):
            return
        try:
            remove_fn(episode_id, namespace=namespace)
        except TypeError:
            remove_fn(episode_id)
        except Exception as e:
            self._emit(
                EventType.SYSTEM,
                "Failed to remove episode from graph store",
                {"error": str(e), "episode_id": episode_id},
                level="warn",
            )

    def _persist_trust_update(self, episode_id: str, trust: float, *, namespace: Optional[str]) -> None:
        """Persist trust updates to the underlying store when available."""
        store = self.store
        if store is None:
            return
        update_fn = getattr(store, "update_trust", None)
        if not callable(update_fn):
            return
        try:
            update_fn(episode_id, trust, namespace=namespace)
        except TypeError:
            update_fn(episode_id, trust)
        except Exception as e:
            self._emit(
                EventType.SYSTEM,
                "Failed to update trust in store",
                {"error": str(e), "episode_id": episode_id},
                level="warn",
            )

    def _persist_batch(self, episodes: Sequence[Episode], *, namespace: Optional[str] = None) -> None:
        """Persist episodes to store with proper namespace handling."""
        if not episodes:
            return
        # Use explicit namespace if provided, fall back to self.namespace
        effective_ns = namespace or self.namespace
        if self.store is not None:
            try:
                # For GraphStore implementations, always use save_episode to ensure
                # namespace is properly propagated for graph persistence
                from hieromem_core.persistence.graph_store import GraphStore
                if isinstance(self.store, GraphStore):
                    save_episode = getattr(self.store, "save_episode", None)
                    if callable(save_episode):
                        for ep in episodes:
                            try:
                                save_episode(ep, namespace=effective_ns)
                            except TypeError:
                                save_episode(ep)
                else:
                    # Non-graph stores can use save_many if available
                    save_many = getattr(self.store, "save_many", None)
                    if callable(save_many):
                        save_many(episodes)
                    else:
                        for ep in episodes:
                            self.store.save_episode(ep)
            except Exception as e:
                self._emit(EventType.SYSTEM, "Failed to persist episodes", {"error": str(e)}, level="error")

        for ep in episodes:
            self._persist_graph_episode(ep, namespace=effective_ns)

        ids = [ep.id for ep in episodes]
        matrix = np.stack([np.asarray(ep.e, dtype=np.float32)
                          for ep in episodes])
        metas = [
    ep.meta if isinstance(
        ep.meta,
         Mapping) else {} for ep in episodes]
        self._index_add_batch(ids, matrix, metas=metas)

    def _index_add_episode(self, ep: Episode) -> None:
        index = self.index
        if index is None:
            return
        ids_attr = getattr(index, "ids", [])
        try:
            exists = ep.id in ids_attr
        except TypeError:
            exists = False
        if exists:
            index.remove(ep.id)
        vector = np.asarray(ep.e, dtype=np.float32)[np.newaxis, :]
        self._index_add_batch([ep.id], vector, metas=[
                              ep.meta if isinstance(ep.meta, Mapping) else {}])

    def _index_remove_episode(self, episode_id: str) -> None:
        index = self.index
        if index is None:
            return
        index.remove(episode_id)

    def _index_add_batch(
        self, ids: Sequence[str], matrix: np.ndarray, *, metas: Optional[Sequence[Mapping[str, Any]]]=None
    ) -> None:
        index = self.index
        if index is None or not ids:
            return
        batched = np.asarray(matrix, dtype=np.float32)
        if self.index_router and metas is not None:
            grouped: dict[VectorIndex, tuple[list[str], list[np.ndarray]]] = {}
            for episode_id, row, meta in zip(ids, batched, metas):
                target_index = self.index_router.route(
                    namespace=meta.get("namespace"), encoder_id=meta.get("encoder_id") or self._default_encoder_id()
                )
                if target_index is None:
                    continue
                bucket_ids, bucket_rows = grouped.setdefault(
                    target_index, ([], []))
                bucket_ids.append(episode_id)
                bucket_rows.append(np.asarray(row, dtype=np.float32))
            for target_index, (bucket_ids, bucket_rows) in grouped.items():
                self._add_to_index(
    target_index, bucket_ids, np.stack(bucket_rows))
            return

        self._add_to_index(index, ids, batched)

    def _add_to_index(
    self,
    index: VectorIndex,
    ids: Sequence[str],
     batched: np.ndarray) -> None:
        add_embeddings = getattr(index, "add_embeddings", None)
        add_many = getattr(index, "add_many", None)
        add_single = getattr(index, "add", None)
        if callable(add_embeddings):
            add_embeddings(ids, batched)
            return
        if callable(add_many):
            add_many(ids, batched)
            return
        if callable(add_single):
            for episode_id, row in zip(ids, batched):
                add_single(episode_id, row)
            return
        raise AttributeError("Index backend does not support add operations")

    def _reconcile_index(self) -> None:
        index = self.index
        if index is None:
            return
        valid_ids = {ep.id for ep in self.episodes}
        index_ids = set(index.ids)

        # Remove entries that are no longer persisted
        for stale_id in index_ids - valid_ids:
            index.remove(stale_id)

        # Add missing episodes from the persistence-backed list
        missing_ids = valid_ids - index_ids
        if not missing_ids:
            return
        batch: list[Episode] = [
    ep for ep in self.episodes if ep.id in missing_ids]
        if not batch:
            return
        ids = [ep.id for ep in batch]
        matrix = np.stack([np.asarray(ep.e, dtype=np.float32) for ep in batch])
        self._index_add_batch(ids, matrix)
