scripts/export_normalization_clusters.py

#!/usr/bin/env python3
"""Export kink title clusters for offline review and normalization benchmarking.

Emits JSONL lines (one JSON object per line) for either:

- ``fingerprint``: same rules as ``scripts/find_name_duplicate_clusters.py`` /
  ``frontend/name-fingerprint.js`` (keep in sync when tuning).
- ``merge_signature``: same rules as ``backend.kink_merge.merge_signature`` /
  ``propose_merge_clusters`` (FMF / threesome token variants).

Usage:
  python scripts/export_normalization_clusters.py --db data/store_slim.db --out /tmp/clusters.jsonl --mode fingerprint
  python scripts/export_normalization_clusters.py --db data/store_slim.db --out /tmp/merge_sig.jsonl --mode merge_signature --min-size 2

For **pair-level** labeling (merge vs alias vs reject), prefer
``scripts/source_normalization_benchmark_pairs.py`` (tiers documented in ``backend/benchmark_pair_sources.py``).
"""
from __future__ import annotations

import argparse
import json
import random
import re
import sqlite3
from collections import defaultdict
from pathlib import Path

from backend.kink_merge import merge_graph_eligible, merge_signature


def fingerprint_play_title(raw: str) -> str:
    s = (raw or "").strip().lower()
    if not s:
        return ""
    s = re.sub(r"[\u2018\u2019\u201c\u201d\u0060\u00b4]", "'", s)
    s = re.sub(r"\s+", " ", s)
    s = re.sub(r"\bgood\s+morning\s+", "goodmorning ", s)
    s = re.sub(r"\bdeep\s*throating\b", "deepthroat", s)
    s = re.sub(r"\bdeepthroating\b", "deepthroat", s)
    s = re.sub(r"\bdeep\s*throat\b", "deepthroat", s)
    s = re.sub(r"\bblow\s+jobs\b", "blowjob", s, flags=re.I)
    s = re.sub(r"\bblow\s+job\b", "blowjob", s, flags=re.I)
    s = re.sub(r"\bblowjobs\b", "blowjob", s)
    s = re.sub(r"\bstrap\s+ons\b", "strapon", s, flags=re.I)
    s = re.sub(r"\bstrap\s+on\b", "strapon", s, flags=re.I)
    s = re.sub(r"\bstrapons\b", "strapon", s)
    s = re.sub(r"[^\w\s]", " ", s)
    s = re.sub(r"\s+", " ", s).strip()
    s = re.sub(r"\s+", "", s)
    return s


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("--db", type=Path, default=Path("data/store_slim.db"))
    ap.add_argument("--out", type=Path, required=True, help="output JSONL path")
    ap.add_argument("--mode", choices=("fingerprint", "merge_signature", "both"), default="fingerprint")
    ap.add_argument("--min-size", type=int, default=2, help="minimum cluster size to emit")
    ap.add_argument("--max-clusters", type=int, default=0, help="cap clusters written (0 = no cap)")
    ap.add_argument("--sample-random", type=int, default=0, help="if set, shuffle clusters then take first N after sort")
    args = ap.parse_args()
    db = args.db.resolve()
    if not db.is_file():
        raise SystemExit(f"database not found: {db}")

    conn = sqlite3.connect(db)
    conn.row_factory = sqlite3.Row
    rows = conn.execute(
        """
        SELECT k.id, k.name, k.cluster, COALESCE(m.popularity, 0.0) AS pop
        FROM kink k
        LEFT JOIN fetlifekinkmeta m ON m.kink_id = k.id
        """,
    ).fetchall()
    conn.close()

    def emit_fingerprint() -> list[dict[str, object]]:
        buckets: dict[str, list[dict[str, object]]] = defaultdict(list)
        for r in rows:
            fp = fingerprint_play_title(str(r["name"]))
            if len(fp) < 4:
                continue
            buckets[fp].append(
                {
                    "kink_id": str(r["id"]),
                    "name": str(r["name"]),
                    "cluster": str(r["cluster"]),
                    "popularity": float(r["pop"] or 0.0),
                },
            )
        out = []
        for key, members in buckets.items():
            if len(members) < args.min_size:
                continue
            out.append({"cluster_key": key, "members": members})
        out.sort(key=lambda c: (-len(c["members"]), str(c["cluster_key"])))
        return out

    def emit_merge_sig() -> list[dict[str, object]]:
        buckets: dict[tuple[str, str], list[dict[str, object]]] = defaultdict(list)
        for r in rows:
            if not merge_graph_eligible(str(r["name"])):
                continue
            sig = merge_signature(str(r["name"]))
            buckets[(str(r["cluster"]), sig)].append(
                {
                    "kink_id": str(r["id"]),
                    "name": str(r["name"]),
                    "cluster": str(r["cluster"]),
                    "popularity": float(r["pop"] or 0.0),
                },
            )
        out = []
        for (cluster, sig), members in sorted(buckets.items(), key=lambda x: (-len(x[1]), x[0][0], x[0][1])):
            if len(members) < args.min_size:
                continue
            out.append({"cluster_key": f"{cluster}\t{sig}", "cluster": cluster, "signature": sig, "members": members})
        return out

    clusters_out: list[tuple[str, list[dict[str, object]]]] = []
    if args.mode in ("fingerprint", "both"):
        for c in emit_fingerprint():
            clusters_out.append(("fingerprint", c))
    if args.mode in ("merge_signature", "both"):
        for c in emit_merge_sig():
            clusters_out.append(("merge_signature", c))

    if args.sample_random > 0 and len(clusters_out) > args.sample_random:
        rng = random.Random(42)
        rng.shuffle(clusters_out)
        clusters_out = clusters_out[: args.sample_random]

    if args.max_clusters > 0:
        clusters_out = clusters_out[: args.max_clusters]

    args.out.parent.mkdir(parents=True, exist_ok=True)
    with args.out.open("w", encoding="utf-8") as fh:
        for mode, cluster in clusters_out:
            rec = {"mode": mode, **cluster}
            fh.write(json.dumps(rec, ensure_ascii=False) + "\n")

    print(json.dumps({"wrote": str(args.out), "lines": len(clusters_out), "mode": args.mode}, indent=2))
    return 0


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