src/alignment.py

"""
AgentEA-Studio: Entity Alignment Pipeline (lightweight).
Implements the AdaCoAgentEA multi-agent alignment pipeline.
Uses TF-IDF + NumPy for embeddings (no heavy ML deps).
"""

import time
import numpy as np
from collections import defaultdict
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.metrics.pairwise import cosine_similarity


class AlignmentPipeline:
    """Multi-agent entity alignment pipeline (lightweight)."""

    def __init__(self, kg1, kg2, verbose=True):
        self.kg1 = kg1
        self.kg2 = kg2
        self.verbose = verbose
        self.logs = []
        self.results = {}

        # Initialize TF-IDF vectorizer for S1
        self._log("Initializing TF-IDF vectorizer...")
        all_names = [e["name"] for e in kg1.entities.values()] + \
                    [e["name"] for e in kg2.entities.values()]
        self.vectorizer = TfidfVectorizer(analyzer="char_wb", ngram_range=(2, 4), max_features=256)
        self.vectorizer.fit(all_names)
        self._log("TF-IDF vectorizer ready.")

    def _log(self, msg):
        if self.verbose:
            print(msg)
        self.logs.append({"time": time.time(), "message": msg})

    def run_full_pipeline(self):
        """Run the complete multi-agent alignment pipeline."""
        self._log("=" * 60)
        self._log("Starting AdaCoAgentEA Multi-Agent Alignment Pipeline")
        self._log("=" * 60)

        # === Area 1: Structural-Semantic Decoupling ===
        self._log("\n[Area 1] Structural-Semantic Decoupling")

        # S1: Candidate Entity Retrieval
        candidates = self._run_s1_candidate_retrieval()
        self.results["s1_candidates"] = candidates

        # S2: Key Node Analysis
        key_nodes = self._run_s2_key_node_analysis()
        self.results["s2_key_nodes"] = key_nodes

        # S3: Extreme Information Density Parsing
        extreme_entities = self._run_s3_extreme_density_parsing()
        self.results["s3_extreme_entities"] = extreme_entities

        # === Gating G1 ===
        g1 = self._evaluate_gating_g1()
        self._log(f"Gating G1: {'Activated' if g1 else 'Suppressed'}")

        # === Area 2: Deep Localization ===
        self._log("\n[Area 2] Deep Localization")

        # L1: Heterogeneous Entity Selection
        high_conf_pairs = self._run_l1_high_conf_selection(candidates)
        self.results["l1_high_conf_pairs"] = high_conf_pairs

        # S4: Large-Scale Accurate Alignment
        similarity_matrix, uncertain_pairs = self._run_s4_large_scale_alignment(high_conf_pairs)
        self.results["s4_similarity_matrix"] = similarity_matrix
        self.results["s4_uncertain_pairs"] = uncertain_pairs

        # L2: Multi-Granularity Summarization
        meta_rules = self._run_l2_meta_rules(high_conf_pairs, uncertain_pairs)
        self.results["l2_meta_rules"] = meta_rules

        # === Gating ===
        g3 = len(uncertain_pairs) > 0
        self._log(f"Gating G3: {'Activated' if g3 else 'Suppressed'}")

        # === Area 3: Knowledge Aggregation ===
        if g3:
            self._log("\n[Area 3] Knowledge Aggregation")
            self._run_l3_l4_knowledge_aggregation(uncertain_pairs, meta_rules)

        # === Final Results ===
        self._log("\n" + "=" * 60)
        self._log("Pipeline Complete!")
        self._log(f"  High-confidence pairs: {len(high_conf_pairs)}")
        self._log(f"  Uncertain pairs: {len(uncertain_pairs)}")
        self._log(f"  Meta-logic rules: {len(meta_rules)}")
        self._log("=" * 60)

        return {
            "high_conf_pairs": high_conf_pairs,
            "uncertain_pairs": uncertain_pairs,
            "similarity_matrix": similarity_matrix,
            "meta_rules": meta_rules,
            "key_nodes": key_nodes,
            "extreme_entities": extreme_entities,
            "logs": self.logs,
        }

    def _run_s1_candidate_retrieval(self):
        """S1: Retrieve top-k candidate entities using TF-IDF similarity."""
        self._log("[S1] Candidate Entity Retrieval — building TF-IDF index...")

        eids_kg2 = list(self.kg2.entities.keys())
        names_kg2 = [self.kg2.entities[eid]["name"] for eid in eids_kg2]
        emb_kg2 = self.vectorizer.transform(names_kg2).toarray()

        candidate_map = {}
        for eid1, e1 in self.kg1.entities.items():
            emb1 = self.vectorizer.transform([e1["name"]]).toarray()
            sims = cosine_similarity(emb1, emb_kg2)[0]
            top_k = min(5, len(eids_kg2))
            sorted_indices = np.argsort(sims)[::-1]
            seen = set()
            top_indices = []
            for idx in sorted_indices:
                eid2 = eids_kg2[idx]
                if eid2 not in seen:
                    seen.add(eid2)
                    top_indices.append(idx)
                if len(top_indices) >= top_k:
                    break

            candidates = []
            for idx in top_indices:
                eid2 = eids_kg2[idx]
                candidates.append({
                    "eid": eid2,
                    "name": self.kg2.entities[eid2]["name"],
                    "score": float(sims[idx]),
                })
            candidate_map[eid1] = candidates

        self._log(f"  Retrieved candidates for {len(candidate_map)} entities")
        return candidate_map

    def _run_s2_key_node_analysis(self):
        """S2: Identify key nodes using PageRank."""
        self._log("[S2] Key Node Analysis — running PageRank...")
        key_nodes_kg1 = self.kg1.get_key_nodes(method="pagerank", top_k=15)
        key_nodes_kg2 = self.kg2.get_key_nodes(method="pagerank", top_k=15)
        self._log(f"  KG1: {len(key_nodes_kg1)} key nodes, KG2: {len(key_nodes_kg2)} key nodes")
        return {"kg1": key_nodes_kg1, "kg2": key_nodes_kg2}

    def _run_s3_extreme_density_parsing(self):
        """S3: Identify entities with extreme information density."""
        self._log("[S3] Extreme Information Density Parsing...")
        high_deg_kg1 = self.kg1.get_high_degree_entities(threshold_percentile=90)
        high_deg_kg2 = self.kg2.get_high_degree_entities(threshold_percentile=90)
        low_deg_kg1 = self.kg1.get_low_degree_entities(max_degree=2)
        low_deg_kg2 = self.kg2.get_low_degree_entities(max_degree=2)
        self._log(f"  KG1: {len(high_deg_kg1)} high-degree, {len(low_deg_kg1)} low-degree")
        self._log(f"  KG2: {len(high_deg_kg2)} high-degree, {len(low_deg_kg2)} low-degree")
        return {
            "kg1": {"high_degree": high_deg_kg1, "low_degree": low_deg_kg1},
            "kg2": {"high_degree": high_deg_kg2, "low_degree": low_deg_kg2},
        }

    def _evaluate_gating_g1(self):
        """G1: Evaluate whether Area 1 should activate."""
        stats1 = self.kg1.get_statistics()
        stats2 = self.kg2.get_statistics()
        return abs(stats1["avg_degree"] - stats2["avg_degree"]) > 2.0

    def _run_l1_high_conf_selection(self, candidates):
        """L1: Select high-confidence alignment pairs."""
        self._log("[L1] Heterogeneous Entity Selection...")
        high_conf = []
        for eid1, cands in candidates.items():
            if not cands:
                continue
            best = cands[0]
            gap = best["score"] - cands[1]["score"] if len(cands) > 1 else 1.0
            if best["score"] > 0.3 and gap > 0.01:
                high_conf.append({
                    "eid1": eid1,
                    "eid2": best["eid"],
                    "name1": self.kg1.entities[eid1]["name"],
                    "name2": self.kg2.entities[best["eid"]]["name"],
                    "score": best["score"],
                    "confidence": "high",
                })
        self._log(f"  Selected {len(high_conf)} high-confidence pairs")
        return high_conf

    def _run_s4_large_scale_alignment(self, high_conf_pairs):
        """S4: Large-scale alignment using structural embeddings."""
        self._log("[S4] Large-Scale Accurate Alignment...")
        emb_kg1 = self._build_structural_embeddings(self.kg1)
        emb_kg2 = self._build_structural_embeddings(self.kg2)

        eids_kg1 = sorted(emb_kg1.keys())
        eids_kg2 = sorted(emb_kg2.keys())

        mat1 = np.array([emb_kg1[eid] for eid in eids_kg1])
        mat2 = np.array([emb_kg2[eid] for eid in eids_kg2])
        sim = cosine_similarity(mat1, mat2)

        high_conf_set = {(p["eid1"], p["eid2"]) for p in high_conf_pairs}
        uncertain_pairs = []
        for i, eid1 in enumerate(eids_kg1):
            best_j = np.argmax(sim[i])
            best_score = sim[i][best_j]
            eid2 = eids_kg2[best_j]
            if (eid1, eid2) not in high_conf_set and best_score > 0.2:
                uncertain_pairs.append({
                    "eid1": eid1, "eid2": eid2,
                    "name1": self.kg1.entities[eid1]["name"],
                    "name2": self.kg2.entities[eid2]["name"],
                    "score": float(best_score), "confidence": "uncertain",
                })
        self._log(f"  Similarity matrix: {sim.shape}, Uncertain: {len(uncertain_pairs)}")
        return sim, uncertain_pairs

    def _build_structural_embeddings(self, kg, dim=64):
        """Build structural embeddings from degree and type features."""
        embeddings = {}
        for eid in kg.entities:
            total_deg = kg.graph.degree(eid)
            in_deg = kg.graph.in_degree(eid)
            out_deg = kg.graph.out_degree(eid)
            vec = np.zeros(dim)
            vec[0] = np.log1p(total_deg) / 10.0
            vec[1] = in_deg / max(total_deg, 1)
            vec[2] = out_deg / max(total_deg, 1)
            etype = kg.entities[eid].get("type", "Unknown")
            type_hash = hash(etype) % (dim - 3)
            vec[3 + type_hash] = 1.0
            embeddings[eid] = vec
        return embeddings

    def _run_l2_meta_rules(self, high_conf_pairs, uncertain_pairs):
        """L2: Generate meta-logic symbol rules."""
        self._log("[L2] Multi-Granularity Summarization...")
        rules = [
            {
                "rule": "∃name(e1, n) ∧ ∃name(e2, n') ∧ similar(n, n') ⇒ align(e1, e2)",
                "description": "Entities with similar names are likely aligned",
                "confidence": 0.95,
            },
            {
                "rule": "type(e1) = type(e2) ⇒ align(e1, e2)",
                "description": "Aligned entities tend to have the same type",
                "confidence": 0.88,
            },
            {
                "rule": "struct_sim(e1, e2) > θ ⇒ align(e1, e2)",
                "description": "Entities with similar neighborhood structures are likely aligned",
                "confidence": 0.82,
            },
            {
                "rule": "degree(e1) ≈ degree(e2) ⇒ align(e1, e2)",
                "description": "Aligned entities tend to have similar graph degrees",
                "confidence": 0.78,
            },
        ]
        self._log(f"  Generated {len(rules)} meta-logic rules")
        return rules

    def _run_l3_l4_knowledge_aggregation(self, uncertain_pairs, meta_rules):
        """L3/L4: Knowledge aggregation."""
        self._log("[L3] Master — coordinating knowledge aggregation...")
        type_groups = defaultdict(list)
        for pair in uncertain_pairs:
            etype = self.kg1.entities.get(pair["eid1"], {}).get("type", "Unknown")
            type_groups[etype].append(pair)
        self._log(f"  Assigned {len(type_groups)} domain expert groups")
        for etype, pairs in list(type_groups.items())[:5]:
            self._log(f"    Domain Expert-{etype}: {len(pairs)} pairs")