scripts/propose_scenario_bridges.py

#!/usr/bin/env python3
"""Propose scenario (long-tail) -> canonical kink bridges from existing catalog edges.

Reads the local SQLite store, flags scenario-like kinks with heuristics, pairs each
with short high-popularity catalog neighbors, and writes CSV (or applies reviewed rows).

Convention: ``left_kink_id`` = scenario-style kink, ``right_kink_id`` = canonical target.
Uses :meth:`Backend.add_similarity_edge` with ``similarity_type="scenario_bridge"``.
"""
from __future__ import annotations

import argparse
import csv
import sys
from pathlib import Path

from sqlmodel import Session, select

from backend import Backend
from models import FetlifeKinkMeta, Kink

ROOT = Path(__file__).resolve().parent.parent
DEFAULT_DB = ROOT / "data" / "store.db"

SCENARIO_SUBSTRINGS = (
    "scenario",
    "story",
    "fantasy",
    "roleplay",
    "role play",
    "scene",
    "acting",
    "narrative",
)


def _meta_maps(session) -> tuple[dict[str, FetlifeKinkMeta], dict[str, Kink]]:
    metas = session.exec(select(FetlifeKinkMeta)).all()
    kinks = session.exec(select(Kink)).all()
    return {m.kink_id: m for m in metas}, {k.id: k for k in kinks}


def is_scenario_candidate(
    kink: Kink,
    meta: FetlifeKinkMeta | None,
    *,
    min_name_len: int,
    min_name_len_keyword_only: int,
    min_popularity: float,
    min_similar_count: int,
    require_keyword: bool,
) -> bool:
    name = kink.name
    pop = float(meta.popularity) if meta else 0.0
    sim = int(meta.similar_count) if meta else 0
    if pop < min_popularity or sim < min_similar_count:
        return False
    lower = name.lower()
    keyword = any(s in lower for s in SCENARIO_SUBSTRINGS)
    if require_keyword:
        return len(name) >= min_name_len_keyword_only and keyword
    long_enough = len(name) >= min_name_len
    very_long = len(name) >= min_name_len + 35
    return long_enough and (keyword or very_long)


def propose_rows(
    backend: Backend,
    *,
    top_n: int,
    canonical_max_name_len: int,
    canonical_min_popularity: float,
    min_name_len: int,
    min_name_len_keyword_only: int,
    min_popularity: float,
    min_similar_count: int,
    require_keyword: bool,
) -> list[dict[str, str]]:
    out: list[dict[str, str]] = []
    with Session(backend.engine) as session:
        meta_by_id, kink_by_id = _meta_maps(session)

    for kid, kink in sorted(kink_by_id.items()):
        meta = meta_by_id.get(kid)
        if not is_scenario_candidate(
            kink,
            meta,
            min_name_len=min_name_len,
            min_name_len_keyword_only=min_name_len_keyword_only,
            min_popularity=min_popularity,
            min_similar_count=min_similar_count,
            require_keyword=require_keyword,
        ):
            continue

        edges = backend._edges_for_kink(kid, limit=80)
        catalog_edges = [e for e in edges if e.get("type") == "catalog"]
        taken = 0
        for edge in catalog_edges:
            if taken >= top_n:
                break
            rid = edge["id"]
            target = kink_by_id.get(rid)
            if not target:
                continue
            rmeta = meta_by_id.get(rid)
            rpop = float(rmeta.popularity) if rmeta else 0.0
            if rpop < canonical_min_popularity:
                continue
            if len(target.name) > canonical_max_name_len:
                continue
            if rid == kid:
                continue
            reason = f"catalog_sim={edge['score']:.4f}; canon_pop={rpop:.0f}"
            out.append(
                {
                    "left_kink_id": kid,
                    "right_kink_id": rid,
                    "proposed_score": f"{float(edge['score']):.6f}",
                    "reason": reason,
                }
            )
            taken += 1
    return out


def write_csv(rows: list[dict[str, str]], dest: Path | None) -> None:
    fieldnames = ["left_kink_id", "right_kink_id", "proposed_score", "reason"]
    if dest:
        with dest.open("w", newline="", encoding="utf-8") as f:
            w = csv.DictWriter(f, fieldnames=fieldnames)
            w.writeheader()
            w.writerows(rows)
    else:
        w = csv.DictWriter(sys.stdout, fieldnames=fieldnames)
        w.writeheader()
        w.writerows(rows)


def apply_whitelist(backend: Backend, path: Path, *, method: str, version: str) -> int:
    """Apply rows from a CSV with the same columns as proposal output."""
    applied = 0
    with path.open(newline="", encoding="utf-8") as f:
        reader = csv.DictReader(f)
        for row in reader:
            left = (row.get("left_kink_id") or "").strip()
            right = (row.get("right_kink_id") or "").strip()
            if not left or not right:
                continue
            score_s = (row.get("proposed_score") or "0.8").strip()
            try:
                score = float(score_s)
            except ValueError:
                continue
            score = max(0.0, min(1.0, score))
            backend.add_similarity_edge(left, right, "scenario_bridge", score, method, version)
            applied += 1
    return applied


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--db", type=Path, default=DEFAULT_DB, help="Path to store.db")
    parser.add_argument("--output", "-o", type=Path, default=None, help="Write CSV here (default: stdout)")
    parser.add_argument("--top-n", type=int, default=3, help="Catalog neighbors per scenario kink")
    parser.add_argument("--canonical-max-name-len", type=int, default=45)
    parser.add_argument("--canonical-min-popularity", type=float, default=400.0)
    parser.add_argument("--min-name-len", type=int, default=48)
    parser.add_argument("--min-name-len-keyword-only", type=int, default=36)
    parser.add_argument("--min-popularity", type=float, default=120.0)
    parser.add_argument("--min-similar-count", type=int, default=4)
    parser.add_argument(
        "--require-keyword",
        action="store_true",
        help="Only treat kinks as scenarios when a scenario keyword matches (still uses length floor)",
    )
    parser.add_argument(
        "--apply",
        action="store_true",
        help="Insert scenario_bridge edges from a reviewed CSV (requires --whitelist)",
    )
    parser.add_argument(
        "--whitelist",
        type=Path,
        default=None,
        help="CSV from a prior run (edited) for --apply",
    )
    parser.add_argument(
        "--method",
        default="catalog_neighbor_v1",
        help="Similarity method string stored on each edge",
    )
    parser.add_argument("--version", default="v1")
    args = parser.parse_args()

    if not args.db.exists():
        print(f"Database not found: {args.db}", file=sys.stderr)
        sys.exit(1)

    backend = Backend(args.db)

    if args.apply:
        if not args.whitelist or not args.whitelist.is_file():
            print("--apply requires an existing --whitelist CSV path.", file=sys.stderr)
            sys.exit(1)
        n = apply_whitelist(backend, args.whitelist, method=args.method, version=args.version)
        print(f"Applied {n} scenario_bridge edges.", file=sys.stderr)
        return

    rows = propose_rows(
        backend,
        top_n=args.top_n,
        canonical_max_name_len=args.canonical_max_name_len,
        canonical_min_popularity=args.canonical_min_popularity,
        min_name_len=args.min_name_len,
        min_name_len_keyword_only=args.min_name_len_keyword_only,
        min_popularity=args.min_popularity,
        min_similar_count=args.min_similar_count,
        require_keyword=args.require_keyword,
    )
    write_csv(rows, args.output)
    print(f"Wrote {len(rows)} proposed rows.", file=sys.stderr)


if __name__ == "__main__":
    main()