backend/recsys_data.py

"""Dataset export helpers for OTS recommender experiments.

The app keeps its production rules in the backend; this module converts that state
into standard implicit-feedback records that external recommender libraries can use.
"""
from __future__ import annotations

from collections import defaultdict
from dataclasses import dataclass, field
from functools import cached_property
from pathlib import Path
from typing import Any

from sqlmodel import Session, select

from backend.constants import RATING_WEIGHTS
from backend.discovery_lanes import discovery_lane_for_kink
from backend.scenarios import scenario_title_fields
from backend.scrape_artifacts import clean_scraped_catalog_text, clean_scraped_notes
from models import (
    FetlifeKinkMeta,
    FetlifeUserFetish,
    Kink,
    KinkContentType,
    KinkScenarioParent,
    PlayPreference,
    ScenarioPreference,
)

APP_RATING_VALUES = {
    "love": 3.0,
    "like": 2.0,
    "curious": 1.0,
    "not_interested": 0.0,
    "hard_no": -2.0,
}

FETLIFE_BUCKET_VALUES = {
    "into": 2.0,
    "curious_about": 1.0,
    "soft_limits": -0.5,
    "hard_limits": -2.0,
}


@dataclass(frozen=True)
class RecsysInteraction:
    user_id: str
    kink_id: str
    rating: float
    source: str
    timestamp: int


@dataclass(frozen=True)
class RecsysItem:
    kink_id: str
    name: str
    content_kind: str
    cluster: str
    lane: str
    popularity: float
    is_scenario: bool
    title_surface_as_scenario: bool
    shared_eligible: bool
    scenario_parent_ids: tuple[str, ...] = field(default_factory=tuple)


@dataclass
class RecsysDataset:
    interactions: list[RecsysInteraction]
    items: dict[str, RecsysItem]

    @cached_property
    def recommendable_item_ids(self) -> set[str]:
        return {
            item.kink_id
            for item in self.items.values()
            if item.content_kind == "play"
            and item.shared_eligible
            and not item.is_scenario
            and not item.title_surface_as_scenario
        }

    def canonical_item_ids(self, kink_id: str) -> tuple[str, ...]:
        item = self.items.get(kink_id)
        if not item:
            return ()
        if item.is_scenario:
            parents = tuple(pid for pid in item.scenario_parent_ids if pid in self.recommendable_item_ids)
            return parents
        if kink_id in self.recommendable_item_ids:
            return (kink_id,)
        return ()

    def positive_training_triples(self, *, min_rating: float = 0.5) -> list[tuple[str, str, float]]:
        """Return canonical positive interactions for OTS implicit-feedback models."""
        by_pair: dict[tuple[str, str], float] = {}
        for row in self.interactions:
            if row.rating < min_rating:
                continue
            for item_id in self.canonical_item_ids(row.kink_id):
                key = (row.user_id, item_id)
                by_pair[key] = max(by_pair.get(key, 0.0), float(row.rating))
        return [(user_id, item_id, rating) for (user_id, item_id), rating in sorted(by_pair.items())]

    def seen_item_ids_by_user(self, *, include_negative: bool = True) -> dict[str, set[str]]:
        seen: dict[str, set[str]] = {}
        for row in self.interactions:
            if not include_negative and row.rating <= 0:
                continue
            canonical_ids = self.canonical_item_ids(row.kink_id)
            if not canonical_ids:
                continue
            bucket = seen.setdefault(row.user_id, set())
            bucket.update(canonical_ids)
        return seen

    def popularity_ranked_items(self) -> list[str]:
        return [
            item.kink_id
            for item in sorted(
                self.items.values(),
                key=lambda item: (item.popularity, item.name.lower(), item.kink_id),
                reverse=True,
            )
            if item.kink_id in self.recommendable_item_ids
        ]


def _rating_value(state: str) -> float:
    if state in APP_RATING_VALUES:
        return APP_RATING_VALUES[state]
    return float(RATING_WEIGHTS.get(state, 0.0))


def _atomic_value(value: Any) -> str:
    return str(value if value is not None else "").replace("\t", " ").replace("\n", " ").strip()


def _item_from_catalog_payload(kink: dict[str, Any]) -> RecsysItem:
    return RecsysItem(
        kink_id=str(kink["id"]),
        name=str(kink.get("name", "")),
        content_kind=str(kink.get("content_kind") or ""),
        cluster=str(kink.get("cluster") or ""),
        lane=discovery_lane_for_kink(kink),
        popularity=float(kink.get("popularity", 0.0) or 0.0),
        is_scenario=bool(kink.get("is_scenario")),
        title_surface_as_scenario=bool(kink.get("title_surface_as_scenario")),
        shared_eligible=bool(kink.get("shared_eligible")),
        scenario_parent_ids=tuple(str(pid) for pid in (kink.get("scenario_parent_ids") or [])),
    )


def _items_from_store(backend: Any) -> dict[str, RecsysItem]:
    with Session(backend.engine) as session:
        kinks = session.exec(select(Kink).order_by(Kink.id)).all()
        type_rows = {row.kink_id: row for row in session.exec(select(KinkContentType)).all()}
        meta_rows = {row.kink_id: row for row in session.exec(select(FetlifeKinkMeta)).all()}
        scenario_rows = session.exec(select(KinkScenarioParent)).all()
    parents_by_scenario: dict[str, list[str]] = defaultdict(list)
    for row in scenario_rows:
        parents_by_scenario[row.scenario_kink_id].append(row.parent_kink_id)
    items: dict[str, RecsysItem] = {}
    for kink in kinks:
        meta = meta_rows.get(kink.id)
        source_backed_popularity = float(meta.popularity) if meta else 0.0
        raw_notes = kink.notes.strip()
        display_notes = clean_scraped_notes(kink.id, kink.name, raw_notes)
        definition = clean_scraped_catalog_text(kink.name, kink.short_definition)
        explicit_popularity = float(backend._extract_popularity_from_notes(raw_notes) or 0.0)
        popularity = source_backed_popularity or explicit_popularity
        type_row = type_rows.get(kink.id)
        base_payload = {
            "id": kink.id,
            "name": kink.name,
            "cluster": kink.cluster,
            "content_kind": type_row.content_kind if type_row else "",
            "definition": definition,
            "summary": definition,
            "notes": display_notes,
            "popularity": popularity,
            "source_backed_popularity": source_backed_popularity,
            "similar_count": int(meta.similar_count) if meta else 0,
            "filtered_asset_count": 0,
            "raw_asset_count": 0,
        }
        content_kind = base_payload["content_kind"] or backend._content_kind(base_payload)
        base_payload["content_kind"] = content_kind
        scenario_parent_ids = tuple(sorted(parents_by_scenario.get(kink.id, [])))
        is_scenario = bool(scenario_parent_ids)
        _, title_surface = (
            scenario_title_fields(kink.name) if content_kind == "play" else (0.0, False)
        )
        flags = backend._derived_product_flags(base_payload)
        shared_eligible = bool(flags.get("shared_eligible"))
        if is_scenario or title_surface:
            shared_eligible = False
        items[kink.id] = RecsysItem(
            kink_id=kink.id,
            name=kink.name,
            content_kind=str(content_kind),
            cluster=kink.cluster,
            lane=discovery_lane_for_kink(base_payload),
            popularity=popularity,
            is_scenario=is_scenario,
            title_surface_as_scenario=title_surface,
            shared_eligible=shared_eligible,
            scenario_parent_ids=scenario_parent_ids,
        )
    return items


def export_recsys_dataset(
    backend: Any,
    *,
    include_fetlife_samples: bool = True,
    include_scenario_preferences: bool = True,
) -> RecsysDataset:
    items = _items_from_store(backend)
    interactions: list[RecsysInteraction] = []
    ts = 1
    with Session(backend.engine) as session:
        for row in session.exec(select(PlayPreference).order_by(PlayPreference.user_id, PlayPreference.kink_id)).all():
            if row.kink_id not in items:
                continue
            interactions.append(
                RecsysInteraction(row.user_id, row.kink_id, _rating_value(row.interest_state), "play_preference", ts)
            )
            ts += 1

        if include_scenario_preferences:
            for row in session.exec(
                select(ScenarioPreference).order_by(
                    ScenarioPreference.user_id,
                    ScenarioPreference.parent_kink_id,
                    ScenarioPreference.scenario_kink_id,
                )
            ).all():
                rating = _rating_value(row.interest_state)
                if row.scenario_kink_id in items:
                    interactions.append(
                        RecsysInteraction(
                            row.user_id,
                            row.scenario_kink_id,
                            rating,
                            "scenario_preference",
                            ts,
                        )
                    )
                    ts += 1
                if row.parent_kink_id in items:
                    interactions.append(
                        RecsysInteraction(
                            row.user_id,
                            row.parent_kink_id,
                            rating * 0.75,
                            "scenario_parent_signal",
                            ts,
                        )
                    )
                    ts += 1

        if include_fetlife_samples:
            for row in session.exec(
                select(FetlifeUserFetish).order_by(FetlifeUserFetish.nickname, FetlifeUserFetish.kink_id)
            ).all():
                if row.kink_id not in items:
                    continue
                rating = FETLIFE_BUCKET_VALUES.get(row.bucket, 0.0)
                interactions.append(
                    RecsysInteraction(
                        f"fetlife:{row.nickname}",
                        row.kink_id,
                        rating,
                        f"fetlife_{row.bucket}",
                        ts,
                    )
                )
                ts += 1

    return RecsysDataset(interactions=interactions, items=items)


def write_recbole_atomic_files(
    dataset: RecsysDataset,
    output_dir: Path,
    *,
    dataset_name: str = "kink_cli",
) -> dict[str, Path]:
    """Write RecBole-compatible atomic files for external training environments."""
    output_dir.mkdir(parents=True, exist_ok=True)
    inter_path = output_dir / f"{dataset_name}.inter"
    item_path = output_dir / f"{dataset_name}.item"

    with inter_path.open("w", encoding="utf-8") as fh:
        fh.write("user_id:token\titem_id:token\trating:float\ttimestamp:float\tsource:token\n")
        for row in dataset.interactions:
            fh.write(
                "\t".join(
                    [
                        _atomic_value(row.user_id),
                        _atomic_value(row.kink_id),
                        f"{row.rating:.6g}",
                        str(row.timestamp),
                        _atomic_value(row.source),
                    ]
                )
                + "\n"
            )

    with item_path.open("w", encoding="utf-8") as fh:
        fh.write(
            "item_id:token\tname:token_seq\tcontent_kind:token\tcluster:token\tlane:token\t"
            "popularity:float\tis_scenario:float\tshared_eligible:float\tparent_ids:token_seq\n"
        )
        for item in sorted(dataset.items.values(), key=lambda item: item.kink_id):
            fh.write(
                "\t".join(
                    [
                        _atomic_value(item.kink_id),
                        _atomic_value(item.name),
                        _atomic_value(item.content_kind),
                        _atomic_value(item.cluster),
                        _atomic_value(item.lane),
                        f"{item.popularity:.6g}",
                        "1" if item.is_scenario else "0",
                        "1" if item.shared_eligible else "0",
                        _atomic_value(" ".join(item.scenario_parent_ids)),
                    ]
                )
                + "\n"
            )

    return {"interactions": inter_path, "items": item_path}