evorm_core/plugin.py

#!/usr/bin/env python3
"""
EvoRM Plugin: Evolvable Neuro-Symbolic Reasoning Framework
===========================================================
A plug-and-play module that implements the EvoRM framework from TKDE paper.
Integrates with AdaCoAgentEA's LLM1_label_selector.py.

Core components:
  - RuleEncoding: distill LLM decisions into FOL rules
  - HypergraphStorage: weighted hypergraph for rules + entity contexts
  - TwoStageInferenceController: symbolic filtering + LLM judgment with feedback
  - RuleMaintenance: dynamic confidence tracking + periodic optimization
"""

import os
import json
import re
import time
import math
import hashlib
import threading
from collections import defaultdict
from typing import Dict, List, Tuple, Set, Optional, Any
# ===== EvoRM MLP Gate (TKDE Section III-E) =====
from evorm_core.mlp_gate import MLPGate
# ===== EvoRM Mlight/Mheavy (TKDE Section III-C) =====
from evorm_core.mlight import MlightRationaleElicitor, MheavyDecisionMaker
from evorm_core.entity_embedding import EntityEmbedder
from evorm_core.config import EvoRMConfig
from dataclasses import dataclass, field


# ==============================================================================
# Data Structures
# ==============================================================================

@dataclass
class ConditionAtom:
    """A single condition atom in a FOL rule."""
    atom_type: str  # e.g., "SameValue", "DifferValue", "ShareNeighbor", "DifferNeighbor"
    attr: str       # attribute or relation name
    value1: str = ""  # value for entity 1
    value2: str = ""  # value for entity 2

    def to_key(self) -> str:
        return f"{self.atom_type}({self.attr})"

    def __hash__(self):
        return hash(self.to_key())

    def __eq__(self, other):
        return self.to_key() == other.to_key()


@dataclass
class FOLRule:
    """A first-order logic rule distilled from LLM inference."""
    rule_id: str
    atoms: List[ConditionAtom]
    conclusion: int  # 1 = match, 0 = non-match
    conf: float = 0.0       # global confidence Conf(R)
    conf0: float = 0.0      # initial confidence
    fresh: float = 1.0      # freshness Fresh(R)
    trigger_count: int = 0  # |Trig(R)|
    used_count: int = 0     # times LLM used this rule
    sR_sum: float = 0.0     # sum of contribution scores
    last_used_time: float = 0.0  # timestamp of last positive contribution
    created_time: float = 0.0

    def update_confidence(self):
        """Update global confidence: Conf(R) = (Conf0 + sum(sR)) / (1 + trigger_count)"""
        self.conf = (self.conf0 + self.sR_sum) / (1.0 + self.trigger_count)

    def update_freshness(self, decay_lambda: float = 0.01):
        """Update freshness: Fresh(R) = exp(-lambda * (t_now - t_last))"""
        if self.last_used_time > 0:
            self.fresh = math.exp(-decay_lambda * (time.time() - self.last_used_time))
        else:
            self.fresh = 1.0

    def atom_set(self) -> Set[str]:
        return {a.to_key() for a in self.atoms}

    def __hash__(self):
        return hash(self.rule_id)

    def __eq__(self, other):
        return self.rule_id == other.rule_id


@dataclass
class Hyperedge:
    """A hyperedge in the weighted hypergraph, representing a matching cluster."""
    hyperedge_id: str
    entity_pairs: Set[Tuple[int, int]] = field(default_factory=set)  # Mk
    node_set: Set[str] = field(default_factory=set)  # Vk
    rules: List[FOLRule] = field(default_factory=list)  # Rk
    weight: float = 0.0  # wk

    def recompute_weight(self, alpha: float = 0.6, beta: float = 0.4,
                         entity_embeddings: Dict = None):
        """
        Recompute hyperedge weight: wk = alpha * max(Conf(R)) + beta * avg_cos_sim(Vk)
        
        Paper Eq. 6: The hyperedge weight combines symbolic rule confidence
        with neural embedding similarity of entity pairs in the cluster.
        """
        symbolic = 0.0
        if self.rules:
            symbolic = max(r.conf for r in self.rules)

        neural = 0.0
        if entity_embeddings and self.entity_pairs:
            # Collect embeddings for entities in this hyperedge
            embs = []
            for ep in self.entity_pairs:
                eid1, eid2 = f"e_{ep[0]}", f"e_{ep[1]}"
                if eid1 in entity_embeddings:
                    embs.append(entity_embeddings[eid1])
                if eid2 in entity_embeddings:
                    embs.append(entity_embeddings[eid2])
            
            if len(embs) >= 2:
                # Compute avg pairwise cosine similarity
                from evorm_core.entity_embedding import EntityEmbedder
                embedder = EntityEmbedder()
                neural = embedder.avg_cosine_similarity(embs)
            elif entity_embeddings:
                # Fallback: use rule confidence as proxy
                neural = symbolic * 0.5

        self.weight = alpha * symbolic + beta * neural


# ==============================================================================
# Rule Encoding Module
# ==============================================================================

class RuleEncoding:
    """Convert LLM matching decisions into explicit FOL rules."""

    # Template-based parsing patterns
    ATTR_SAME_PATTERN = re.compile(
        r'(\w+)\s*=\s*same\s*(?:\(([^)]*)\))?', re.IGNORECASE)
    ATTR_DIFFER_PATTERN = re.compile(
        r'(\w+)\s*=\s*differ\s*(?:\(([^)]*)\))?', re.IGNORECASE)
    NEIGHBOR_PATTERN = re.compile(
        r'neighbor\s*:\s*(\w+)', re.IGNORECASE)
    NEIGHBOR_DIFFER_PATTERN = re.compile(
        r'neighbor\s*:\s*differ\s*(\w*)', re.IGNORECASE)
    SHARE_NEIGHBOR_PATTERN = re.compile(
        r'(?:share|has)\s+(?:common\s+)?(?:neighbor|relation)\s*(?::\s*)?(\w*)', re.IGNORECASE)

    def __init__(self, client=None):
        self.client = client
        self.rule_counter = 0

    def _generate_rule_id(self, atoms: List[ConditionAtom]) -> str:
        """Generate a deterministic rule ID from atoms."""
        key = "|".join(sorted(a.to_key() for a in atoms))
        return hashlib.md5(key.encode()).hexdigest()[:16]

    def parse_rationale(self, rationale: str, conclusion: int) -> List[ConditionAtom]:
        """
        Parse a natural language rationale into typed condition atoms.

        Args:
            rationale: natural language rationale from LLM
            conclusion: 1 (match) or 0 (non-match)

        Returns:
            List of ConditionAtom objects
        """
        atoms = []

        # Extract [DECISIVE] and [SUPPORTING] sections
        decisive_section = ""
        supporting_section = ""

        decisive_match = re.search(
            r'\[DECISIVE\](.*?)(?:\[SUPPORTING\]|$)', rationale, re.DOTALL)
        if decisive_match:
            decisive_section = decisive_match.group(1)

        supporting_match = re.search(
            r'\[SUPPORTING\](.*?)$', rationale, re.DOTALL)
        if supporting_match:
            supporting_section = supporting_match.group(1)

        all_text = decisive_section + " " + supporting_section

        # Parse attr=same patterns (V7: extract values from rationale)
        for m in self.ATTR_SAME_PATTERN.finditer(all_text):
            attr = m.group(1).strip().lower()
            val_str = m.group(2).strip() if m.group(2) else ""
            atoms.append(ConditionAtom(
                atom_type="SameValue",
                attr=attr,
                value1=val_str,
                value2=val_str,  # SameValue: both entities have same value
            ))

        # Parse attr=differ patterns (V7: extract values from rationale)
        for m in self.ATTR_DIFFER_PATTERN.finditer(all_text):
            attr = m.group(1).strip().lower()
            val_str = m.group(2).strip() if m.group(2) else ""
            # Parse "value1 vs value2" format
            if " vs " in val_str:
                parts = val_str.split(" vs ", 1)
                v1, v2 = parts[0].strip(), parts[1].strip()
            else:
                v1, v2 = val_str, val_str
            atoms.append(ConditionAtom(
                atom_type="DifferValue",
                attr=attr,
                value1=v1,
                value2=v2,
            ))

        # Parse neighbor:rel patterns
        for m in self.NEIGHBOR_PATTERN.finditer(all_text):
            rel = m.group(1).strip().lower()
            if rel and 'differ' not in rel.lower():
                atoms.append(ConditionAtom(
                    atom_type="ShareNeighbor",
                    attr=rel
                ))

        # Parse neighbor:differ patterns
        for m in self.NEIGHBOR_DIFFER_PATTERN.finditer(all_text):
            rel = m.group(1).strip().lower()
            if rel:
                atoms.append(ConditionAtom(
                    atom_type="DifferNeighbor",
                    attr=rel
                ))

        # If no atoms parsed, create a generic atom based on the rationale
        if not atoms:
            # Try to extract any attribute mentions
            words = re.findall(r'\b\w+\b', all_text)
            # Create a simple atom based on the conclusion
            if conclusion == 1:
                atoms.append(ConditionAtom(
                    atom_type="SameValue",
                    attr="entity_name"
                ))
            else:
                atoms.append(ConditionAtom(
                    atom_type="DifferValue",
                    attr="entity_name"
                ))

        return atoms
    
    def _find_best_key(self, attr: str, es_ctx: Dict, et_ctx: Dict,
                       atom: 'ConditionAtom') -> Optional[str]:
        """
        Find the best matching record key for a rationale attribute name.
        
        Rationale attributes (e.g., "name", "title", "manufacturer") may not
        match the actual schema keys (e.g., "Beer_Name", "Product_Title").
        This method normalizes rationale attrs to actual record keys.
        
        Returns:
            Best matching key, or None if no match found.
        """
        attr_lower = attr.lower().replace('_', ' ').strip()
        all_keys = list(es_ctx.keys()) + list(et_ctx.keys())
        
        # Strategy 1: Exact match
        for key in all_keys:
            if key.lower().replace('_', ' ') == attr_lower:
                return key
        
        # Strategy 2: Substring match (rationale attr is substring of key or vice versa)
        for key in all_keys:
            key_lower = key.lower().replace('_', ' ')
            if attr_lower in key_lower or key_lower in attr_lower:
                return key
        
        # Strategy 3: Token overlap
        attr_tokens = set(attr_lower.split())
        best_key = None
        best_overlap = 0
        for key in all_keys:
            key_tokens = set(key.lower().replace('_', ' ').split())
            overlap = len(attr_tokens & key_tokens)
            if overlap > best_overlap and overlap >= 1:
                best_overlap = overlap
                best_key = key
        
        if best_key:
            return best_key
        
        # Strategy 4: Value search - check if the atom's value1 appears in any key's value
        if atom.value1:
            for key in all_keys:
                val = es_ctx.get(key, '') or et_ctx.get(key, '')
                if val and atom.value1.lower() in str(val).lower():
                    return key
        
        return None
    
    def create_rule(self, rationale: str, conclusion: int,
                    used_rules: List[Tuple['FOLRule', float]] = None,
                    es_context: Dict = None, et_context: Dict = None) -> FOLRule:
        """
        Create a new FOL rule from rationale and LLM feedback.

        Args:
            rationale: natural language rationale
            conclusion: 1 (match) or 0 (non-match)
            used_rules: list of (rule, contribution_score) from LLM feedback

        Returns:
            New FOLRule object
        """
        atoms = self.parse_rationale(rationale, conclusion)
        
        # Enhance generic fallback atoms with context-aware multi-attribute rules
        # When stage2_simple is used, the rationale is just "MATCH"/"NON-MATCH"
        # and parse_rationale creates generic SameValue/DifferValue(entity_name) atoms.
        # These would match ALL pairs, making routing impossible.
        # Instead, create a SINGLE rule with ALL matching/differing attributes.
        # This makes the rule highly specific: it only fires when ALL conditions hold.
        if len(atoms) == 1 and atoms[0].attr == 'entity_name':
            atoms = []
            if es_context and et_context:
                for key in es_context:
                    if key in et_context and not key.startswith('neighbors_'):
                        es_val = str(es_context.get(key, '')).lower().strip()
                        et_val = str(et_context.get(key, '')).lower().strip()
                        if es_val and et_val:
                            if conclusion == 1 and es_val == et_val:
                                atoms.append(ConditionAtom(atom_type="SameValue", attr=key))
                            elif conclusion == 0 and es_val != et_val:
                                atoms.append(ConditionAtom(atom_type="DifferValue", attr=key))
                # V12: Also check keys only in one context (schema mismatch)
                for key in set(list(es_context.keys()) + list(et_context.keys())):
                    if key.startswith('neighbors_'):
                        continue
                    es_val = str(es_context.get(key, '')).lower().strip()
                    et_val = str(et_context.get(key, '')).lower().strip()
                    if es_val and et_val:
                        if conclusion == 1 and es_val == et_val:
                            if not any(a.attr == key for a in atoms):
                                atoms.append(ConditionAtom(atom_type="SameValue", attr=key))
                        elif conclusion == 0 and es_val != et_val:
                            if not any(a.attr == key for a in atoms):
                                atoms.append(ConditionAtom(atom_type="DifferValue", attr=key))
            # If still no atoms, keep the generic one
            if not atoms:
                atoms.append(ConditionAtom(
                    atom_type="SameValue" if conclusion == 1 else "DifferValue",
                    attr="entity_name"))
        
        rule_id = self._generate_rule_id(atoms)

        rule = FOLRule(
            rule_id=rule_id,
            atoms=atoms,
            conclusion=conclusion,
            created_time=time.time(),
            last_used_time=time.time(),
        )

        # Warm-start confidence from overlapping used rules
        if used_rules:
            max_s = 0.0
            for used_rule, sR in used_rules:
                overlap = len(set(a.to_key() for a in atoms) &
                              set(a.to_key() for a in used_rule.atoms))
                union = len(set(a.to_key() for a in atoms) |
                            set(a.to_key() for a in used_rule.atoms))
                if union > 0 and overlap / union > 0:
                    max_s = max(max_s, sR)
            rule.conf0 = max(0.5, max_s)  # minimum 0.5 from LLM feedback
        elif es_context and et_context:
            # Value-level Jaccard overlap for warm-start (paper formula 4)
            es_vals = set(str(v).lower()[:50] for v in es_context.values() if v)
            et_vals = set(str(v).lower()[:50] for v in et_context.values() if v)
            val_overlap = len(es_vals & et_vals)
            val_union = len(es_vals | et_vals)
            val_jaccard = val_overlap / max(1, val_union) if val_union > 0 else 0.0
            # Cap initial confidence at 0.6 - requires LLM validation to increase
            rule.conf0 = min(0.7, max(0.55, val_jaccard))
        else:
            # Default: moderate confidence for new rules
            rule.conf0 = 0.5

        rule.conf = rule.conf0
        return rule

    def elicit_rationale(self, entity_context: str, decision: int,
                         prompt_template: str = None) -> str:
        """
        Use LLM to elicit rationale for a matching decision.

        Args:
            entity_context: serialized entity pair context
            decision: 1 (match) or 0 (non-match)
            prompt_template: optional custom template

        Returns:
            Natural language rationale
        """
        if prompt_template is None:
            decision_str = "MATCH" if decision == 1 else "NON-MATCH"
            prompt_template = f"""Analyze the following entity pair and explain why they are a {decision_str}.

Entity Context:
{entity_context}

Please output your analysis in the following format:
[DECISIVE] List the core attributes or relations that are pivotal to the {decision_str} decision.
[SUPPORTING] List any auxiliary evidence that corroborates the decision.

Use the format: attr=same(value), attr=differ(value1 vs value2), or neighbor:rel_name."""
        
        if self.client:
            try:
                response = self.client.chat.completions.create(
                    model="gpt-3.5-turbo",
                    messages=[{'role': 'user', 'content': prompt_template}],
                    temperature=0.1
                )
                return response.choices[0].message.content.strip()
            except Exception as e:
                print(f"RuleEncoding: rationale elicitation failed: {e}")
                return f"[DECISIVE] entity_name=same\n[SUPPORTING] automatic fallback"

        # Fallback without client
        return f"[DECISIVE] entity_name={'same' if decision == 1 else 'differ'}\n[SUPPORTING] automatic fallback"


# ==============================================================================
# Hypergraph Storage Module
# ==============================================================================

class HypergraphStorage:
    """Weighted hypergraph for organizing rules and entity contexts."""

    def __init__(self, alpha: float = 0.6, beta: float = 0.4,
                 merge_threshold: float = 0.5,
                 entity_embedder: 'EntityEmbedder' = None):
        self.hyperedges: Dict[str, Hyperedge] = {}
        self.rules: Dict[str, FOLRule] = {}
        self.inverted_index: Dict[str, Set[str]] = defaultdict(set)  # entity -> hyperedge IDs
        self.alpha = alpha
        self.beta = beta
        self.merge_threshold = merge_threshold  # eta_h
        self.edge_counter = 0
        self.entity_embedder = entity_embedder
        self.entity_embeddings: Dict[str, 'np.ndarray'] = {}  # entity_id -> embedding
        self.flat_mode = False  # w/o Hypergraph ablation (Table VI "w/o U")

    def _generate_edge_id(self) -> str:
        self.edge_counter += 1
        return f"HE_{self.edge_counter:06d}"

    def get_or_create_hyperedge(self, entity_pair: Tuple[int, int],
                                node_set: Set[str],
                                rule: FOLRule) -> Hyperedge:
        """
        Find the most similar existing hyperedge or create a new one.

        Args:
            entity_pair: (es, et) tuple
            node_set: VR - nodes extracted from entity context
            rule: the new FOL rule

        Returns:
            The matched or newly created hyperedge
        """
        # Flat mode (w/o Hypergraph ablation): don't create hyperedges
        if self.flat_mode:
            # Still register the rule globally
            if rule.rule_id not in self.rules:
                self.rules[rule.rule_id] = rule
            # Return a dummy hyperedge (not stored)
            dummy = Hyperedge(
                hyperedge_id='flat',
                entity_pairs={entity_pair},
                node_set=node_set,
                rules=[rule],
            )
            return dummy

        # Find the most structurally similar hyperedge
        best_id = None
        best_overlap = 0.0

        for he_id, he in self.hyperedges.items():
            if not he.node_set:
                continue
            overlap = len(node_set & he.node_set)
            union = len(node_set | he.node_set)
            if union > 0:
                jaccard = overlap / union
                if jaccard > best_overlap:
                    best_overlap = jaccard
                    best_id = he_id

        if best_id and best_overlap >= self.merge_threshold:
            # Update existing hyperedge
            he = self.hyperedges[best_id]
            he.entity_pairs.add(entity_pair)
            he.node_set.update(node_set)
            if rule not in he.rules:
                he.rules.append(rule)
            he.recompute_weight(self.alpha, self.beta, self.entity_embeddings)

            # Update inverted index
            for node in node_set:
                self.inverted_index[node].add(best_id)
            # Also index by rule atom attribute names for cross-query matching
            for atom in rule.atoms:
                attr_key = f'__attr__:{atom.attr.lower()}'
                self.inverted_index[attr_key].add(best_id)

            return he
        else:
            # Create new hyperedge
            new_id = self._generate_edge_id()
            he = Hyperedge(
                hyperedge_id=new_id,
                entity_pairs={entity_pair},
                node_set=node_set,
                rules=[rule],
            )
            he.recompute_weight(self.alpha, self.beta, self.entity_embeddings)
            self.hyperedges[new_id] = he

            # Update inverted index
            for node in node_set:
                self.inverted_index[node].add(new_id)
            # Also index by rule atom attribute names for cross-query matching
            for atom in rule.atoms:
                attr_key = f'__attr__:{atom.attr.lower()}'
                self.inverted_index[attr_key].add(new_id)

            # Also register the rule
            if rule.rule_id not in self.rules:
                self.rules[rule.rule_id] = rule

            return he

    def get_candidate_hyperedges(self, entity_ids: List[str],
                                    entity_contexts: List[Dict] = None) -> List[Hyperedge]:
        """
        Retrieve candidate hyperedges containing any of the given entity IDs.
        O(1) complexity via inverted index.

        Args:
            entity_ids: list of entity ID strings

        Returns:
            List of candidate hyperedges
        """
        candidate_ids: Set[str] = set()
        for eid in entity_ids:
            candidate_ids.update(self.inverted_index.get(eid, set()))

        # Attribute-value-based fallback for cross-query rule matching
        if entity_contexts:
            for ctx in entity_contexts:
                if not ctx:
                    continue
                for key in ['entity_name', 'name', 'title', 'description']:
                    val = ctx.get(key, '')
                    if val and isinstance(val, str):
                        candidate_ids.update(self.inverted_index.get(val.lower()[:100], set()))
                        for token in val.lower().split()[:5]:
                            if len(token) > 2:
                                candidate_ids.update(self.inverted_index.get(token, set()))

        # V8 FIX: Look up by attribute keys using the SAME format as indexing.
        # Indexing uses f"a_{{key}}" (see record_trajectory node_set).
        # Previously used f"__attr__:{{key.lower()}}" which NEVER matched.
        # Also removed "if not candidate_ids" guard — attribute lookup must
        # always run for cross-query matching in ER where entity IDs and values
        # are unique per pair.
        if entity_contexts:
            attr_keys_seen = set()
            for ctx in entity_contexts:
                if not ctx:
                    continue
                for key in ctx:
                    if key.startswith('__') or key.startswith('neighbors_'):
                        continue
                    attr_keys_seen.add(key)
            for key in attr_keys_seen:
                idx_key = f'a_{key}'
                found = self.inverted_index.get(idx_key, set())
                if found:
                    candidate_ids.update(found)

        return [self.hyperedges[hid] for hid in candidate_ids
                if hid in self.hyperedges]



    def store_entity_embedding(self, entity_id: str, embedding: 'np.ndarray'):
        """Store entity embedding for neural similarity computation."""
        self.entity_embeddings[entity_id] = embedding

    def get_entity_embeddings(self) -> Dict:
        """Get all stored entity embeddings."""
        return self.entity_embeddings

    def get_candidate_rules(self, entity_ids: List[str],
                             entity_contexts: List[Dict] = None) -> List[FOLRule]:
        """
        Get all candidate rules from hyperedges containing the entity IDs.

        Args:
            entity_ids: list of entity ID strings
            entity_contexts: optional entity context dicts for attribute-based matching

        Returns:
            List of candidate FOL rules
        """
        # Flat mode (w/o Hypergraph): return all rules
        if self.flat_mode:
            return list(self.rules.values())

        candidates = self.get_candidate_hyperedges(entity_ids, entity_contexts)
        rules_set: Dict[str, FOLRule] = {}
        for he in candidates:
            for rule in he.rules:
                if rule.rule_id not in rules_set:
                    rules_set[rule.rule_id] = rule
        return list(rules_set.values())

    def register_rule(self, rule: FOLRule):
        """Register a rule in the global rule set."""
        if rule.rule_id not in self.rules:
            self.rules[rule.rule_id] = rule

    def get_rule(self, rule_id: str) -> Optional[FOLRule]:
        return self.rules.get(rule_id)

    def stats(self) -> Dict:
        return {
            'num_hyperedges': len(self.hyperedges),
            'num_rules': len(self.rules),
            'num_indexed_entities': len(self.inverted_index),
            'avg_rules_per_edge': (
                sum(len(he.rules) for he in self.hyperedges.values()) /
                max(1, len(self.hyperedges))
            ),
        }


# ==============================================================================
# Two-Stage Inference Controller
# ==============================================================================

class TwoStageInferenceController:
    """Two-stage controller: symbolic filtering + LLM judgment with feedback."""

    def __init__(self, hypergraph: HypergraphStorage,
                 client=None,
                 theta_hi: float = 0.65,  # V11: balanced threshold for match routing
                 theta_prune: float = 0.75,  # V12: conservative non-match routing
                 theta_gate: float = 0.3,
                 route_nonmatch: bool = True,  # V12: allow S1 to route non-match decisions
                 K: int = 5,
                 use_mlight: bool = True):
        self.hypergraph = hypergraph
        self.client = client
        self.theta_hi = theta_hi       # high confidence threshold for direct match
        self.theta_prune = theta_prune # threshold for direct non-match
        self.theta_gate = theta_gate   # MLP gate threshold
        self.route_nonmatch = route_nonmatch  # V12: allow S1 to route non-match decisions
        self.K = K                     # top-K hyperedges for evidence subgraph

        # MLP Gate (gϕ) — paper Section III-E
        self.mlp_gate = None  # Set externally by EvoRMPlugin
        self.use_mlp_gate = False  # Disabled by default (enable after training)

        # Mlight / Mheavy — paper Section III-C (rationale elicitation separated from decision)
        self.use_mlight = use_mlight
        self.mlight = None  # Set externally by EvoRMPlugin
        self.mheavy = None  # Set externally by EvoRMPlugin

        # Statistics
        self.stage1_hits = 0
        self.stage1_total = 0
        self.stage1_correct = 0  # V5: track correct Stage-1 decisions
        self.stage2_hits = 0
        self.stage2_total = 0

    def verify_rule(self, rule: FOLRule, es_context: Dict,
                    et_context: Dict) -> bool:
        """
        Verify if a rule is triggered for the given entity pair.

        Args:
            rule: FOL rule to verify
            es_context: source entity context (dict with attributes)
            et_context: target entity context (dict with attributes)

        Returns:
            True if all atoms in the rule hold
        """
        for atom in rule.atoms:
            if not self._verify_atom(atom, es_context, et_context):
                return False
        return True

    @staticmethod
    def _token_overlap(v1: str, v2: str) -> float:
        """Compute Jaccard token overlap between two strings."""
        t1 = set(v1.lower().split())
        t2 = set(v2.lower().split())
        if not t1 or not t2:
            return 0.0
        return len(t1 & t2) / len(t1 | t2)


    def _verify_atom(self, atom: ConditionAtom, es_ctx: Dict,
                     et_ctx: Dict) -> bool:
        """Verify a single condition atom (V12: _find_best_key for schema normalization)."""
        attr = atom.attr

        if atom.atom_type == "SameValue":
            v1 = es_ctx.get(attr, "")
            v2 = et_ctx.get(attr, "")
            
            # V12: Use _find_best_key to normalize schema-mismatched attributes
            if not v1 and not v2:
                best_key = self._find_best_key(attr, es_ctx, et_ctx, atom)
                if best_key:
                    v1 = es_ctx.get(best_key, "")
                    v2 = et_ctx.get(best_key, "")
                if not v1 and not v2:
                    # Both missing as keys — use fallback comparison
                    return self._verify_samevalue_fallback(attr, atom.value1, atom.value2, es_ctx, et_ctx)
            elif not v1 or not v2:
                # One side has the key, other doesn't — try normalize
                best_key = self._find_best_key(attr, es_ctx, et_ctx, atom)
                if best_key:
                    v1 = es_ctx.get(best_key, "") or v1
                    v2 = et_ctx.get(best_key, "") or v2
                if not v1 or not v2:
                    return False
            
            # Direct key comparison (existing logic)
            v1l, v2l = v1.lower().strip(), v2.lower().strip()
            if v1l == v2l:
                return True
            if v1l in v2l or v2l in v1l:
                return True
            overlap = self._token_overlap(v1, v2)
            return overlap >= 0.5

        elif atom.atom_type == "DifferValue":
            v1 = es_ctx.get(attr, "")
            v2 = et_ctx.get(attr, "")
            
            # V12: Use _find_best_key to normalize schema-mismatched attributes
            if not v1 and not v2:
                best_key = self._find_best_key(attr, es_ctx, et_ctx, atom)
                if best_key:
                    v1 = es_ctx.get(best_key, "")
                    v2 = et_ctx.get(best_key, "")
                if not v1 and not v2:
                    return self._verify_differvalue_fallback(attr, atom.value1, atom.value2, es_ctx, et_ctx)
            elif not v1 or not v2:
                best_key = self._find_best_key(attr, es_ctx, et_ctx, atom)
                if best_key:
                    v1 = es_ctx.get(best_key, "") or v1
                    v2 = et_ctx.get(best_key, "") or v2
                if not v1 or not v2:
                    return False
            
            # Direct key comparison (existing logic)
            v1l, v2l = v1.lower().strip(), v2.lower().strip()
            if v1l != v2l:
                overlap = self._token_overlap(v1, v2)
                if overlap < 0.4:
                    return True
            return False

        elif atom.atom_type == "ShareNeighbor":
            es_neighbors = es_ctx.get(f"neighbors_{attr}", set())
            et_neighbors = et_ctx.get(f"neighbors_{attr}", set())
            return bool(es_neighbors & et_neighbors)

        elif atom.atom_type == "DifferNeighbor":
            es_neighbors = es_ctx.get(f"neighbors_{attr}", set())
            et_neighbors = et_ctx.get(f"neighbors_{attr}", set())
            return not bool(es_neighbors & et_neighbors) and es_neighbors and et_neighbors

        elif atom.atom_type == "SemanticEquiv":
            v1 = es_ctx.get(attr, "")
            v2 = et_ctx.get(attr, "")
            if v1 and v2:
                tokens1 = set(v1.lower().split())
                tokens2 = set(v2.lower().split())
                if tokens1 and tokens2:
                    overlap = len(tokens1 & tokens2) / len(tokens1 | tokens2)
                    return overlap > 0.5
            return False

        elif atom.atom_type == "SemanticConflict":
            v1 = es_ctx.get(attr, "")
            v2 = et_ctx.get(attr, "")
            if v1 and v2:
                tokens1 = set(v1.lower().split())
                tokens2 = set(v2.lower().split())
                if tokens1 and tokens2:
                    overlap = len(tokens1 & tokens2) / len(tokens1 | tokens2)
                    return overlap < 0.2
            return False

        return False

    def _verify_samevalue_fallback(self, attr: str, val1: str, val2: str,
                                    es_ctx: Dict, et_ctx: Dict) -> bool:
        """Fallback SameValue verification when attribute is not a direct key.
        
        Strategies (in order):
        1. If value hints provided: check if both contexts contain the same value
        2. If no value hints: check if both contexts contain the attribute name
           and compare the surrounding tokens
        """
        # Strategy 1: Value-based comparison
        if val1 and val2 and len(val1) > 2 and len(val2) > 2:
            val1_lower = val1.lower().strip()
            val2_lower = val2.lower().strip()
            found_in_es = self._find_value_in_context(val1_lower, es_ctx)
            found_in_et = self._find_value_in_context(val2_lower, et_ctx)
            if found_in_es and found_in_et:
                # Both values found — check if they refer to the same thing
                if val1_lower == val2_lower:
                    return True
                # Different values but both found — check surrounding context
                w1 = self._extract_context_window(found_in_es, val1_lower)
                w2 = self._extract_context_window(found_in_et, val2_lower)
                if w1 and w2:
                    return self._token_overlap(w1, w2) >= 0.5
                return False
            return False
        
        # Strategy 2: Attribute name based — both contexts mention the attribute
        if len(attr) > 2:
            found_in_es = self._find_value_in_context(attr.lower(), es_ctx)
            found_in_et = self._find_value_in_context(attr.lower(), et_ctx)
            if found_in_es and found_in_et:
                # Both mention the attribute — compare surrounding context
                w1 = self._extract_context_window(found_in_es, attr.lower())
                w2 = self._extract_context_window(found_in_et, attr.lower())
                if w1 and w2:
                    return self._token_overlap(w1, w2) >= 0.4
                return True  # Both mention the attribute, weak signal
            return False
        
        return False

    def _verify_differvalue_fallback(self, attr: str, val1: str, val2: str,
                                      es_ctx: Dict, et_ctx: Dict) -> bool:
        """Fallback DifferValue verification when attribute is not a direct key."""
        # Strategy 1: Value-based comparison
        if val1 and val2 and len(val1) > 2 and len(val2) > 2:
            val1_lower = val1.lower().strip()
            val2_lower = val2.lower().strip()
            if val1_lower == val2_lower:
                return False  # Same values = not a differ
            found_in_es = self._find_value_in_context(val1_lower, es_ctx)
            found_in_et = self._find_value_in_context(val2_lower, et_ctx)
            if found_in_es and found_in_et:
                # Different values, both found in respective contexts
                return True
            return False
        
        # Strategy 2: Attribute name based — both mention the attribute
        if len(attr) > 2:
            found_in_es = self._find_value_in_context(attr.lower(), es_ctx)
            found_in_et = self._find_value_in_context(attr.lower(), et_ctx)
            if found_in_es and found_in_et:
                # Both mention attribute — compare surrounding context
                w1 = self._extract_context_window(found_in_es, attr.lower())
                w2 = self._extract_context_window(found_in_et, attr.lower())
                if w1 and w2:
                    return self._token_overlap(w1, w2) < 0.3
                return False
            return False
        
        return False

    @staticmethod
    def _find_value_in_context(target: str, ctx: Dict) -> str:
        """Find a target string in context values. Returns the containing value.
        
        Uses four strategies:
        1. Exact substring match in values (fast path)
        2. Token-level exact match in values
        3. Token-level substring match in values (handles singular/plural)
        4. Key name matching: if target tokens overlap with a context KEY,
           return the value of that key (handles 'artist_name' -> 'artist' key)
        """
        if not ctx or not target:
            return ""
        target_lower = target.lower().strip()
        target_tokens = set(target_lower.replace('_', ' ').split())
        
        for key, val in ctx.items():
            if not isinstance(val, str):
                continue
            val_lower = val.lower()
            # Strategy 1: Exact substring in value
            if target_lower in val_lower:
                return val
            # Strategy 2 & 3: Token-level matching in value
            if target_tokens:
                val_tokens = set(val_lower.split())
                # Strategy 2: Exact token match
                if target_tokens & val_tokens:
                    return val
                # Strategy 3: Substring token match (handles stemming variations)
                for tt in target_tokens:
                    if len(tt) < 3:
                        continue
                    for vt in val_tokens:
                        if len(vt) < 3:
                            continue
                        if tt in vt or vt in tt:
                            return val
        
        # Strategy 4: Key name matching (handles 'artist_name' -> 'artist' key)
        if target_tokens:
            for key, val in ctx.items():
                if not isinstance(val, str) or not val.strip():
                    continue
                key_lower = key.lower().strip()
                key_tokens = set(key_lower.replace('_', ' ').split())
                # Check if target tokens overlap with key tokens
                # e.g., target="artist_name" (tokens: artist, name), key="artist" (tokens: artist)
                if target_tokens & key_tokens:
                    return val
                # Also check substring: "artist" in "artist_name" or vice versa
                for tt in target_tokens:
                    if len(tt) < 3:
                        continue
                    for kt in key_tokens:
                        if len(kt) < 3:
                            continue
                        if tt in kt or kt in tt:
                            return val
        
        return ""

    @staticmethod
    def _extract_context_window(text: str, target: str, window: int = 6) -> str:
        """Extract a window of words around the target in text.
        
        Uses token-level matching to handle singular/plural variations.
        """
        import re as _re
        words = text.lower().split()
        target_tokens = set(target.lower().replace('_', ' ').split())
        
        # First try exact substring match
        for i, w in enumerate(words):
            if target in w:
                start = max(0, i - window)
                end = min(len(words), i + window + 1)
                return " ".join(words[start:end])
        
        # Fallback: token-level match (strip punctuation, check overlap)
        if target_tokens:
            for i, w in enumerate(words):
                w_clean = _re.sub(r'[^a-z0-9]', '', w)
                if w_clean in target_tokens or any(t in w_clean for t in target_tokens if len(t) > 3):
                    start = max(0, i - window)
                    end = min(len(words), i + window + 1)
                    return " ".join(words[start:end])
        
        return text.lower()

    def stage1_symbolic_filtering(self, es_id: str, et_id: str,
                                  es_context: Dict, et_context: Dict
                                  ) -> Tuple[Optional[int], List[FOLRule], List[FOLRule]]:
        """
        Stage 1: Symbolic filtering and gated routing.

        Returns:
            (decision, triggered_rules, candidate_rules)
            decision: None if survival (needs Stage 2), 0 or 1 if direct routing
        """
        self.stage1_total += 1

        # Candidate retrieval
        candidate_rules = self.hypergraph.get_candidate_rules([es_id, et_id], [es_context, et_context])

        # Rule verification
        triggered: List[FOLRule] = []
        for rule in candidate_rules:
            if self.verify_rule(rule, es_context, et_context):
                triggered.append(rule)

        if not triggered:
            return None, [], candidate_rules

        # V8: Self-supervised confidence boost for verified rules.
        # When a rule's atoms are verified against actual data, it's evidence
        # the rule is applicable. Boost confidence by a small amount.
        # This is critical for cold-start: rules start at conf=0.55 and
        # need to accumulate evidence to reach routing thresholds.
        for rule in triggered:
            if rule.trigger_count > 0:
                # Move confidence 5% toward the verification signal
                rule.conf = rule.conf + 0.05 * (1.0 - rule.conf)

        # Check verdicts
        match_rules = [r for r in triggered if r.conclusion == 1]
        nonmatch_rules = [r for r in triggered if r.conclusion == 0]

        # V9: Only route when there's sufficient evidence diversity.
        # Cold-start protection: if all triggered rules are from one class
        # and total trigger count is low, be conservative.
        total_triggers = sum(r.trigger_count for r in triggered)
        cold_start = False  # V10: No cold-start, rely on confidence thresholds alone
        
        # Direct Matching
        # V9: Route when confidence is high enough. During cold start, require
        # higher confidence (0.7) to prevent premature routing.
        if match_rules and not nonmatch_rules:
            max_conf = max(r.conf for r in match_rules)
            req_conf = 0.70 if cold_start else self.theta_hi
            if max_conf >= req_conf:
                self.stage1_hits += 1
                return 1, triggered, candidate_rules

        # Direct Non-Matching
        # V12: Only route non-match when route_nonmatch is enabled
        # Non-match routing is conservative: requires theta_prune >= 0.75
        # and rule confidence >= 0.75 (5+ successful triggers)
        if self.route_nonmatch and nonmatch_rules and not match_rules:
            max_conf = max(r.conf for r in nonmatch_rules)
            # Require at least 5 triggers for non-match routing confidence
            max_triggers = max(r.trigger_count for r in nonmatch_rules)
            req_conf = self.theta_prune
            if max_conf >= req_conf and max_triggers >= 5:
                self.stage1_hits += 1
                return 0, triggered, candidate_rules

        # Survival - check MLP gate before deferring to Stage 2
        if self.use_mlp_gate and self.mlp_gate is not None and self.mlp_gate.is_trained:
            try:
                features = self.mlp_gate.extract_features(es_context, et_context, triggered)
                if not self.mlp_gate.should_invoke_llm(features):
                    self.stage1_hits += 1
                    return 0, triggered, candidate_rules
            except Exception as _mlp_e:
                pass  # MLP training error (non-critical)
        
        return None, triggered, candidate_rules

    def construct_evidence_subgraph(self, es_id: str, et_id: str,
                                    es_context: Dict, et_context: Dict,
                                    candidate_rules: List[FOLRule]) -> Dict:
        """
        Construct a compact evidence subgraph for LLM inference.

        Args:
            es_id: source entity ID
            et_id: target entity ID
            candidate_rules: candidate rules

        Returns:
            Evidence subgraph data
        """
        # Get hyperedges
        candidate_edges = self.hypergraph.get_candidate_hyperedges([es_id, et_id], [es_context, et_context])

        # Prioritize edges containing both entities
        joint_edges = []
        other_edges = []
        for he in candidate_edges:
            if es_id in he.node_set and et_id in he.node_set:
                joint_edges.append(he)
            else:
                other_edges.append(he)

        # Sort by weight
        joint_edges.sort(key=lambda x: x.weight, reverse=True)
        other_edges.sort(key=lambda x: x.weight, reverse=True)

        # Select top-K
        selected = joint_edges[:self.K]
        if len(selected) < self.K:
            selected.extend(other_edges[:self.K - len(selected)])

        return {
            'hyperedges': selected,
            'rules': candidate_rules,
        }

    def build_stage2_prompt(self, es_id: str, et_id: str,
                            es_context: Dict, et_context: Dict,
                            evidence: Dict,
                            triggered_rules: List[FOLRule],
                            base_prompt: str) -> str:
        """
        Build the Stage 2 prompt with evidence subgraph and rule metadata.

        Args:
            es_id, et_id: entity IDs
            es_context, et_context: entity contexts
            evidence: evidence subgraph from construct_evidence_subgraph
            triggered_rules: triggered rules
            base_prompt: original prompt from AdaCoAgentEA

        Returns:
            Augmented prompt
        """
        # Build rule metadata section
        rule_section = ""
        if triggered_rules:
            rule_section = "\n\n**Triggered Rules (from historical matching experience):**\n"
            for i, rule in enumerate(triggered_rules, 1):
                atom_strs = []
                for a in rule.atoms:
                    atom_strs.append(f"{a.atom_type}({a.attr})")
                atoms_str = " ∧ ".join(atom_strs)
                verdict = "MATCH" if rule.conclusion == 1 else "NON-MATCH"
                rule_section += (
                    f"  Rule {i}: {atoms_str} → {verdict}\n"
                    f"    Confidence: {rule.conf:.3f} | "
                    f"Triggered: {rule.trigger_count} times\n"
                )

        # Build evidence summary
        evidence_section = ""
        if evidence.get('hyperedges'):
            evidence_section = "\n\n**Evidence from Similar Historical Cases:**\n"
            for he in evidence['hyperedges'][:3]:
                matching_pairs = len(he.entity_pairs)
                evidence_section += (
                    f"  - Cluster (weight={he.weight:.3f}): "
                    f"{matching_pairs} similar entity pairs stored\n"
                )

        # Build the feedback instruction
        feedback_instruction = """
        
**IMPORTANT: Rule Contribution Feedback**
After making your decision, you MUST also provide a JSON object with contribution scores (0.0 to 1.0) for each triggered rule that influenced your decision. Format:
```json
{
  "decision": "MATCH" or "NON-MATCH",
  "rule_feedback": {
    "rule_1": 0.8,
    "rule_2": 0.3
  },
  "reasoning": "[DECISIVE] ... [SUPPORTING] ..."
}
```
Score meaning: 1.0 = DECISIVE (rule was the key factor), 0.1-0.9 = SUPPORTING (rule partially influenced), 0.0 = IRRELEVANT (rule was considered but not used).

You MUST respond with valid JSON only."""

        return base_prompt + rule_section + evidence_section + feedback_instruction

    def parse_llm_feedback(self, response_text: str
                           ) -> Tuple[Optional[int], Dict[str, float], str]:
        """
        Parse the LLM's JSON response to extract decision, rule feedback, and rationale.

        Returns:
            (decision, rule_feedback dict, rationale)
        """
        # Try to extract JSON from response
        try:
            # Find JSON block
            json_match = re.search(r'```json\s*(.*?)\s*```', response_text, re.DOTALL)
            if json_match:
                json_str = json_match.group(1)
            else:
                # Try to find bare JSON
                json_match = re.search(r'\{.*"decision".*\}', response_text, re.DOTALL)
                if json_match:
                    json_str = json_match.group(0)
                else:
                    json_str = response_text

            data = json.loads(json_str)
            decision_str = data.get('decision', '').upper()
            decision = 1 if 'MATCH' in decision_str and 'NON' not in decision_str else 0

            rule_feedback = data.get('rule_feedback', {})
            reasoning = data.get('reasoning', '')

            return decision, rule_feedback, reasoning
        except (json.JSONDecodeError, KeyError):
            pass

        # Fallback: try to parse from plain text
        decision = None
        if 'match' in response_text.lower() and 'non-match' not in response_text.lower():
            decision = 1
        elif 'non-match' in response_text.lower() or 'no match' in response_text.lower():
            decision = 0

        return decision, {}, response_text

    def stage2_elicit_rationale(self, es_id: str, et_id: str,
                                es_context: Dict, et_context: Dict,
                                triggered_rules: List[FOLRule]):
        """
        Stage 2a: Elicit rationale using Mlight (independent from decision).

        Paper Section III-C Step 1: Mlight generates a natural language
        rationale explaining the matching evidence, WITHOUT making a decision.

        Returns:
            RationaleResult with structured analysis
        """
        if self.mlight is not None:
            return self.mlight.elicit(es_context, et_context, triggered_rules)
        else:
            from evorm_core.mlight import RationaleResult
            return RationaleResult(
                rationale="[DECISIVE] entity_name=same\\n[SUPPORTING] automatic fallback (no Mlight)",
                token_usage=0,
            )

    def stage2_llm_judgment(self, es_id: str, et_id: str,
                            es_context: Dict, et_context: Dict,
                            triggered_rules: List[FOLRule],
                            candidate_rules: List[FOLRule],
                            base_prompt: str) -> Dict:
        """
        Stage 2: LLM judgment with rule feedback.

        Returns:
            Dict with decision, rule_feedback, rationale, token_usage
        """
        self.stage2_total += 1

        # Increment trigger counts
        for rule in triggered_rules:
            rule.trigger_count += 1

        # ================================================================
        # Two-step Mlight-Mheavy pipeline (Paper Section III-C)
        # ================================================================
        if self.use_mlight and self.mlight is not None and self.mheavy is not None:
            # OPTIMIZATION: Skip Mlight when no rules triggered
            # Mlight is designed to analyze triggered rules; without rules,
            # its rationale adds little value while costing 1 LLM call.
            if triggered_rules:
                # Step 1: Mlight - elicit rationale (no decision)
                rationale_result = self.stage2_elicit_rationale(
                    es_id, et_id, es_context, et_context, triggered_rules)
                rationale_text = rationale_result.rationale
                total_tokens = rationale_result.token_usage
            else:
                # No rules to analyze - use a minimal rationale
                rationale_text = "[DECISIVE] Direct attribute comparison\n[SUPPORTING] No historical rules available"
                total_tokens = 0

            # Step 2: Mheavy - make decision with rationale as context
            decision_result = self.mheavy.decide(
                es_context, et_context,
                rationale_text,
                triggered_rules,
                base_prompt=base_prompt)

            total_tokens += decision_result.get('token_usage', 0)

            # Track tokens for EvoRM experiments
            try:
                from utils import tokens as tokens_cal
                tokens_cal.update_add_var(total_tokens)
            except Exception as _mlp_e:
                pass

            if decision_result.get('decision') is not None:
                self.stage2_hits += 1

            # When Mlight is skipped, use Mheavy summary as rationale for better rule creation
            effective_rationale = rationale_text
            if not triggered_rules:
                mheavy_summary = decision_result.get('summary', '')
                if mheavy_summary and mheavy_summary != 'no client available':
                    effective_rationale = f"[DECISIVE] {mheavy_summary}\n[SUPPORTING] Direct LLM reasoning"

            return {
                'decision': decision_result.get('decision'),
                'rule_feedback': decision_result.get('rule_feedback', {}),
                'rationale': effective_rationale,
                'confidence': decision_result.get('confidence', 0.5),
                'token_usage': total_tokens,
                'raw_response': decision_result.get('raw_response', ''),
                'mlight_raw': rationale_text if not triggered_rules else rationale_result.raw_response,
                'method': 'mheavy-only' if not triggered_rules else 'mlight+mheavy',
            }

        # ================================================================
        # Legacy single-call approach (fallback)
        # ================================================================
        # Construct evidence subgraph
        evidence = self.construct_evidence_subgraph(
            es_id, et_id, candidate_rules)

        # Build augmented prompt
        prompt = self.build_stage2_prompt(
            es_id, et_id, es_context, et_context,
            evidence, triggered_rules, base_prompt)

        if self.client:
            try:
                response = self.client.chat.completions.create(
                    model="gpt-3.5-turbo",
                    messages=[{'role': 'user', 'content': prompt}],
                    temperature=0.1,
                    response_format={"type": "json_object"},
                )
                response_text = response.choices[0].message.content.strip()
                tokens = response.usage.total_tokens

                decision, rule_feedback, rationale = self.parse_llm_feedback(
                    response_text)

                if decision is not None:
                    self.stage2_hits += 1

                return {
                    'decision': decision,
                    'rule_feedback': rule_feedback,
                    'rationale': rationale,
                    'token_usage': tokens,
                    'raw_response': response_text,
                    'method': 'single-call',
                }
            except Exception as e:
                print(f"Stage 2 LLM call failed: {e}")
                return {
                    'decision': None,
                    'rule_feedback': {},
                    'rationale': '',
                    'token_usage': 0,
                    'raw_response': '',
                    'error': str(e),
                    'method': 'error',
                }

        # Fallback without client
        return {
            'decision': None,
            'rule_feedback': {},
            'rationale': 'no client available',
            'token_usage': 0,
            'raw_response': '',
            'method': 'no-client',
        }

    def get_stats(self) -> Dict:
        stats = {
            'stage1_total': self.stage1_total,
            'stage1_hits': self.stage1_hits,
            'stage1_correct': self.stage1_correct,
            'stage1_accuracy': self.stage1_correct / max(1, self.stage1_hits),
            'stage1_rate': self.stage1_hits / max(1, self.stage1_total),
            'stage2_total': self.stage2_total,
            'stage2_hits': self.stage2_hits,
            'use_mlight': self.use_mlight,
        }
        if self.mlight is not None:
            stats['mlight'] = self.mlight.get_stats()
        if self.mheavy is not None:
            stats['mheavy'] = self.mheavy.get_stats()
        return stats


# ==============================================================================
# Rule Maintenance Module
# ==============================================================================

class RuleMaintenance:
    """Dynamic rule maintenance: confidence tracking + periodic optimization."""

    def __init__(self, hypergraph: HypergraphStorage,
                 decay_lambda: float = 0.01,
                 freshness_threshold: float = 0.1,
                 confidence_threshold: float = 0.2,
                 eval_triggers: int = 10,
                 merge_similarity: float = 0.7,
                 max_rules: int = 10000,
                 evict_percentile: float = 0.1):
        self.hypergraph = hypergraph
        self.decay_lambda = decay_lambda
        self.theta_f = freshness_threshold
        self.theta_c = confidence_threshold
        self.N_eval = eval_triggers
        self.eta_m = merge_similarity
        self.max_rules = max_rules
        self.evict_percentile = evict_percentile

        self.last_optimization_time = time.time()
        self.optimization_interval = 300  # 5 minutes

    def update_rule_confidence(self, rule_id: str, sR: float):
        """
        Update rule confidence with LLM contribution score.

        Args:
            rule_id: the rule to update
            sR: contribution score from LLM (0.0 to 1.0)
        """
        rule = self.hypergraph.get_rule(rule_id)
        if rule is None:
            return

        rule.sR_sum += sR
        if sR > 0:
            rule.used_count += 1
            rule.last_used_time = time.time()

        rule.update_confidence()
        rule.update_freshness(self.decay_lambda)

    def should_optimize(self) -> bool:
        """Check if it's time for periodic optimization."""
        return (time.time() - self.last_optimization_time) >= self.optimization_interval

    def optimize(self):
        """Execute periodic set optimization: pruning, flipping, merging, forgetting."""
        self.last_optimization_time = time.time()

        rules = list(self.hypergraph.rules.values())
        if not rules:
            return

        # 1. Pruning stale rules
        stale_rules = []
        for rule in rules:
            rule.update_freshness(self.decay_lambda)
            if rule.fresh < self.theta_f:
                stale_rules.append(rule.rule_id)

        for rid in stale_rules:
            self._remove_rule(rid)
        if stale_rules:
            print(f"RuleMaintenance: pruned {len(stale_rules)} stale rules")

        # 2. Conclusion flipping for low-confidence rules
        remaining = [r for r in self.hypergraph.rules.values()
                     if r.rule_id not in stale_rules]
        for rule in remaining:
            if (rule.conf < self.theta_c and
                    rule.trigger_count >= self.N_eval and
                    rule.used_count < rule.trigger_count * 0.3):
                # Flip conclusion
                old_conclusion = rule.conclusion
                rule.conclusion = 1 - rule.conclusion
                rule.conf = 0.3  # warm-start reset
                rule.conf0 = 0.3
                rule.sR_sum = 0.0
                rule.used_count = 0
                print(f"RuleMaintenance: flipped conclusion of {rule.rule_id} "
                      f"from {old_conclusion} to {rule.conclusion}")

        # 3. Rule merging
        self._merge_similar_rules()

        # 4. Capacity-based forgetting
        if len(self.hypergraph.rules) > self.max_rules:
            self._evict_low_utility_rules()

    def _remove_rule(self, rule_id: str):
        """Remove a rule from the hypergraph."""
        if rule_id in self.hypergraph.rules:
            del self.hypergraph.rules[rule_id]

        # Remove from hyperedges
        for he in self.hypergraph.hyperedges.values():
            he.rules = [r for r in he.rules if r.rule_id != rule_id]

    def _merge_similar_rules(self):
        """Merge similar rules using Jaccard similarity on atom sets."""
        rule_list = list(self.hypergraph.rules.values())
        merged = set()

        for i, r1 in enumerate(rule_list):
            if r1.rule_id in merged:
                continue
            cluster = [r1]
            for j, r2 in enumerate(rule_list):
                if i >= j or r2.rule_id in merged:
                    continue
                if r1.conclusion == r2.conclusion:
                    set1 = r1.atom_set()
                    set2 = r2.atom_set()
                    if set1 and set2:
                        union = len(set1 | set2)
                        if union > 0:
                            jaccard = len(set1 & set2) / union
                            if jaccard >= self.eta_m:
                                cluster.append(r2)

            if len(cluster) > 1:
                # Merge cluster into r1 (keep the one with higher confidence)
                best = max(cluster, key=lambda r: r.conf * r.trigger_count)
                for r in cluster:
                    if r.rule_id != best.rule_id:
                        merged.add(r.rule_id)
                        # Transfer stats
                        best.trigger_count += r.trigger_count
                        best.sR_sum += r.sR_sum
                        best.used_count += r.used_count
                        best.last_used_time = max(
                            best.last_used_time, r.last_used_time)
                        self._remove_rule(r.rule_id)

                # Recompute best rule's confidence as weighted average
                total_trig = sum(r.trigger_count for r in cluster)
                if total_trig > 0:
                    best.conf = sum(r.conf * r.trigger_count
                                    for r in cluster) / total_trig
                    best.conf0 = best.conf
                best.update_freshness(self.decay_lambda)
                print(f"RuleMaintenance: merged {len(cluster)} rules into {best.rule_id}")

        # Clean up merged rules
        for rid in list(merged):
            if rid in self.hypergraph.rules:
                del self.hypergraph.rules[rid]

    def _evict_low_utility_rules(self):
        """Evict low-utility rules when capacity is exceeded."""
        rules = list(self.hypergraph.rules.values())
        # Sort by utility: Conf(R) * Fresh(R)
        rules.sort(key=lambda r: r.conf * r.fresh)

        num_to_evict = int(len(rules) * self.evict_percentile)
        for rule in rules[:num_to_evict]:
            self._remove_rule(rule.rule_id)

        print(f"RuleMaintenance: evicted {num_to_evict} low-utility rules")


# ==============================================================================
# EvoRM Plugin - Main Interface
# ==============================================================================

class EvoRMPlugin:
    """
    Main EvoRM plugin class for integration with AdaCoAgentEA.

    Usage:
        plugin = EvoRMPlugin(client=openai_client)
        
        # Stage 1: Check if symbolic routing can decide
        decision, triggered, candidates = plugin.stage1(es_id, et_id, es_ctx, et_ctx)
        
        if decision is not None:
            # Direct routing - no LLM needed
            return decision
        else:
            # Stage 2: LLM with evidence and rule feedback
            result = plugin.stage2(es_id, et_id, es_ctx, et_ctx, 
                                   triggered, candidates, base_prompt)
            # Process result['decision'] and result['rule_feedback']
    """

    def __init__(self, client=None,
                 mlp_n_warmup: int = None,
                 theta_hi: float = 0.65,  # V11: balanced threshold for match routing
                 theta_prune: float = 0.75,  # V12: conservative non-match routing
                 alpha: float = 0.6,
                 beta: float = 0.4,
                 merge_threshold: float = 0.5,
                 persistence_dir: str = None,
                 ablation_mode: str = None,
                 enable_mlp_gate: bool = True,
                 eval_mode: bool = True,  # V14: eval mode (no rationale suffix, same prompt as baseline)
                 config: 'EvoRMConfig' = None):
        # Use config if provided, otherwise use individual params
        if config is not None:
            self.config = config
            alpha = config.alpha
            beta = config.beta
            merge_threshold = config.merge_threshold
            theta_hi = config.theta_hi
            theta_prune = config.theta_prune
            persistence_dir = config.persistence_dir or persistence_dir
            ablation_mode = config.ablation_mode or ablation_mode
        else:
            self.config = EvoRMConfig(
                alpha=alpha, beta=beta, merge_threshold=merge_threshold,
                theta_hi=theta_hi, theta_prune=theta_prune,
                persistence_dir=persistence_dir, ablation_mode=ablation_mode)
        
        self.client = client
        self.persistence_dir = persistence_dir
        self.ablation_mode = ablation_mode  # None, 'no_stage1', 'no_maintenance', etc.
        self._enable_mlp_gate = enable_mlp_gate

        # Initialize components
        self.rule_encoding = RuleEncoding(client=client)
        self.hypergraph = HypergraphStorage(
            alpha=alpha, beta=beta, merge_threshold=merge_threshold)
        self.controller = TwoStageInferenceController(
            self.hypergraph, client=client,
            theta_hi=theta_hi, theta_prune=theta_prune,
            route_nonmatch=self.config.route_nonmatch)
        self.maintenance = RuleMaintenance(
            self.hypergraph)

        # Entity Embedder — paper Section III-D (demb=1024, cosine similarity)
        self.entity_embedder = EntityEmbedder(
            demb=self.config.demb, n_features=self.config.n_features)
        self.hypergraph.entity_embedder = self.entity_embedder

        # Mlight / Mheavy — paper Section III-C (rationale + decision separation)
        self.mlight = MlightRationaleElicitor(
            client=client, model=self.config.model,
            temperature=self.config.mlight_temperature,
            max_retries=self.config.api_max_retries,
            timeout=self.config.api_timeout)
        self.mheavy = MheavyDecisionMaker(
            client=client, model=self.config.model,
            temperature=self.config.mheavy_temperature,
            max_retries=self.config.api_max_retries,
            timeout=self.config.api_timeout)
        self.controller.mlight = self.mlight
        self.controller.mheavy = self.mheavy

        # Ablation: disable Mlight (use legacy single-call)
        if self.ablation_mode == 'no_mlight':
            self.controller.use_mlight = False

        # Ablation: w/o Hypergraph (flat layout) — Table VI "w/o U"
        if self.ablation_mode == 'no_hypergraph':
            self.hypergraph.flat_mode = True

        # MLP Gate (gϕ) — paper Section III-E
        # Disabled by default: requires training data and causes incorrect routing
        # when trained on insufficient data. Enable explicitly with enable_mlp_gate=True.
        self.mlp_gate = None
        self._mlp_samples_collected = 0
        self.controller.mlp_gate = None
        self.controller.use_mlp_gate = False

        # Only enable MLP Gate when explicitly requested (via enable_mlp_gate=True or ablation_mode)
        if enable_mlp_gate or self.ablation_mode == 'enable_mlp_gate':
            self.mlp_gate = MLPGate(
                input_dim=self.config.mlp_input_dim,
                hidden_dims=self.config.mlp_hidden_dims,
                theta_gate=self.config.theta_gate,
                n_warmup=mlp_n_warmup if mlp_n_warmup is not None else self.config.mlp_n_warmup,
                device='cpu',  # CPU avoids CUDA threading conflicts with ThreadPoolExecutor
            )
            self.controller.mlp_gate = self.mlp_gate
            self.controller.use_mlp_gate = True

        # Statistics
        self.total_queries = 0
        self.llm_calls_saved = 0
        self.total_tokens = 0
        
        # V5: Training/eval mode — controls rationale suffix in stage2_simple
        self._eval_mode = eval_mode

        # Load persisted state if available
        if persistence_dir:
            self._load_state()

    def stage1(self, es_id: str, et_id: str,
               es_context: Dict, et_context: Dict
               ) -> Tuple[Optional[int], List[FOLRule], List[FOLRule]]:
        """
        Stage 1: Symbolic filtering.

        Returns:
            (decision, triggered_rules, candidate_rules)
            decision: None if needs Stage 2, 0/1 if direct routing
        """
        # Ablation: skip Stage 1
        if self.ablation_mode == 'no_stage1':
            return None, [], []
        self.total_queries += 1
        return self.controller.stage1_symbolic_filtering(
            es_id, et_id, es_context, et_context)

    def record_stage1_result(self, true_label: int, stage1_decision: int):
        """V5: Track Stage-1 routing accuracy by comparing with true label."""
        if stage1_decision is not None and true_label == stage1_decision:
            self.controller.stage1_correct += 1

    def set_eval_mode(self, eval_mode: bool = True):
        """V5: Set eval mode. In eval mode, stage2_simple uses baseline prompt
        without rationale suffix to avoid LLM behavior drift."""
        self._eval_mode = eval_mode

    def stage2_simple(self, es_id: str, et_id: str,
                     es_context: Dict, et_context: Dict,
                     triggered_rules: List[FOLRule],
                     base_prompt: str,
                     use_rationale_suffix: bool = True) -> Dict:
        """
        Simplified Stage 2: Use baseline LLM call for decision.
        This makes EvoRM a true plugin - Stage 1 routes when confident,
        otherwise falls back to the original baseline LLM call.
        No Mlight/Mheavy overhead.
        
        Appends a rationale elicitation suffix to the prompt so the LLM
        explains its reasoning. This enables parse_rationale to extract
        [DECISIVE] and [SUPPORTING] attribute-level conditions, creating
        specific rules that improve Stage-1 routing accuracy.
        
        Returns the raw response text so the calling code can parse it
        according to the task format (MATCH/NON-MATCH for ER, C1/C2/NIL for EL, etc.)
        """
        if self.client is None:
            return {'decision': None, 'rationale': '', 'token_usage': 0, 'raw_response': ''}
        
        # ---- Rationale Elicitation Suffix (V5: optional) ----
        # When use_rationale_suffix=True (warmup), append a structured reasoning
        # request so parse_rationale() can extract [DECISIVE]/[SUPPORTING] conditions.
        # When False (eval), use the baseline prompt as-is to avoid LLM behavior drift.
        if use_rationale_suffix:
            rationale_suffix = (
                "\n\nAfter giving your answer, briefly explain your reasoning "
                "using the format below. This helps verify your decision.\n"
                "[DECISIVE]\n"
                "Key attributes that determined your decision. Write each as:\n"
                "  attribute_name = same (actual_value)\n"
                "  attribute_name = differ (value1 vs value2)\n"
                "Examples: name = same (Apple Inc.), city = differ (Cupertino vs Seattle)\n"
                "[SUPPORTING]\n"
                "Additional evidence supporting your decision."
            )
            full_prompt = base_prompt + rationale_suffix
        else:
            full_prompt = base_prompt
        
        try:
            response = self.client.chat.completions.create(
                model="gpt-3.5-turbo-1106",
                messages=[{'role': 'user', 'content': full_prompt}],
                temperature=0.1,
            )
            response_text = response.choices[0].message.content.strip()
            tokens = response.usage.total_tokens
            
            # Track tokens
            try:
                from utils import tokens as tokens_cal
                tokens_cal.update_add_var(tokens)
            except Exception as _mlp_e:
                pass
            
            self.total_tokens += tokens
            
            # Parse decision from the FIRST LINE of the response.
            # The rationale sections ([DECISIVE], [SUPPORTING]) follow on
            # subsequent lines, so we extract the decision from line 1.
            first_line = response_text.split('\n')[0].strip().upper()
            full_upper = response_text.upper()
            
            if 'MATCH' in first_line and 'NON' not in first_line:
                decision = 1
            elif 'NON-MATCH' in first_line or 'NON_MATCH' in first_line:
                decision = 0
            elif first_line.startswith('C') and first_line[1:].isdigit():
                # Entity Linking format: C1, C2, etc.
                try:
                    decision = int(first_line[1:])
                except ValueError:
                    decision = 0
            elif 'NIL' in first_line:
                decision = 0
            elif 'MATCH' in full_upper and 'NON' not in full_upper:
                decision = 1
            elif 'NON-MATCH' in full_upper or 'NON_MATCH' in full_upper:
                decision = 0
            elif 'NIL' in full_upper:
                decision = 0
            else:
                # Try to find candidate ID in EL format
                import re as _re
                cm = _re.search(r'\bC(\d+)\b', full_upper)
                if cm:
                    decision = int(cm.group(1))
                else:
                    decision = 0  # Default to non-match
            
            # V12: Increment trigger_count for triggered rules
            # This is critical for rule confidence boost and routing.
            # Without this, rules never accumulate trigger_count and S1 rate stays at 0%.
            for rule in triggered_rules:
                rule.trigger_count += 1

            return {
                'decision': decision,
                'rationale': response_text,
                'raw_response': response_text,
                'token_usage': tokens,
                'method': 'baseline-fallback',
            }
        except Exception as e:
            return {
                'decision': None,
                'rationale': '',
                'raw_response': '',
                'token_usage': 0,
                'error': str(e),
                'method': 'error',
            }

    def stage2(self, es_id: str, et_id: str,
               es_context: Dict, et_context: Dict,
               triggered_rules: List[FOLRule],
               candidate_rules: List[FOLRule],
               base_prompt: str) -> Dict:
        """
        Stage 2: LLM judgment with evidence and rule feedback.

        Returns:
            Dict with 'decision', 'rule_feedback', 'rationale', 'token_usage'
        """
        result = self.controller.stage2_llm_judgment(
            es_id, et_id, es_context, et_context,
            triggered_rules, candidate_rules, base_prompt)

        self.total_tokens += result.get('token_usage', 0)

        # Process rule feedback
        if result.get('rule_feedback'):
            for rule_key, sR in result['rule_feedback'].items():
                # Try to find matching rule
                rule_id = self._find_rule_by_key(rule_key, triggered_rules)
                if rule_id:
                    self.maintenance.update_rule_confidence(rule_id, float(sR))

        return result

    def _find_rule_by_key(self, key: str, rules: List[FOLRule]) -> Optional[str]:
        """Find a rule ID by its index key (e.g., 'rule_1')."""
        match = re.match(r'rule_(\d+)', key)
        if match:
            idx = int(match.group(1)) - 1
            if 0 <= idx < len(rules):
                return rules[idx].rule_id
        return None

    def record_trajectory(self, es_id: str, et_id: str,
                          es_context: Dict, et_context: Dict,
                          decision: int, rationale: str,
                          triggered_rules: List[FOLRule],
                          rule_feedback: Dict[str, float]):
        """
        Record a matching trajectory after LLM inference.
        Creates/updates hyperedges and rules in the hypergraph.

        Args:
            es_id: source entity ID
            et_id: target entity ID
            es_context: source entity context
            et_context: target entity context
            decision: 1 (match) or 0 (non-match)
            rationale: LLM rationale
            triggered_rules: triggered rules
            rule_feedback: dict of rule_key -> contribution score
        """
        # Ablation: skip rule maintenance
        if self.ablation_mode == 'no_maintenance':
            return
        # Create a new rule from the rationale
        used_rules = []
        for rule_key, sR in rule_feedback.items():
            rule_id = self._find_rule_by_key(rule_key, triggered_rules)
            if rule_id:
                rule = self.hypergraph.get_rule(rule_id)
                if rule:
                    used_rules.append((rule, float(sR)))

        new_rule = self.rule_encoding.create_rule(
            rationale, decision, used_rules,
            es_context=es_context, et_context=et_context)

        # Build node set from entity contexts
        node_set = set()
        node_set.add(f"e_{es_id}")
        node_set.add(f"e_{et_id}")
        
        # Attribute-value-based indexing for cross-query rule matching
        for ctx in [es_context, et_context]:
            if ctx:
                for key in ['entity_name', 'name', 'title', 'description']:
                    val = ctx.get(key, '')
                    if val and isinstance(val, str):
                        node_set.add(val.lower()[:100])
                        for token in val.lower().split()[:5]:
                            if len(token) > 2:
                                node_set.add(token)
        
        for key, val in es_context.items():
            if not key.startswith('neighbors_'):
                node_set.add(f"a_{key}")
                if val:
                    node_set.add(f"v_{key}_{str(val)[:50]}")

        for key, val in et_context.items():
            if not key.startswith('neighbors_'):
                node_set.add(f"a_{key}")
                if val:
                    node_set.add(f"v_{key}_{str(val)[:50]}")

        # Compute and store entity embeddings
        if hasattr(self, 'entity_embedder') and self.entity_embedder is not None:
            try:
                emb_s = self.entity_embedder.embed_entity(es_context, f"e_{es_id}")
                emb_t = self.entity_embedder.embed_entity(et_context, f"e_{et_id}")
                self.hypergraph.store_entity_embedding(f"e_{es_id}", emb_s)
                self.hypergraph.store_entity_embedding(f"e_{et_id}", emb_t)
            except Exception as _mlp_e:
                pass

        # Create or update hyperedge
        self.hypergraph.get_or_create_hyperedge(
            (int(es_id) if es_id.isdigit() else hash(es_id),
             int(et_id) if et_id.isdigit() else hash(et_id)),
            node_set, new_rule)

        # Collect MLP gate training sample (thread-safe)
        if hasattr(self, 'mlp_gate') and self.mlp_gate is not None:
            if not self.mlp_gate.is_trained:
                # Collect samples from ALL Survival pairs, not just those with triggered rules
                try:
                    features = self.mlp_gate.extract_features(
                        es_context, et_context, triggered_rules if triggered_rules else [])
                    self.mlp_gate.collect_sample(features, float(1 if decision > 0 else 0))  # V10: binarize for EL compat
                    self._mlp_samples_collected += 1
                except Exception as e:
                    print(f'[MLPGate] Sample collection error: {e}')

        # Periodic maintenance
        if self.maintenance.should_optimize():
            self.maintenance.optimize()

    def build_context_dict(self, entity_name: str, relations: List[str] = None,
                           descriptions: str = "") -> Dict:
        """
        Build a context dictionary from entity information.

        Args:
            entity_name: entity name
            relations: list of relation strings (e.g., "Has relation 'X' with Y")
            descriptions: entity description text

        Returns:
            Context dict suitable for rule verification
        """
        ctx = {'entity_name': entity_name, 'description': descriptions}

        if relations:
            for rel_str in relations:
                # Parse "Has relation 'R' with E" or "Is R of E"
                m1 = re.match(r"Has relation '([^']+)' with (.+)", rel_str)
                if m1:
                    rel = m1.group(1).lower()
                    neighbor = m1.group(2).strip()
                    key = f"neighbors_{rel}"
                    if key not in ctx:
                        ctx[key] = set()
                    ctx[key].add(neighbor)

                m2 = re.match(r"Is ([^']+) of (.+)", rel_str)
                if m2:
                    rel = m2.group(1).lower()
                    neighbor = m2.group(2).strip()
                    key = f"neighbors_{rel}"
                    if key not in ctx:
                        ctx[key] = set()
                    ctx[key].add(neighbor)

        return ctx


    def train_mlp_if_ready(self):
        """Train MLP Gate if ready. Thread-safe: call from main thread only."""
        if hasattr(self, "mlp_gate") and self.mlp_gate is not None:
            if not self.mlp_gate.is_trained and self.mlp_gate.is_ready_to_train():
                try:
                    self.mlp_gate.fit_model(epochs=self.config.mlp_epochs,
                                            batch_size=self.config.mlp_batch_size,
                                            verbose=True)
                except Exception as e:
                    print(f"[MLPGate] Training error (main thread): {e}")
    def get_stats(self) -> Dict:
        """Get comprehensive statistics."""
        controller_stats = self.controller.get_stats()
        hypergraph_stats = self.hypergraph.stats()

        stats = {
            **controller_stats,
            **hypergraph_stats,
            'total_queries': self.total_queries,
            'llm_calls_saved': self.llm_calls_saved,
            'llm_save_rate': (self.llm_calls_saved /
                              max(1, self.llm_calls_saved + self.controller.stage2_total)),
            'total_tokens': self.total_tokens,
        }
        if hasattr(self, 'mlp_gate') and self.mlp_gate is not None:
            stats['mlp_gate'] = self.mlp_gate.get_stats()
        if hasattr(self, 'entity_embedder') and self.entity_embedder is not None:
            stats['entity_embedder'] = self.entity_embedder.get_stats()
        return stats

    def _save_state(self):
        """Persist hypergraph state to disk."""
        if not self.persistence_dir:
            return
        os.makedirs(self.persistence_dir, exist_ok=True)

        state = {
            'rules': {},
            'hyperedges': {},
            'inverted_index': {k: list(v) for k, v in self.hypergraph.inverted_index.items()},
            'stats': self.get_stats(),
        }

        for rid, rule in self.hypergraph.rules.items():
            state['rules'][rid] = {
                'rule_id': rule.rule_id,
                'atoms': [{'atom_type': a.atom_type, 'attr': a.attr,
                           'value1': a.value1, 'value2': a.value2}
                          for a in rule.atoms],
                'conclusion': rule.conclusion,
                'conf': rule.conf,
                'conf0': rule.conf0,
                'trigger_count': rule.trigger_count,
                'used_count': rule.used_count,
                'sR_sum': rule.sR_sum,
                'last_used_time': rule.last_used_time,
                'created_time': rule.created_time,
            }

        for hid, he in self.hypergraph.hyperedges.items():
            state['hyperedges'][hid] = {
                'hyperedge_id': he.hyperedge_id,
                'entity_pairs': list(he.entity_pairs),
                'node_set': list(he.node_set),
                'rule_ids': [r.rule_id for r in he.rules],
                'weight': he.weight,
            }

        state_path = os.path.join(self.persistence_dir, 'evorm_state.json')
        with open(state_path, 'w', encoding='utf-8') as f:
            json.dump(state, f, ensure_ascii=False, indent=2)

        # Save MLP gate
        if hasattr(self, 'mlp_gate') and self.mlp_gate is not None:
            try:
                mlp_path = os.path.join(self.persistence_dir, 'mlp_gate.pt')
                self.mlp_gate.save(mlp_path)
            except Exception as e:
                print(f"Failed to save MLP gate: {e}")

        print(f"EvoRM state saved to {state_path}")

    def _load_state(self):
        """Load persisted hypergraph state from disk."""
        if not self.persistence_dir:
            return
        state_path = os.path.join(self.persistence_dir, 'evorm_state.json')
        if not os.path.exists(state_path):
            return

        try:
            with open(state_path, 'r', encoding='utf-8') as f:
                state = json.load(f)

            # Restore rules
            for rid, rdata in state.get('rules', {}).items():
                atoms = [ConditionAtom(**a) for a in rdata['atoms']]
                rule = FOLRule(
                    rule_id=rdata['rule_id'],
                    atoms=atoms,
                    conclusion=rdata['conclusion'],
                    conf=rdata['conf'],
                    conf0=rdata['conf0'],
                    trigger_count=rdata['trigger_count'],
                    used_count=rdata['used_count'],
                    sR_sum=rdata['sR_sum'],
                    last_used_time=rdata['last_used_time'],
                    created_time=rdata['created_time'],
                )
                self.hypergraph.rules[rid] = rule

            # Restore hyperedges
            for hid, hdata in state.get('hyperedges', {}).items():
                he = Hyperedge(
                    hyperedge_id=hdata['hyperedge_id'],
                    entity_pairs=set(
                        tuple(p) for p in hdata['entity_pairs']),
                    node_set=set(hdata['node_set']),
                    weight=hdata['weight'],
                )
                for rid in hdata['rule_ids']:
                    if rid in self.hypergraph.rules:
                        he.rules.append(self.hypergraph.rules[rid])
                self.hypergraph.hyperedges[hid] = he

            # Restore inverted index
            for k, v in state.get('inverted_index', {}).items():
                self.hypergraph.inverted_index[k] = set(v)

            # Load MLP gate
            if hasattr(self, 'mlp_gate') and self.mlp_gate is not None:
                try:
                    mlp_path = os.path.join(self.persistence_dir, 'mlp_gate.pt')
                    self.mlp_gate.load(mlp_path)
                except Exception as e:
                    print(f"Failed to load MLP gate: {e}")

            print(f"EvoRM state loaded from {state_path}: "
                  f"{len(self.hypergraph.rules)} rules, "
                  f"{len(self.hypergraph.hyperedges)} hyperedges")

        except Exception as e:
            print(f"Failed to load EvoRM state: {e}")


# ==============================================================================
# Test / Demo
# ==============================================================================

if __name__ == "__main__":
    print("EvoRM Plugin - Self Test")
    print("=" * 60)

    # Create plugin without OpenAI client (offline test)
    plugin = EvoRMPlugin(client=None)

    # Test rule encoding
    print("\n1. Testing Rule Encoding...")
    rationale = """
    [DECISIVE] title=same ("locating data sources"), year=same (2003)
    [SUPPORTING] authors=same
    """
    atoms = plugin.rule_encoding.parse_rationale(rationale, conclusion=1)
    print(f"   Parsed atoms: {[a.to_key() for a in atoms]}")

    # Test rule creation
    rule = plugin.rule_encoding.create_rule(rationale, conclusion=1)
    print(f"   Created rule: {rule.rule_id}, conf={rule.conf:.3f}")

    # Test hypergraph storage
    print("\n2. Testing Hypergraph Storage...")
    node_set = {"e_1", "e_2", "a_title", "a_year", "a_authors"}
    he = plugin.hypergraph.get_or_create_hyperedge(
        (1, 100), node_set, rule)
    print(f"   Created hyperedge: {he.hyperedge_id}, weight={he.weight:.3f}")

    # Test candidate retrieval
    candidates = plugin.hypergraph.get_candidate_rules(["e_1"])
    print(f"   Candidates for e_1: {len(candidates)} rules")

    # Test stage 1
    print("\n3. Testing Stage 1 (Symbolic Filtering)...")
    es_ctx = {"entity_name": "Test E1", "title": "locating data sources",
              "year": "2003", "authors": "Smith et al"}
    et_ctx = {"entity_name": "Test E2", "title": "locating data sources",
              "year": "2003", "authors": "Smith et al"}
    decision, triggered, cands = plugin.stage1("e_1", "e_2", es_ctx, et_ctx)
    print(f"   Decision: {decision}, Triggered: {len(triggered)} rules")

    # Test trajectory recording
    print("\n4. Testing Trajectory Recording...")
    plugin.record_trajectory(
        "e_1", "e_2", es_ctx, et_ctx,
        decision=1, rationale=rationale,
        triggered_rules=triggered,
        rule_feedback={"rule_1": 0.8})

    # Test stats
    print("\n5. Stats:", json.dumps(plugin.get_stats(), indent=2))

    print("\n✅ All tests passed!")