scripts/merge_dedup_eval.py

"""Evaluate kink-name merge techniques on a hand-labeled ground-truth set.

Uses OTS tooling — no hand-rolled stemmer:
  - **spaCy** (`en_core_web_sm`) for proper lemmatization + POS-aware token filtering
  - **rapidfuzz** for normalized-edit-distance fuzzy matching (catches typos, doubled letters)
  - **sentence-transformers** (`all-MiniLM-L6-v2`) for semantic paraphrase scoring (lazy-loaded;
    skip via ``--no-embeddings`` if you only want the deterministic passes)

Each technique produces a verdict per pair (merge / no-merge); we report precision, recall,
F1 against the labeled set so you can see which combination of signals generalizes.

Run:
  .venv/bin/python scripts/merge_dedup_eval.py
  .venv/bin/python scripts/merge_dedup_eval.py --no-embeddings  # skip the SBERT pass
"""
from __future__ import annotations

import argparse
import json
import sqlite3
import sys
from collections import defaultdict
from pathlib import Path

REPO = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(REPO))

from backend.kink_merge import merge_signature as merge_signature_v1  # noqa: E402

GROUND_TRUTH_PATH = REPO / "tests" / "data" / "merge_dedup_ground_truth.json"


# ---------------------------------------------------------------------------
# OTS library wrappers — module-level so the eval shows clearly which signal
# each technique uses. None of these are hand-rolled.
# ---------------------------------------------------------------------------

_NLP = None
_SBERT = None


def _spacy() -> "spacy.language.Language":
    global _NLP
    if _NLP is None:
        import spacy
        # Disable parser/NER — we only need the tagger/lemmatizer.
        _NLP = spacy.load("en_core_web_sm", disable=["parser", "ner"])
    return _NLP


def _sbert():
    global _SBERT
    if _SBERT is None:
        from sentence_transformers import SentenceTransformer
        _SBERT = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")
    return _SBERT


# Spacy POS tags whose tokens carry kink-identity meaning. We drop DET / ADP / AUX / PRON / etc
# automatically — those are spaCy's stopword-equivalent classes.
_KEEP_POS = frozenset({"NOUN", "PROPN", "VERB", "ADJ", "ADV", "NUM", "X"})


def lemma_signature(name: str) -> str:
    """spaCy-lemmatized signature: {sorted unique lemmas of content tokens}."""
    base = merge_signature_v1(name)  # carry over FFM/MMF + threesome canonicalization
    if not base:
        return ""
    doc = _spacy()(base)
    lemmas = sorted({tok.lemma_.lower() for tok in doc if tok.pos_ in _KEEP_POS and tok.lemma_.strip()})
    if not lemmas:
        return ""
    if len(lemmas) == 1 and len(lemmas[0]) < 5:
        return ""  # avoid degenerate single-letter collisions
    return " ".join(lemmas)


def fuzzy_merge(a: str, b: str, *, threshold: int = 90) -> bool:
    """rapidfuzz token_set_ratio — catches typos / doubled letters / minor reorderings."""
    from rapidfuzz import fuzz
    return fuzz.token_set_ratio(a.lower(), b.lower()) >= threshold


def sbert_merge(a: str, b: str, *, threshold: float = 0.80) -> bool:
    """SBERT cosine similarity — catches paraphrases ('cuddling' ↔ 'snuggling')."""
    from sentence_transformers import util
    model = _sbert()
    emb = model.encode([a, b], normalize_embeddings=True)
    return float(util.cos_sim(emb[0], emb[1])) >= threshold


# ---------------------------------------------------------------------------
# Ground truth — pulled from real catalog names (verified against
# data/store_slim_b2.db) plus user-reported failures.
# ---------------------------------------------------------------------------

GROUND_TRUTH: list[dict] = [
    # POSITIVE: morphological variants of single-word kinks
    {"a": "Cuddling", "b": "Cuddles", "merge": True, "rationale": "real: fetlife_2275 + fetlife_590"},
    {"a": "Kissing", "b": "Kisses", "merge": True, "rationale": "plural vs gerund"},
    {"a": "Spanking", "b": "Spanks", "merge": True, "rationale": "gerund vs plural"},
    {"a": "Biting", "b": "Bites", "merge": True, "rationale": "gerund vs plural"},
    {"a": "Choking", "b": "Chokes", "merge": True, "rationale": "gerund vs verb"},
    {"a": "Tickling", "b": "Tickles", "merge": True, "rationale": "gerund vs plural"},
    {"a": "Licking", "b": "Licks", "merge": True, "rationale": "gerund vs plural"},

    # POSITIVE: word-order permutations (v1 already catches)
    {"a": "Cum Swallowing", "b": "Swallowing Cum", "merge": True, "rationale": "real: fetlife_19006 + fetlife_8045"},
    {"a": "Hair Pulling", "b": "Pulling Hair", "merge": True, "rationale": "tokens equal"},
    {"a": "Nipple Sucking", "b": "Sucking Nipples", "merge": True, "rationale": "morph + perm"},

    # POSITIVE: morphology + word-order (the hard case)
    {"a": "Cuddling after sex", "b": "After sex cuddles", "merge": True, "rationale": "user-reported"},
    {"a": "Cuddling after sex", "b": "Cuddly after sex", "merge": True, "rationale": "user-reported"},
    {"a": "After sex cuddles", "b": "Cuddly after sex", "merge": True, "rationale": "user-reported"},
    {"a": "Spanking with a paddle", "b": "Paddle spanking", "merge": True, "rationale": "morph + perm + stopword"},
    {"a": "Sucking on nipples", "b": "Nipple sucking", "merge": True, "rationale": "stopword + morph + perm"},

    # POSITIVE: morphology siblings of action verbs
    {"a": "Tease", "b": "Teasing", "merge": True, "rationale": "morph variants"},
    {"a": "Tease", "b": "Teases", "merge": True, "rationale": "noun vs verb, same act"},
    {"a": "MMF Threesome", "b": "FMM Threesomes", "merge": True, "rationale": "v1's gender-balance canonicalization"},

    # NEGATIVE: same root, different concept
    {"a": "Anal Sex", "b": "Anal Beads", "merge": False, "rationale": "act vs toy"},
    {"a": "Cuddling", "b": "Cuddly Mogwai", "merge": False, "rationale": "specific named scenario"},
    {"a": "Cum on Face", "b": "Cum on Tits", "merge": False, "rationale": "target body part is the kink"},
    {"a": "Choking", "b": "Chokers", "merge": False, "rationale": "act vs accessory"},
    {"a": "Whipping", "b": "Whips", "merge": False, "rationale": "act vs toy"},
    {"a": "Bondage", "b": "Bonding", "merge": False, "rationale": "different roots (bond vs bondage)"},
    {"a": "Tease", "b": "Tearing", "merge": False, "rationale": "tear ≠ tease"},
    {"a": "MMF Threesome", "b": "FFM Threesome", "merge": False, "rationale": "different gender balance"},
]


def save_ground_truth() -> None:
    GROUND_TRUTH_PATH.parent.mkdir(parents=True, exist_ok=True)
    GROUND_TRUTH_PATH.write_text(json.dumps(GROUND_TRUTH, indent=2), encoding="utf-8")


# ---------------------------------------------------------------------------
# Technique evaluation harness.
# Each technique is a name + a "would these merge" predicate.
# ---------------------------------------------------------------------------


def eval_technique(name: str, predicate) -> dict:
    tp = fp = fn = tn = 0
    misses: list[tuple[dict, bool]] = []
    for row in GROUND_TRUTH:
        merged = predicate(row["a"], row["b"])
        if row["merge"] and merged: tp += 1
        elif row["merge"] and not merged:
            fn += 1
            misses.append((row, merged))
        elif not row["merge"] and merged:
            fp += 1
            misses.append((row, merged))
        else: tn += 1
    n = len(GROUND_TRUTH)
    precision = tp / (tp + fp) if (tp + fp) else 0.0
    recall = tp / (tp + fn) if (tp + fn) else 0.0
    f1 = 2 * precision * recall / (precision + recall) if (precision + recall) else 0.0
    return {
        "name": name, "n": n,
        "tp": tp, "fp": fp, "fn": fn, "tn": tn,
        "accuracy": (tp + tn) / n,
        "precision": precision, "recall": recall, "f1": f1,
        "misses": misses,
    }


def print_report(res: dict) -> None:
    print(f"\n=== {res['name']} ===")
    print(f"  acc={res['accuracy']:.2%}  P={res['precision']:.2f}  R={res['recall']:.2f}  F1={res['f1']:.2f}  "
          f"tp={res['tp']} fp={res['fp']} fn={res['fn']} tn={res['tn']}")
    for row, merged in res["misses"][:8]:
        kind = "FN" if row["merge"] and not merged else "FP"
        print(f"   {kind}: {row['a']!r} ↔ {row['b']!r}  ({row['rationale']})")


def show_v2_finds_in_catalog(db_path: Path, *, limit: int = 12) -> None:
    if not db_path.is_file():
        print(f"\n(skip catalog scan — {db_path} not present)")
        return
    print(f"\n=== Catalog scan: {db_path} ===")
    conn = sqlite3.connect(f"file:{db_path}?mode=ro", uri=True)
    rows = conn.execute("SELECT id, name FROM kink WHERE name IS NOT NULL").fetchall()
    conn.close()
    print(f"  rows: {len(rows)}")

    v1_groups = defaultdict(list)
    v2_groups = defaultdict(list)
    for kid, name in rows:
        v1_groups[merge_signature_v1(name)].append((kid, name))
        v2_groups[lemma_signature(name)].append((kid, name))

    new_merges = []
    for sig, members in v2_groups.items():
        if len(members) < 2:
            continue
        v1_sigs = {merge_signature_v1(name) for _kid, name in members}
        if len(v1_sigs) == 1:
            continue
        new_merges.append((sig, members))

    print(f"  v1 buckets ≥2: {sum(1 for v in v1_groups.values() if len(v) >= 2)}")
    print(f"  spacy-lemma buckets ≥2: {sum(1 for v in v2_groups.values() if len(v) >= 2)}")
    print(f"  NEW merges (lemma-only): {len(new_merges)}")
    new_merges.sort(key=lambda x: -len(x[1]))
    print(f"\n  --- top {min(limit, len(new_merges))} clusters lemma catches that v1 misses ---")
    for sig, members in new_merges[:limit]:
        names = sorted({n for _kid, n in members})
        if len(names) < 2:
            continue
        print(f"    sig={sig!r}")
        for n in names[:5]:
            print(f"      • {n}")
        if len(names) > 5:
            print(f"      … +{len(names) - 5} more")


def main() -> int:
    p = argparse.ArgumentParser()
    p.add_argument("--no-embeddings", action="store_true",
                   help="Skip the slow SBERT pass; useful when iterating on the deterministic stages.")
    p.add_argument("--save-truth", action="store_true",
                   help="Persist GROUND_TRUTH to tests/data/merge_dedup_ground_truth.json")
    args = p.parse_args()

    if args.save_truth:
        save_ground_truth()
        print(f"wrote {GROUND_TRUTH_PATH}")

    print("Techniques (all OTS, no hand-rolled stemmers):")
    print("  v1: backend.kink_merge.merge_signature  — sorted unique tokens, FFM canonicalization")
    print("  lemma: spaCy en_core_web_sm  — POS-filtered + lemmatized signature equality")
    print("  fuzzy: rapidfuzz.fuzz.token_set_ratio ≥ 90")
    if not args.no_embeddings:
        print("  semantic: sentence-transformers/all-MiniLM-L6-v2  — cosine ≥ 0.80")

    techniques = [
        ("v1 token-set", lambda a, b: merge_signature_v1(a) == merge_signature_v1(b) and merge_signature_v1(a) != ""),
        ("spaCy lemma signature", lambda a, b: lemma_signature(a) == lemma_signature(b) and lemma_signature(a) != ""),
        ("rapidfuzz token_set_ratio≥90", lambda a, b: fuzzy_merge(a, b, threshold=90)),
        # Combined: lemma OR fuzzy (OR pass — if EITHER thinks merge, merge)
        ("lemma OR fuzzy(≥90)", lambda a, b: (
            (lemma_signature(a) == lemma_signature(b) and lemma_signature(a) != "")
            or fuzzy_merge(a, b, threshold=90)
        )),
    ]
    if not args.no_embeddings:
        techniques.append(
            ("lemma OR sbert(≥0.80)", lambda a, b: (
                (lemma_signature(a) == lemma_signature(b) and lemma_signature(a) != "")
                or sbert_merge(a, b, threshold=0.80)
            )),
        )

    results = [eval_technique(name, fn) for name, fn in techniques]
    for r in results:
        print_report(r)

    print("\n=== summary table ===")
    print(f"  {'technique':<35}  acc      P      R      F1")
    for r in results:
        print(f"  {r['name']:<35}  {r['accuracy']:.2%}  {r['precision']:.2f}  {r['recall']:.2f}  {r['f1']:.2f}")

    show_v2_finds_in_catalog(REPO / "data" / "store_slim_b2.db")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())