src/knowledge_extraction/cluster/cluster.py

"""Mention[] → TermCluster[].

Matching order, cheapest and most certain first:
    1. exact match after normalisation
    2. abbreviation ↔ expansion (from legend blocks)
    3. conservative fuzzy (token_set_ratio >= 92, and only above 5 characters)

**Over-merging is much worse than under-merging.** An under-merge costs one
extra LLM call and one extra review-queue row. A wrong merge silently destroys a
distinct term, and no downstream stage recovers it — the expert never learns the
term existed. Every threshold here is set on that asymmetry.

Clustering is also what makes conflict detection possible at all: contradictory
definitions can only be compared if all evidence for a term reaches the same
call.
"""

from __future__ import annotations

from rapidfuzz import fuzz

from ..models import AbbrevPair, ClusterResult, Mention, TermCluster
from ..settings import FUZZY_MIN_LEN, FUZZY_THRESHOLD
from .normalize import AbbrevIndex, is_noise


def cluster_mentions(
    mentions: list[Mention],
    abbrev_pairs: list[AbbrevPair],
    doc_id: str,
    fuzzy_threshold: int = FUZZY_THRESHOLD,
) -> ClusterResult:
    index = AbbrevIndex(abbrev_pairs)
    kept = [m for m in mentions if not is_noise(m.surface)]

    buckets: dict[str, dict] = {}

    for mention in kept:
        key = index.canonical_key(mention.surface)

        if key in buckets:
            _add(buckets[key], mention, "exact")
            continue

        linked = next((k for k in buckets if index.linked(mention.surface, k)), None)
        if linked:
            _add(buckets[linked], mention, "abbrev")
            continue

        match = _fuzzy_match(key, buckets.keys(), fuzzy_threshold)
        if match:
            _add(buckets[match], mention, "fuzzy")
            continue

        buckets[key] = {
            "surfaces": {mention.surface},
            "mentions": [mention],
            "reasons": set(),
        }

    clusters: list[TermCluster] = []
    ordered = sorted(buckets.items(), key=lambda kv: -len(kv[1]["mentions"]))
    for i, (key, data) in enumerate(ordered):
        variants = sorted(data["surfaces"], key=lambda s: (len(s), s))
        clusters.append(
            TermCluster(
                cluster_id=f"{doc_id}#c{i:03d}",
                canonical=_canonical(variants, key),
                variants=variants,
                mentions=data["mentions"],
                mention_count=len(data["mentions"]),
                merge_reasons=sorted(data["reasons"]),
            )
        )

    n_mentions, n_clusters = len(kept), len(clusters)
    return ClusterResult(
        doc_id=doc_id,
        clusters=clusters,
        n_mentions=n_mentions,
        n_clusters=n_clusters,
        compression_ratio=round(n_mentions / n_clusters, 3) if n_clusters else 0.0,
    )


def _add(bucket: dict, mention: Mention, reason: str) -> None:
    bucket["surfaces"].add(mention.surface)
    bucket["mentions"].append(mention)
    bucket["reasons"].add(reason)


def _fuzzy_match(key: str, existing, threshold: int) -> str | None:
    best, best_score = None, 0.0
    for other in existing:
        # Short strings fuzzy-match far too easily: "PA" vs "UA" scores high on
        # token_set_ratio. Below FUZZY_MIN_LEN only exact matching is allowed.
        if min(len(key), len(other)) < FUZZY_MIN_LEN:
            continue
        score = fuzz.token_set_ratio(key, other)
        if score >= threshold and score > best_score:
            best, best_score = other, score
    return best


def _canonical(variants: list[str], key: str) -> str:
    """Prefer the shortest non-trivial surface — usually the term as a reader
    would look it up ("PA", not "Physical Availability (PA) untuk ...")."""
    for v in variants:
        if len(v) >= 2:
            return v
    return key