app/graph/alignment.py

"""Entity + relation alignment for the per-user knowledge graph.

Without alignment, doc A emits "GWU" and doc B emits "GWU " (trailing space)
or "gwu" as separate MERGE keys, so each doc contributes its own island of
nodes. This module normalizes names and fuzz-matches near-duplicates before
the write hits Neo4j, so the KG converges into one connected graph as docs
accumulate.

Three layers, in order of increasing cost / recall (see #43a in TASKS):
  L1  case + unicode + punctuation normalization (deterministic)
  L2  fuzzy string match against same-type existing entities (rapidfuzz)
  L4  relation predicate canonicalization (small mapping table)

L3 (semantic alignment via bge embeddings for abbreviation-vs-expansion pairs)
is planned separately as #43b.
"""

from __future__ import annotations

import re
import unicodedata
from collections.abc import Iterable
from dataclasses import dataclass

from rapidfuzz import fuzz

from app.graph.schema import GraphEntity, GraphRelationship

# ---------------------------------------------------------------------------
# L1 — name normalization
# ---------------------------------------------------------------------------

_WHITESPACE_RE = re.compile(r"\s+")
_TRAILING_PUNCT_RE = re.compile(r"[.,;:!?'\"()\[\]]+$")
_LEADING_ARTICLES: tuple[str, ...] = ("the ", "a ", "an ")


def normalize_name(name: str) -> str:
    """Canonical form used as the Neo4j MERGE key for entities.

    Order matters: NFKC first (folds fullwidth/ligatures), then lowercase,
    then strip leading articles and trailing punctuation, then collapse
    internal whitespace. Empty input returns empty string.
    """
    if not name:
        return ""
    s = unicodedata.normalize("NFKC", name).strip().lower()
    for art in _LEADING_ARTICLES:
        if s.startswith(art):
            s = s[len(art) :]
            break
    s = _TRAILING_PUNCT_RE.sub("", s)
    s = _WHITESPACE_RE.sub(" ", s).strip()
    return s


# ---------------------------------------------------------------------------
# L2 — fuzzy alias detection
# ---------------------------------------------------------------------------

# 0-100 similarity. 90 is safe for named entities: catches typo/spacing
# variants like "Kamalasankaris" vs "Kamalasankari" (1 char off, ~96),
# without collapsing genuinely-different short names like "Java" vs
# "JavaScript" (~67, both Technology). Adjust per per-user tuning later.
FUZZY_THRESHOLD = 90


@dataclass(frozen=True)
class Candidate:
    """A same-type existing entity that a new entity may alias to. Only
    carries the fields alignment needs — display name lives in Neo4j."""

    name_lower: str
    type: str


def find_alias(
    normalized: str,
    entity_type: str,
    candidates: Iterable[Candidate],
) -> str | None:
    """Return the best same-type candidate whose fuzzy similarity clears
    FUZZY_THRESHOLD, or None. Exact matches short-circuit.

    Type-scoping prevents false positives: "Java" and "JavaScript" are both
    Technology at ~67% similarity (would be dropped), but if one were
    Person and the other Technology they'd never be compared. Token-sort
    ratio is order-invariant so "John Smith" vs "Smith, John" aligns.
    """
    if not normalized:
        return None
    best_score = 0.0
    best_match: str | None = None
    for cand in candidates:
        if cand.type != entity_type:
            continue
        if cand.name_lower == normalized:
            return cand.name_lower
        score = fuzz.token_sort_ratio(normalized, cand.name_lower)
        if score > best_score:
            best_score = score
            best_match = cand.name_lower
    if best_score >= FUZZY_THRESHOLD:
        return best_match
    return None


# ---------------------------------------------------------------------------
# L4 — relation predicate normalization
# ---------------------------------------------------------------------------

# Maps common raw predicates to a canonical form. Only forward-direction
# variants are grouped; inverse forms ("authored by" as the inverse of
# "authored") would need source/target swapping which is out of scope for
# L4 — those still create their own edges. Additions welcome as the corpus
# surfaces more variants.
RELATION_ALIASES: dict[str, str] = {
    # authorship (creator → work)
    "author of": "authored",
    "authors": "authored",
    "wrote": "authored",
    "writes": "authored",
    "co-authored": "co-authored",
    "coauthored": "co-authored",
    "co-author of": "co-authored",
    # affiliation (person → org)
    "works at": "affiliated with",
    "works for": "affiliated with",
    "employed by": "affiliated with",
    "employee of": "affiliated with",
    "member of": "affiliated with",
    # location (entity → location)
    "based in": "located in",
    "headquartered in": "located in",
    "situated in": "located in",
    # usage (user → tool)
    "using": "uses",
    "utilizes": "uses",
    "utilises": "uses",
    "leverages": "uses",
    # dependency
    "requires": "depends on",
    "relies on": "depends on",
    "built on": "based on",
    # comparison (source outperforms target)
    "better than": "outperforms",
    "exceeds": "outperforms",
    "surpasses": "outperforms",
    # composition
    "consists of": "contains",
    "comprised of": "contains",
    "includes": "contains",
}


def normalize_relation(relation: str) -> str:
    """Map a raw relation phrase to its canonical form. Unknown phrases are
    lowercased+trimmed and passed through unchanged — better a slightly
    noisy predicate than a lossy remap."""
    if not relation:
        return ""
    key = _WHITESPACE_RE.sub(" ", relation.strip().lower())
    key = _TRAILING_PUNCT_RE.sub("", key)
    return RELATION_ALIASES.get(key, key)


# ---------------------------------------------------------------------------
# Batch aligner — the entrypoint called by the worker
# ---------------------------------------------------------------------------


@dataclass(frozen=True)
class AlignedEntity:
    """An entity resolved to canonical form and ready for upsert."""

    name: str  # display form (raw name from extraction; Neo4j preserves first-write)
    name_lower: str  # MERGE key (normalized or aliased-to)
    type: str
    description: str


@dataclass(frozen=True)
class AlignedRelationship:
    """A relationship with source/target rewritten to canonical name_lower
    keys and predicate mapped through RELATION_ALIASES."""

    source_lower: str
    target_lower: str
    relation: str


def align_batch(
    entities: list[GraphEntity],
    relationships: list[GraphRelationship],
    existing_candidates: list[Candidate],
) -> tuple[list[AlignedEntity], list[AlignedRelationship]]:
    """Align one doc's extraction against the existing user KG.

    Steps:
    1. Normalize each incoming entity's name.
    2. Fuzzy-match against the existing candidates AND against entities
       already aligned earlier in this same batch — so a doc that mentions
       both "GWU" and "George Washington University" (rare but possible)
       collapses to one node when they cross the threshold.
    3. Build a lookup: raw extraction name → canonical name_lower.
    4. Rewrite relationships through that lookup; normalize their
       predicates; drop self-loops and endpoint-missing edges.
    """
    aligned_entities: list[AlignedEntity] = []
    # Rolling map from the raw extraction string (lowercased) to the
    # canonical name_lower it aligned to. Populated as we walk entities.
    raw_to_canonical: dict[str, str] = {}
    # Grows as we align each entity so later ones in this batch can alias
    # to earlier ones.
    live_candidates: list[Candidate] = list(existing_candidates)
    seen_canonical: set[str] = set()

    for e in entities:
        raw = (e.name or "").strip()
        if not raw:
            continue
        normalized = normalize_name(e.name)
        if not normalized:
            continue
        raw_lower = raw.lower()

        alias = find_alias(normalized, e.type, live_candidates)
        canonical_lower = alias or normalized
        raw_to_canonical[raw_lower] = canonical_lower

        if canonical_lower in seen_canonical:
            # Already emitted (either same normalized form or a fuzzy
            # match to an already-aligned entity in this batch). Don't
            # emit a second AlignedEntity for it — Neo4j MERGE would
            # dedupe anyway, but this keeps the log counts honest.
            continue
        seen_canonical.add(canonical_lower)
        aligned_entities.append(
            AlignedEntity(
                name=raw,
                name_lower=canonical_lower,
                type=e.type,
                description=e.description,
            )
        )
        live_candidates.append(Candidate(name_lower=canonical_lower, type=e.type))

    aligned_rels: list[AlignedRelationship] = []
    seen_rels: set[tuple[str, str, str]] = set()
    for r in relationships:
        src_raw = (r.source or "").strip().lower()
        tgt_raw = (r.target or "").strip().lower()
        src_canonical = raw_to_canonical.get(src_raw)
        tgt_canonical = raw_to_canonical.get(tgt_raw)
        if not src_canonical or not tgt_canonical:
            continue
        if src_canonical == tgt_canonical:
            continue
        rel_canonical = normalize_relation(r.relation)
        if not rel_canonical:
            continue
        key3 = (src_canonical, tgt_canonical, rel_canonical)
        if key3 in seen_rels:
            continue
        seen_rels.add(key3)
        aligned_rels.append(
            AlignedRelationship(
                source_lower=src_canonical,
                target_lower=tgt_canonical,
                relation=rel_canonical,
            )
        )
    return aligned_entities, aligned_rels