"""
Compositional matching engine.

Takes extraction output (concepts with operational signatures) and matches
against the formalism knowledge base using three strategies:

1. DIRECT MATCH: signature matches a known formalism exactly → "≡ X"
2. COMPOSITIONAL MATCH: signature = compose(f₁, f₂, ...) → "≡ X ∘ Y"
3. ANALOGY MATCH: same meso/macro type as known formalism → "≈ X (Δ: ...)"

Composition is type-checked: rules specify input/output signatures, and the
engine verifies that each rule's input constraints are satisfied before
applying it. The result is a valid composition tree, not just a bag of rules.

Returns UNKNOWN when no match found. Returns CONFUSED when the concept's
own operations are internally contradictory.
"""

from __future__ import annotations

import json
import re
from collections.abc import Sequence
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Optional

import yaml

# ---------------------------------------------------------------------------
# Type aliases (match the YAML schemas)
# ---------------------------------------------------------------------------

Operation = str       # maximize | minimize | transform | project | decompose | sample | aggregate | match | propagate
DomainType = str       # vector | matrix | graph | distribution | sequence | manifold | scalar_field | set | function | latent
CodomainType = str     # vector | matrix | graph | distribution | sequence | scalar | embedding | assignment | latent
ObjectiveFamily = str  # divergence | likelihood | energy | correlation | information | none | adversarial | reconstruction
MesoType = str         # joint_embedding | spectral_method | energy_model | dynamical_system | ...
MacroType = str        # optimization | eigenvalue_problem | stochastic_process | statistical_inference | hamiltonian_system | none


@dataclass
class Signature:
    """Typed operational signature extracted from a concept or formalism."""
    operation: Operation | None = None
    domain: DomainType | None = None
    codomain: CodomainType | None = None
    objective_family: ObjectiveFamily | None = None

    @classmethod
    def from_concept(cls, concept: dict) -> "Signature":
        """Infer a typed signature from the LLM's extracted concept fields."""
        sig = cls()
        # Infer operation from the mathematical_operation string
        op_text = (concept.get("mathematical_operation") or "").lower()
        sig.operation = _infer_operation(op_text)
        sig.domain = _normalize_domain(concept.get("domain"))
        sig.codomain = _normalize_domain(concept.get("codomain"))
        sig.objective_family = _infer_objective(concept.get("objective") or "", op_text)
        return sig

    @classmethod
    def from_formalism(cls, fm: dict) -> "Signature":
        """Extract signature from a formalism YAML entry."""
        sig_raw = fm.get("signature", {})
        return cls(
            operation=sig_raw.get("operation"),
            domain=sig_raw.get("domain"),
            codomain=sig_raw.get("codomain"),
            objective_family=sig_raw.get("objective_family"),
        )

    def matches(self, other: "Signature", *, strict: bool = True) -> float:
        """Return a match score 0.0–1.0 between this and another signature.

        strict=True: all non-None fields must match exactly; score = proportion matched.
        strict=False: fuzzy — operation + codomain weighted higher.
        """
        fields = [
            ("operation", 0.35),
            ("domain", 0.15),
            ("codomain", 0.30),
            ("objective_family", 0.20),
        ]
        total = 0.0
        matched = 0.0
        pairs = []
        for attr, weight in fields:
            mine = getattr(self, attr)
            theirs = getattr(other, attr)
            pairs.append((attr, mine, theirs, weight))

        if strict:
            # Only score on fields where both sides have a value
            scorable = [(a, m, t, w) for a, m, t, w in pairs if m is not None and t is not None]
            if not scorable:
                return 0.0
            total = sum(w for _, _, _, w in scorable)
            matched = sum(w for _, m, t, w in scorable if m == t)
        else:
            total = sum(weight for _, _, _, weight in pairs)
            for _, mine_val, theirs_val, weight in pairs:
                if mine_val is None or theirs_val is None:
                    matched += weight * 0.5  # neutral on unknown
                elif mine_val == theirs_val:
                    matched += weight
                # else: 0 on mismatch

        return matched / total if total > 0 else 0.0


@dataclass
class Formalism:
    """A known formalism from the KB."""
    id: str
    name: str
    year: int | None
    origin: str
    signature: Signature
    meso_type: MesoType | None
    macro_type: MacroType | None
    canonical_reference: str
    status: str

    @classmethod
    def from_yaml(cls, entry: dict) -> "Formalism":
        return cls(
            id=entry["id"],
            name=entry["name"],
            year=entry.get("year"),
            origin=entry.get("origin", ""),
            signature=Signature.from_formalism(entry),
            meso_type=entry.get("meso_type"),
            macro_type=entry.get("macro_type"),
            canonical_reference=entry.get("canonical_reference", ""),
            status=entry.get("status", "seed"),
        )


@dataclass
class CompositionRule:
    """A composition rule from the KB."""
    id: str
    name: str
    description: str
    decomposes_to: list[str]          # formalism_ids this rule expands to
    input_constraints: dict[str, Any]
    output_signature: dict[str, Any]
    preserves: list[str]
    introduces: list[str]
    examples: list[str]
    status: str

    @classmethod
    def from_yaml(cls, entry: dict) -> "CompositionRule":
        return cls(
            id=entry["id"],
            name=entry.get("name", entry["id"]),
            description=entry.get("description", ""),
            decomposes_to=entry.get("decomposes_to", []),
            input_constraints=entry.get("input_constraints", {}),
            output_signature=entry.get("output_signature", {}),
            preserves=entry.get("preserves", []),
            introduces=entry.get("introduces", []),
            examples=entry.get("examples", []),
            status=entry.get("status", "seed"),
        )

    def accepts(self, formalism: Formalism) -> bool:
        """Check whether this rule can be applied to the given formalism."""
        constraints = self.input_constraints
        if not constraints:
            return True  # universal rule

        # Check specific formalism IDs
        req_ids = constraints.get("formalism_ids")
        if req_ids is not None:
            if isinstance(req_ids, list) and req_ids and formalism.id not in req_ids:
                return False

        # Check meso types
        req_meso = constraints.get("meso_types")
        if req_meso is not None:
            if isinstance(req_meso, list) and req_meso:
                if formalism.meso_type not in req_meso:
                    return False

        # Check macro types
        req_macro = constraints.get("macro_types")
        if req_macro is not None:
            if isinstance(req_macro, list) and req_macro:
                if formalism.macro_type not in req_macro:
                    return False

        return True


@dataclass
class CompositionNode:
    """A node in a composition tree: formalism + list of applied rules."""
    formalism: Formalism
    rules: list[CompositionRule] = field(default_factory=list)

    @property
    def name(self) -> str:
        if not self.rules:
            return self.formalism.name
        rule_names = " ∘ ".join(r.name for r in self.rules)
        return f"{self.formalism.name} ∘ {rule_names}"

    @property
    def is_direct(self) -> bool:
        return len(self.rules) == 0


@dataclass
class MatchResult:
    """The result of matching a concept against the KB."""
    concept_name: str
    result_type: str  # "identity" | "compositional" | "analogy" | "unknown" | "confused"
    reduction: str     # e.g., "CCA ∘ neuralize ∘ predict_in_codomain"
    reduction_expanded: str  # e.g., "CCA ∘ Gradient Descent ∘ CCA" (rules expanded)
    canonical_analog: str  # e.g., "Kernel CCA (Bach & Jordan, 2002)"
    genuine_delta: str  # what's actually new, if anything
    micro: str          # fine-grained: what operation happens at the lowest level
    meso: str           # mid-level: what structural family this belongs to
    macro: str          # top-level: what grand tradition this sits in
    confidence: float
    nodes: list[CompositionNode] = field(default_factory=list)
    match_scores: list[float] = field(default_factory=list)
    notes: list[str] = field(default_factory=list)

    @property
    def display(self) -> str:
        """Sous rature display: ~~AI Term~~ → Mathematical Operation"""
        if self.result_type == "identity":
            return f"~~{self.concept_name}~~ ≡ {self.reduction}"
        elif self.result_type == "compositional":
            base = f"~~{self.concept_name}~~ ≡ {self.reduction}"
            if self.reduction_expanded and self.reduction_expanded != self.reduction:
                base += f"\x00EXPAND\x00{self.reduction_expanded}\x00/EXPAND\x00"
            return base
        elif self.result_type == "analogy":
            return f"~~{self.concept_name}~~ ≈ {self.reduction} (Δ: {self.genuine_delta})"
        elif self.result_type == "confused":
            return f"~~{self.concept_name}~~ → CONFUSED: {self.notes[0] if self.notes else 'terminology overload'}"
        else:
            return f"~~{self.concept_name}~~ → UNKNOWN"


# ---------------------------------------------------------------------------
# KB loading
# ---------------------------------------------------------------------------

def _kb_dir() -> Path:
    return Path(__file__).resolve().parent / "kb"


def load_formalisms(path: Path | None = None) -> list[Formalism]:
    """Load the formalism KB."""
    if path is None:
        path = _kb_dir() / "formalisms.yaml"
    with open(path) as f:
        data = yaml.safe_load(f)
    return [Formalism.from_yaml(e) for e in data.get("formalisms", [])]


def load_composition_rules(path: Path | None = None) -> list[CompositionRule]:
    """Load the composition rules KB."""
    if path is None:
        path = _kb_dir() / "composition_rules.yaml"
    with open(path) as f:
        data = yaml.safe_load(f)
    return [CompositionRule.from_yaml(e) for e in data.get("composition_rules", [])]


# ---------------------------------------------------------------------------
# Signature inference helpers (parse LLM output into typed fields)
# ---------------------------------------------------------------------------

_OP_PATTERNS: list[tuple[str, str]] = [
    (r"\b(minimi[zs]e|minimi[zs]ation|minimi[zs]ing)\b", "minimize"),
    (r"\b(maximi[zs]e|maximi[zs]ation|maximi[zs]ing)\b", "maximize"),
    (r"\b(project|projection|projecting)\b", "project"),
    (r"\b(decompose|decomposition|factorize|factorization|eigen)\b", "decompose"),
    (r"\b(sample|sampling|generate|generating|generative)\b", "sample"),
    (r"\b(aggregate|aggregation|weighted\s+sum|pooling)\b", "aggregate"),
    (r"\b(transform|transformations?|map|mapping)\b", "transform"),
    (r"\b(match|matching|align|alignment)\b", "match"),
    (r"\b(propagat|diffuse|random\s+walk)\b", "propagate"),
]


def _infer_operation(text: str) -> Operation | None:
    text_lower = text.lower()
    for pattern, op in _OP_PATTERNS:
        if re.search(pattern, text_lower):
            return op
    return None


_DOMAIN_MAP: dict[str, DomainType] = {
    "vector": "vector", "vectors": "vector", "embedding": "vector",
    "matrix": "matrix", "matrices": "matrix",
    "graph": "graph",
    "distribution": "distribution", "probability": "distribution",
    "sequence": "sequence", "token": "sequence", "time series": "sequence",
    "manifold": "manifold",
    "set": "set",
    "function": "function", "scalar field": "scalar_field",
    "latent": "latent", "latent space": "latent",
}


def _normalize_domain(text: str | None) -> DomainType | None:
    if not text:
        return None
    t = text.strip().lower()
    # Try exact match first
    for key, val in _DOMAIN_MAP.items():
        if key in t:
            return val
    return t  # pass through — might be a valid value we just don't have mapped


_OBJ_PATTERNS: list[tuple[str, ObjectiveFamily]] = [
    (r"\b(kl\b|kullback|divergence|kl\s*divergence)\b", "divergence"),
    (r"\b(likelihood|log\s*likelihood|mle|maximum\s*likelihood)\b", "likelihood"),
    (r"\b(energy|free\s*energy|hamiltonian)\b", "energy"),
    (r"\b(correlation|canonical\s*correlation|cca|cross.correlation)\b", "correlation"),
    (r"\b(mutual\s*information|mi\b|infonce|information\s*max)\b", "information"),
    (r"\b(adversarial|minimax|min.max|gan\b|discriminator)\b", "adversarial"),
    (r"\b(reconstruction|autoencod|encode.decode|mse\b|squared\s*error)\b", "reconstruction"),
]


def _infer_objective(obj_text: str, op_text: str) -> ObjectiveFamily | None:
    combined = (obj_text + " " + op_text).lower()
    for pattern, obj in _OBJ_PATTERNS:
        if re.search(pattern, combined):
            return obj
    return None


def _infer_meso_type(sig: Signature, concept: dict) -> MesoType | None:
    """Infer meso-type from signature and concept text."""
    text = (
        f"{concept.get('mathematical_operation', '')} "
        f"{concept.get('canonical_analog', '')}"
    ).lower()

    if any(w in text for w in ("kernel", "rkhs", "nyström", "nystrom")):
        return "kernel_method"
    if any(w in text for w in ("spectral", "eigen", "laplacian", "fourier")):
        return "spectral_method"
    if any(w in text for w in ("energy", "free energy", "boltzmann", "hamiltonian")):
        return "energy_model"
    if any(w in text for w in ("diffusion", "sde", "langevin", "score-based", "ddpm")):
        return "diffusion_process"
    if any(w in text for w in ("variational", "elbo", "vi ")):
        return "variational"
    if any(w in text for w in ("optimal transport", "wasserstein", "sinkhorn")):
        return "optimal_transport"
    if any(w in text for w in ("contrastive", "siamese", "infonce")):
        return "joint_embedding"
    if any(w in text for w in ("gan", "adversarial", "minimax", "generator")):
        return "game_theoretic"
    if any(w in text for w in ("mean field", "mean-field")):
        return "mean_field"
    if any(w in text for w in ("joint embedding", "multi.view", "multiview", "cca")):
        return "joint_embedding"
    if any(w in text for w in ("pca", "projection", "linear", "svd")):
        return "linear_projection"
    if any(w in text for w in ("spin", "ising", "hopfield")):
        return "spin_system"
    return None


# ---------------------------------------------------------------------------
# Matching engine
# ---------------------------------------------------------------------------

# Keyword → rule triggers for canonical analog path.
# When the LLM says "this is essentially X," but the concept text
# describes specific modifications, these keyword sets determine which
# rules describe the paper's actual delta from the canonical analog.
_RULE_KEYWORDS: dict[str, list[str]] = {
    "neuralize": [
        "learned", "learnable", "deep", "encoder", "neural", "network",
        "parameterized", "differentiable", "end-to-end", "trained", "φ_θ",
        "f_θ", "g_θ", "dnn", "backprop",
    ],
    "predict_in_codomain": [
        "predict", "predictive", "predicting", "prediction",
        "latent space", "embedding space", "representation space",
        "in latent", "in embedding", "codomain",
        "future embedding", "future representation",
    ],
    "contrastivize": [
        "contrastive", "contrastively", "positive pair", "negative pair",
        "infonce", "noise contrastive", "nce",
    ],
    "diffuse": [
        "diffusion", "denoising", "denoise", "score-based",
        "reverse process", "forward process", "sde", "ddpm",
    ],
    "adversarize": [
        "adversarial", "gan", "discriminator", "generator",
        "minimax", "min-max",
    ],
    "variational_bound": [
        "variational", "elbo", "vae", "auto-encoding", "autoencoding",
        "amortized inference", "inference network",
    ],
    "regularize": [
        "regularize", "regularization", "l1 ", "l2 ", "weight decay",
        "dropout", "sparsity",
    ],
    "attention_wrap": [
        "attention", "self-attention", "transformer", "attend",
    ],
}


@dataclass
class MatchEngine:
    """The compositional matching engine."""

    formalisms: list[Formalism]
    rules: list[CompositionRule]
    config: dict = field(default_factory=dict)

    # Indexes for fast lookup
    _by_id: dict[str, Formalism] = field(default_factory=dict)
    _by_meso: dict[MesoType, list[Formalism]] = field(default_factory=dict)
    _by_macro: dict[MacroType, list[Formalism]] = field(default_factory=dict)
    _by_name: dict[str, Formalism] = field(default_factory=dict)  # fuzzy name index
    _name_tokens: dict[str, list[Formalism]] = field(default_factory=dict)

    def __post_init__(self):
        self._build_indexes()

    def _build_indexes(self):
        for fm in self.formalisms:
            self._by_id[fm.id] = fm
            if fm.meso_type:
                self._by_meso.setdefault(fm.meso_type, []).append(fm)
            if fm.macro_type:
                self._by_macro.setdefault(fm.macro_type, []).append(fm)
            # Name index: lowercase the name and each token
            name_lower = fm.name.lower()
            self._by_name[name_lower] = fm
            for token in name_lower.replace("(", "").replace(")", "").replace("/", " ").split():
                token = token.strip().rstrip(".,;:")
                if len(token) >= 3:
                    self._name_tokens.setdefault(token, []).append(fm)

    # ---- Canonical analog resolution ----

    def _resolve_canonical_analog(self, analog_text: str) -> Formalism | None:
        """Parse the LLM's canonical_analog field and find the matching formalism.

        Handles formats like:
        - "Kernel CCA (Bach & Jordan, 2002)"
        - "Kernel Canonical Correlation Analysis"
        - "CCA — Bach & Jordan 2002"
        """
        if not analog_text:
            return None
        text_lower = analog_text.lower().strip()

        # 1. Exact name match
        if text_lower in self._by_name:
            return self._by_name[text_lower]

        # 2. Try stripping parenthetical citations
        no_parens = re.sub(r"\([^)]*\)", "", text_lower).strip()
        if no_parens in self._by_name:
            return self._by_name[no_parens]

        # 3. Token intersection scoring
        tokens = set(t.strip().rstrip(".,;:") for t in no_parens.replace("/", " ").split() if len(t.strip()) >= 3)
        if not tokens:
            return None

        scored: list[tuple[int, Formalism]] = []
        for fm in self.formalisms:
            fm_tokens = set(t.strip().rstrip(".,;:") for t in fm.name.lower().replace("(", "").replace(")", "").replace("/", " ").split() if len(t.strip()) >= 3)
            intersection = tokens & fm_tokens
            if intersection:
                scored.append((len(intersection), fm))

        if scored:
            scored.sort(key=lambda x: x[0], reverse=True)
            if scored[0][0] >= 2:
                return scored[0][1]
            # Single-token match only if the matched token is distinctive
            top_token = max(tokens, key=len) if tokens else ""
            for score, fm in scored:
                if score >= 1 and len(top_token) >= 4:  # e.g., "canonical", "correlation"
                    return fm

        return None

    # ---- Main entry point ----

    def match_concept(self, concept: dict) -> MatchResult:
        """Match a single extracted concept against the KB.

        Returns a MatchResult with the best available decomposition.
        """
        name = concept.get("name", "unknown")
        sig = Signature.from_concept(concept)
        meso = _infer_meso_type(sig, concept)

        # Route based on concept flags from extraction
        flags = concept.get("flags", []) or []
        if "cannot_determine_from_abstract" in flags:
            return MatchResult(
                concept_name=name,
                result_type="unknown",
                reduction="cannot determine from abstract",
                reduction_expanded="cannot determine from abstract",
                canonical_analog="",
                genuine_delta="",
                micro=concept.get("mathematical_operation", ""),
                meso=meso or "unknown",
                macro="unknown",
                confidence=0.0,
                notes=["LLM extraction flagged: cannot determine from abstract"],
            )
        if "terminology_overload" in flags or "claim_operation_mismatch" in flags:
            pass  # Still attempt match but note flags

        # Strategy 0: Canonical analog from LLM extraction (highest-weight signal)
        canonical_analog_text = concept.get("canonical_analog", "") or ""
        base_fm = self._resolve_canonical_analog(canonical_analog_text)

        if base_fm is not None:
            # The LLM says this is essentially X. Now find rules that account
            # for what makes it "novel" beyond X.
            rules = self._find_rules_to_match(base_fm, sig, concept)
            if not rules:
                # No rules needed — the LLM-identified analog is the answer
                return MatchResult(
                    concept_name=name,
                    result_type="identity",
                    reduction=base_fm.name,
                    reduction_expanded=base_fm.name,
                    canonical_analog=f"{base_fm.name} ({base_fm.canonical_reference})",
                    genuine_delta="LLM-identified rebranding of known formalism",
                    micro=f"{base_fm.signature.operation}({base_fm.signature.domain} → {base_fm.signature.codomain})",
                    meso=base_fm.meso_type or "none",
                    macro=base_fm.macro_type or "none",
                    confidence=0.85,  # LLM identification is high-confidence
                    nodes=[CompositionNode(formalism=base_fm)],
                    match_scores=[0.85],
                    notes=["matched via LLM canonical_analog field"],
                )
            else:
                # Rules account for the delta from the canonical analog
                rule_names = " ∘ ".join(r.name for r in rules)
                expanded_fms = self.expand_rules(rules)
                expanded = base_fm.name + " ∘ " + " ∘ ".join(expanded_fms)
                return MatchResult(
                    concept_name=name,
                    result_type="compositional",
                    reduction=f"{base_fm.name} ∘ {rule_names}",
                    reduction_expanded=expanded,
                    canonical_analog=f"{base_fm.name} ({base_fm.canonical_reference})",
                    genuine_delta=" ∘ ".join(r.name for r in rules),
                    micro=f"{base_fm.signature.operation}({base_fm.signature.domain} → {base_fm.signature.codomain})",
                    meso=base_fm.meso_type or "none",
                    macro=base_fm.macro_type or "none",
                    confidence=0.80,
                    nodes=[CompositionNode(formalism=base_fm, rules=rules)],
                    match_scores=[0.80],
                    notes=["matched via LLM canonical_analog with rule delta"],
                )

        # Strategy 1: Direct identity match (signature only)
        direct = self._match_direct(sig)
        if direct and direct[1] >= 0.85:
            fm, score = direct
            return MatchResult(
                concept_name=name,
                result_type="identity",
                reduction=fm.name,
                reduction_expanded=fm.name,
                canonical_analog=f"{fm.name} ({fm.canonical_reference})",
                genuine_delta="none — this is a direct rebranding",
                micro=f"exactly {fm.name}: {fm.signature.operation}({fm.signature.domain} → {fm.signature.codomain})",
                meso=fm.meso_type or "none",
                macro=fm.macro_type or "none",
                confidence=score,
                nodes=[CompositionNode(formalism=fm)],
                match_scores=[score],
            )

        # Strategy 2: Compositional match
        comp = self._match_compositional(sig, meso)
        if comp and comp.confidence >= 0.5:
            return comp

        # Strategy 3: Analogy match
        analogy = self._match_analogy(sig, meso)
        if analogy and analogy.confidence >= 0.4:
            return analogy

        # Give up
        return MatchResult(
            concept_name=name,
            result_type="unknown",
            reduction="no match in KB",
            reduction_expanded="no match in KB",
            canonical_analog="",
            genuine_delta="",
            micro=concept.get("mathematical_operation", ""),
            meso=meso or "unknown",
            macro="unknown",
            confidence=0.0,
            notes=["concept signature does not match any formalism or valid composition"],
        )

    # ---- Strategy 1: Direct match ----

    def _match_direct(self, sig: Signature) -> tuple[Formalism, float] | None:
        best: tuple[Formalism, float] | None = None
        best_score = 0.0
        for fm in self.formalisms:
            score = sig.matches(fm.signature, strict=True)
            if score > best_score:
                best_score = score
                best = (fm, score)
        if best and best_score >= 0.5:
            return best
        return None

    # ---- Rule expansion ----

    def expand_rules(self, rules: list[CompositionRule]) -> list[str]:
        """Recursively expand composition rules into their constituent formalism names.

        Each rule's decomposes_to field lists formalism IDs it's composed of.
        This method looks up those formalisms by ID and returns their display names,
        giving the full decomposition chain beneath what appears as a single rule.
        """
        names: list[str] = []
        for rule in rules:
            expanded = False
            for fm_id in rule.decomposes_to:
                fm = self._by_id.get(fm_id)
                if fm:
                    names.append(fm.name)
                    expanded = True
            if not expanded:
                # Rule has no decomposition — use the rule name itself
                names.append(rule.name)
        return names

    # ---- Rule-to-target matching (for canonical analog case) ----

    # _RULE_KEYWORDS is a module-level constant; see below.
    # Keyword → rule triggers for canonical analog path.

    def _find_rules_to_match(
        self,
        base_fm: Formalism,
        target_sig: Signature,
        concept: dict | None = None,
    ) -> list[CompositionRule]:
        """Find composition rules that describe the paper's delta from the
        canonical analog. Uses two strategies:

        1. Signature-distance minimization (greedy, depth ≤ 2)
        2. Keyword-triggered rules from concept text (when LLM has already
           identified the base formalism — the keywords describe what the
           paper actually changed)
        """
        base_sig = base_fm.signature
        base_dist = 1.0 - base_sig.matches(target_sig, strict=False)

        # Strategy A: Signature-distance minimization
        best_rules: list[CompositionRule] = []
        best_dist = base_dist

        for rule in self.rules:
            if not rule.accepts(base_fm):
                continue
            transformed = self._apply_rule_signature(base_sig, rule)
            dist = 1.0 - transformed.matches(target_sig, strict=False)
            if dist < best_dist:
                best_dist = dist
                best_rules = [rule]

        for rule1 in self.rules:
            if not rule1.accepts(base_fm):
                continue
            inter = self._apply_rule_signature(base_sig, rule1)
            for rule2 in self.rules:
                if rule2 is rule1:
                    continue
                if not self._rule_accepts_signature(rule2, inter):
                    continue
                transformed = self._apply_rule_signature(inter, rule2)
                dist = 1.0 - transformed.matches(target_sig, strict=False)
                if dist < best_dist:
                    best_dist = dist
                    best_rules = [rule1, rule2]

        sig_improved = best_dist < base_dist - 0.05

        # Strategy B: Keyword-triggered rules from concept text
        if concept is not None:
            text = " ".join([
                concept.get("mathematical_operation") or "",
                concept.get("objective") or "",
                concept.get("claimed_novelty_text") or "",
                concept.get("confidence_rationale") or "",
            ]).lower()

            keyword_rules: list[CompositionRule] = []
            for rule in self.rules:
                if rule in best_rules:
                    continue
                keywords = _RULE_KEYWORDS.get(rule.id, [])
                if any(kw in text for kw in keywords):
                    if rule.accepts(base_fm) or not rule.input_constraints:
                        keyword_rules.append(rule)

            # Merge: signature-driven rules first, then keyword rules
            # that don't duplicate. Prefer keyword rules when the LLM
            # has identified the base (they're semantically richer).
            if keyword_rules:
                if sig_improved:
                    # Both strategies agree — merge, deduplicate
                    merged = list(best_rules)
                    for kr in keyword_rules:
                        if kr not in merged:
                            merged.append(kr)
                    return merged
                else:
                    # Only keyword strategy fires — use those
                    return keyword_rules

        if sig_improved:
            return best_rules
        return []

    # ---- Strategy 2: Compositional match ----

    def _match_compositional(self, sig: Signature, meso: MesoType | None) -> MatchResult | None:
        """Try to decompose the concept as formalism + composition rules.

        For each formalism whose signature is close, see if applying
        available rules transforms it toward the concept's signature.
        """
        candidates: list[tuple[Formalism, list[CompositionRule], float]] = []

        # For each formalism that could be a base
        for fm in self.formalisms:
            for rule in self.rules:
                if not rule.accepts(fm):
                    continue
                # Apply rule conceptually and score
                composed_sig = self._apply_rule_signature(fm.signature, rule)
                score = sig.matches(composed_sig, strict=True)
                if score >= 0.5:
                    candidates.append((fm, [rule], score))

                # Try two-rule compositions
                for rule2 in self.rules:
                    if rule2 is rule:
                        continue
                    # Check if rule2 accepts the output type of rule1
                    intermediate = self._apply_rule_signature(fm.signature, rule)
                    if not self._rule_accepts_signature(rule2, intermediate):
                        continue
                    composed2 = self._apply_rule_signature(intermediate, rule2)
                    score2 = sig.matches(composed2, strict=True)
                    if score2 >= 0.5:
                        candidates.append((fm, [rule, rule2], score2))

        if not candidates:
            return None

        # Pick best
        best_fm, best_rules, best_score = max(candidates, key=lambda c: c[2])
        rule_names = " ∘ ".join(r.name for r in best_rules)
        reduction = f"{best_fm.name} ∘ {rule_names}"
        expanded_fms = self.expand_rules(best_rules)
        expanded = best_fm.name + " ∘ " + " ∘ ".join(expanded_fms)

        return MatchResult(
            concept_name="",
            result_type="compositional",
            reduction=reduction,
            reduction_expanded=expanded,
            canonical_analog=f"{best_fm.name} ({best_fm.canonical_reference})",
            genuine_delta=" ∘ ".join(r.name for r in best_rules),
            micro=f"{best_fm.signature.operation}({best_fm.signature.domain} → {best_fm.signature.codomain})",
            meso=best_fm.meso_type or "none",
            macro=best_fm.macro_type or "none",
            confidence=best_score,
            nodes=[CompositionNode(formalism=best_fm, rules=best_rules)],
            match_scores=[best_score],
        )

    def _apply_rule_signature(self, sig: Signature, rule: CompositionRule) -> Signature:
        """Compute the approximate output signature after applying a rule.

        Rules modify operation/codomain/objective_family/meso/macro.
        Fields not mentioned in output_signature pass through unchanged.
        """
        out = rule.output_signature
        return Signature(
            operation=out.get("operation", sig.operation),
            domain=sig.domain,  # domain typically preserved
            codomain=out.get("codomain", sig.codomain),
            objective_family=out.get("objective_family", sig.objective_family),
        )

    def _rule_accepts_signature(self, rule: CompositionRule, sig: Signature) -> bool:
        """Check if a rule can be applied to an intermediate signature.

        This is a looser check than rule.accepts(formalism) since we
        don't have a Formalism object — we check meso/macro constraints.
        """
        constraints = rule.input_constraints
        if not constraints:
            return True

        req_ids = constraints.get("formalism_ids")
        if req_ids is not None and isinstance(req_ids, list) and req_ids:
            return False  # specific formalism constraint can't be satisfied by signature alone

        # If rule requires specific meso_types and we can't determine them, be permissive
        req_meso = constraints.get("meso_types")
        if req_meso is not None and isinstance(req_meso, list) and req_meso:
            # Without a formalism we can't enforce meso_type constraints tightly
            # For intermediate nodes, be permissive
            pass

        req_macro = constraints.get("macro_types")
        if req_macro is not None and isinstance(req_macro, list) and req_macro:
            pass  # same reasoning

        return True

    # ---- Strategy 3: Analogy match ----

    def _match_analogy(self, sig: Signature, meso: MesoType | None) -> MatchResult | None:
        """Find formalisms with the same meso/macro type but different specifics."""
        if not meso:
            return None

        candidates = self._by_meso.get(meso, [])
        if not candidates:
            return None

        # Find the best signature match among same-meso formalisms
        best_score = 0.0
        best_fm: Formalism | None = None
        for fm in candidates:
            score = sig.matches(fm.signature, strict=False)
            if score > best_score:
                best_score = score
                best_fm = fm

        if best_fm is None or best_score < 0.3:
            return None

        # Compute the delta: what's different?
        deltas: list[str] = []
        if sig.operation and sig.operation != best_fm.signature.operation:
            deltas.append(f"operation: {best_fm.signature.operation} → {sig.operation}")
        if sig.codomain and sig.codomain != best_fm.signature.codomain:
            deltas.append(f"codomain: {best_fm.signature.codomain} → {sig.codomain}")
        if sig.objective_family and sig.objective_family != best_fm.signature.objective_family:
            deltas.append(f"objective: {best_fm.signature.objective_family} → {sig.objective_family}")
        delta_str = "; ".join(deltas) if deltas else "minor variation"

        return MatchResult(
            concept_name="",
            result_type="analogy",
            reduction=f"{best_fm.name}",
            reduction_expanded=f"{best_fm.name}",
            canonical_analog=f"{best_fm.name} ({best_fm.canonical_reference})",
            genuine_delta=delta_str,
            micro=f"Shares meso-type '{meso}' with {best_fm.name}",
            meso=meso,
            macro=best_fm.macro_type or "none",
            confidence=best_score,
            nodes=[CompositionNode(formalism=best_fm)],
            match_scores=[best_score],
        )

    # ---- Batch matching ----

    def match_paper(self, extraction: dict) -> dict:
        """Match all concepts extracted from a paper.

        Returns the extraction dict augmented with match results.
        """
        concepts = extraction.get("concepts", [])
        matched = []
        for concept in concepts:
            result = self.match_concept(concept)
            matched.append(result)

        extraction["_matches"] = [self._result_to_dict(r) for r in matched]
        extraction["_match_summary"] = self._summarize(matched)
        return extraction

    def _result_to_dict(self, r: MatchResult) -> dict:
        return {
            "concept_name": r.concept_name,
            "result_type": r.result_type,
            "reduction": r.reduction,
            "reduction_expanded": r.reduction_expanded,
            "canonical_analog": r.canonical_analog,
            "genuine_delta": r.genuine_delta,
            "micro": r.micro,
            "meso": r.meso,
            "macro": r.macro,
            "confidence": r.confidence,
            "display": r.display,
            "notes": r.notes,
        }

    def _summarize(self, results: list[MatchResult]) -> dict:
        identity = sum(1 for r in results if r.result_type == "identity")
        compositional = sum(1 for r in results if r.result_type == "compositional")
        analogy = sum(1 for r in results if r.result_type == "analogy")
        unknown = sum(1 for r in results if r.result_type == "unknown")
        confused = sum(1 for r in results if r.result_type == "confused")
        total = len(results)

        return {
            "total_concepts": total,
            "identity_reductions": identity,
            "compositional_reductions": compositional,
            "analogy_matches": analogy,
            "unknown": unknown,
            "confused": confused,
            "reduction_rate": (identity + compositional) / total if total else 0,
        }


# ---------------------------------------------------------------------------
# Convenience: load engine from defaults
# ---------------------------------------------------------------------------

def load_engine(
    formalism_path: Path | None = None,
    rules_path: Path | None = None,
) -> MatchEngine:
    """Load the match engine from the default KB files."""
    formalisms = load_formalisms(formalism_path)
    rules = load_composition_rules(rules_path)
    return MatchEngine(formalisms=formalisms, rules=rules)


# ---------------------------------------------------------------------------
# Self-test: JEPA canonical decomposition
# ---------------------------------------------------------------------------

def _test_jepa():
    """Verify the canonical JEPA decomposition: CCA ∘ neuralize ∘ predict_in_codomain."""
    engine = load_engine()

    # Simulate an extraction result for JEPA
    jepa_concept = {
        "name": "Joint Embedding Predictive Architecture",
        "is_claimed_novel": True,
        "claimed_novelty_text": "predicts representations in latent space rather than raw inputs",
        "mathematical_operation": "maximize mutual information between joint embeddings of x and y, then predict future embedding from past embedding in the joint space",
        "domain": "vector",
        "codomain": "vector",
        "objective": "maximize I(Z_x; Z_y) — mutual information between embeddings, plus prediction error in latent space",
        "constraints": [],
        "canonical_analog": "Kernel CCA (Bach & Jordan 2002)",
        "deconstructive_move": "binary_overturn",
        "confidence": "high",
        "confidence_rationale": "abstract explicitly describes joint embedding + prediction in latent space",
        "flags": [],
    }

    result = engine.match_concept(jepa_concept)
    print("=== JEPA Canonical Decomposition Test ===")
    print(f"Concept: {result.concept_name}")
    print(f"Result type: {result.result_type}")
    print(f"Reduction: {result.reduction}")
    print(f"Display: {result.display}")
    print(f"Micro: {result.micro}")
    print(f"Meso: {result.meso}")
    print(f"Macro: {result.macro}")
    print(f"Confidence: {result.confidence}")
    print(f"Canonical analog: {result.canonical_analog}")
    print(f"Genuine delta: {result.genuine_delta}")
    if result.notes:
        print(f"Notes: {result.notes}")
    print()
    return result


if __name__ == "__main__":
    _test_jepa()