src/knowledge_extraction/cluster/cluster.py

"""Mention[] → TermCluster[].

Matching order, cheapest and most certain first:
    1. exact match after normalisation (plural-folded — `normalize(lemma=True)`)
    2. abbreviation ↔ expansion (from legend blocks)

**There is deliberately no fuzzy pass** (retired 2026-09-10; `_fuzzy_match` is
parked below with the measurement that retired it). Every merge it made that was
correct is now done deterministically by plural folding; 94% of what it made was
a token-subset match, which is always wrong.

**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 src.middlewares.logging import get_logger

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

logger = get_logger("knowledge_cluster")


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)]
    # Pairs the document declares inline - "Waste Removal (WR)" - are as good as
    # a legend entry and this corpus uses them constantly. Learned in one pass
    # BEFORE bucketing, so the bracketed form, the bare expansion and the bare
    # abbreviation all reach the same key instead of splitting three ways.
    inline = index.learn_inline(m.surface for m in kept)
    if inline:
        logger.info("cluster_inline_pairs_learned", extra={"learned": inline})

    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

        # --- FUZZY PASS RETIRED 2026-09-10 (measured, not a style call) ---------
        # `_fuzzy_match` is kept below for the record; this call site is not.
        #
        # `token_set_ratio` returns EXACTLY 100 whenever one string's token set is
        # a subset of the other's, so a head noun always absorbed its own
        # compounds and the threshold of 92 was inert. Measured across the BUMA
        # standard, a Komatsu shop manual and an Open Pit textbook chapter:
        #
        #   131 of the 140 merges this pass made were token-subset matches — ALL
        #   wrong (`Radiator` ate `radiator cap`; `level` ate `oil level sensor`;
        #   `Working Hours` ate `effective working hours`, which cost the gold
        #   term EWH its own entry entirely).
        #
        #   Of the 9 genuine near-misses, 6 were plain singular/plural and 3 were
        #   `gain/loss pty|qty|ua grouping` — three DISTINCT metrics, also wrong.
        #
        # So the pass produced no correct merge that plural folding does not now
        # do deterministically (`normalize(lemma=True)`). Retiring it drove
        # subset-merges to 0 on all three documents and lifted post-cluster gold
        # recall 0.8293 -> 0.8537, while also giving the FEWEST clusters of any
        # correct configuration (Open Pit 431, against 434 for the best guarded
        # fuzzy variant and 445 for no-fuzzy-without-lemma).
        #
        # Matching is now: exact (normalised, plural-folded) -> abbreviation
        # expansion. Nothing else. If you are tempted to restore this, the thing
        # to beat is subset_pairs = 0; see FREE_STAGE_RESULTS.md.
        #
        # 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