backend/mpa_classifier.py

"""
ReelShield MPA-rating classifier.

A multinomial logistic regression that predicts a film's MPA rating bucket
from its Gemini-generated warning vector. Trained on the cached movies that
have both a stored us_certification and an avg-confidence above the ghost
threshold.

Two uses:
  1. Sanity check on Gemini outputs. When predict_one(film) disagrees with
     the stored us_certification, the film is flagged for re-assessment —
     Gemini probably under- or over-reported severity somewhere.
  2. Real-ML defensibility. Unlike the sentence-transformer (which is a
     pretrained black box used zero-shot), this model has its own train/test
     split, learned coefficients, and reported metrics — i.e. classic ML.

Categories (the 9-dim warning vector):
    violence_gore, self_harm_suicide, miscarriage_pregnancy_loss,
    sexual_content_nudity, animal_abuse, substances, language,
    horror_intensity, flashing_lights

Feature vector per film (19 dimensions):
    [severity_0..severity_8, confidence_0..confidence_8, year_normalized]

Label buckets:
    0 = family    (G, PG)
    1 = teen      (PG-13)
    2 = adult     (R, NC-17)
NR and empty MPA fields are dropped.
"""

from __future__ import annotations

import json
import os
from dataclasses import dataclass
from typing import Iterable, Sequence

import numpy as np

# Lazy-imported sklearn / joblib so importing this module is cheap until you
# actually train or predict.

WARNING_CATEGORIES = [
    "violence_gore", "self_harm_suicide",
    "miscarriage_pregnancy_loss", "sexual_content_nudity",
    "animal_abuse", "substances", "language",
    "horror_intensity", "flashing_lights",
]

LABEL_NAMES = ["family", "teen", "adult"]
CERT_TO_LABEL = {
    "G": 0, "PG": 0,
    "PG-13": 1,
    "R": 2, "NC-17": 2,
}
GHOST_CONFIDENCE_FLOOR = 0.4

MODEL_PATH_DEFAULT = os.environ.get("MPA_MODEL_PATH", "/data/mpa_classifier.pkl")


@dataclass
class TrainingData:
    X: np.ndarray            # shape (n_samples, 19)
    y: np.ndarray            # shape (n_samples,)
    tmdb_ids: list[int]
    feature_names: list[str]


def _year_norm(year_str: str | None) -> float:
    """Map a 4-digit year string to a [0, 1] feature.
    Roughly: 1920 -> 0.0, 2025 -> 1.0. Out-of-range or empty -> 0.5."""
    try:
        y = int(str(year_str or "")[:4])
    except (ValueError, TypeError):
        return 0.5
    if y < 1920 or y > 2025:
        return 0.5
    return (y - 1920) / (2025 - 1920)


def _film_features(warnings_json: dict, year: str | None) -> tuple[np.ndarray, float]:
    """Return (feature_vector_19d, avg_confidence)."""
    sf = warnings_json.get("spoiler_free") or {}
    sevs = np.zeros(9, dtype=np.float32)
    confs = np.zeros(9, dtype=np.float32)
    for i, cat in enumerate(WARNING_CATEGORIES):
        d = sf.get(cat) or {}
        sevs[i] = float(d.get("severity") or 0)
        confs[i] = float(d.get("confidence") or 0.0)
    avg_conf = float(np.mean(confs)) if confs.size else 0.0
    yn = _year_norm(year)
    vec = np.concatenate([sevs, confs, np.array([yn], dtype=np.float32)])
    return vec, avg_conf


def feature_names() -> list[str]:
    sev_names = [f"sev_{c}" for c in WARNING_CATEGORIES]
    conf_names = [f"conf_{c}" for c in WARNING_CATEGORIES]
    return sev_names + conf_names + ["year_norm"]


def extract_training_data(db_path: str) -> TrainingData:
    """Pull all films whose MPA rating maps to a label AND whose avg
    confidence is above the ghost floor. Returns aligned arrays."""
    import sqlite3
    conn = sqlite3.connect(db_path)
    rows = conn.execute(
        "SELECT m.tmdb_id, m.year, m.metadata_json, cw.warnings_json "
        "FROM movies m JOIN content_warnings cw USING(tmdb_id)"
    ).fetchall()
    conn.close()

    X_list: list[np.ndarray] = []
    y_list: list[int] = []
    ids: list[int] = []
    for tmdb_id, year, mj, wj in rows:
        try:
            meta = json.loads(mj or "{}")
            warn = json.loads(wj or "{}")
        except json.JSONDecodeError:
            continue
        cert = (meta.get("us_certification") or "").strip().upper()
        if cert not in CERT_TO_LABEL:
            continue
        feat, avg_conf = _film_features(warn, year)
        if avg_conf < GHOST_CONFIDENCE_FLOOR:
            continue
        X_list.append(feat)
        y_list.append(CERT_TO_LABEL[cert])
        ids.append(tmdb_id)

    if not X_list:
        raise RuntimeError("No labelable films found — check cache + confidence floor.")
    return TrainingData(
        X=np.vstack(X_list),
        y=np.array(y_list, dtype=np.int64),
        tmdb_ids=ids,
        feature_names=feature_names(),
    )


def train_model(X: np.ndarray, y: np.ndarray, random_state: int = 42):
    """Fit a multinomial LogisticRegression. Uses class_weight='balanced'
    so the relative class sizes (R-rated tends to dominate) don't bias
    the decision boundary."""
    from sklearn.linear_model import LogisticRegression
    # sklearn >= 1.7 dropped the multi_class kwarg; lbfgs auto-detects
    # multinomial when there are >2 classes.
    model = LogisticRegression(
        solver="lbfgs",
        class_weight="balanced",
        C=1.0,
        max_iter=1000,
        random_state=random_state,
    )
    model.fit(X, y)
    return model


def evaluate(model, X_test: np.ndarray, y_test: np.ndarray) -> dict:
    from sklearn.metrics import (
        accuracy_score, classification_report, confusion_matrix, f1_score,
    )
    y_pred = model.predict(X_test)
    return {
        "accuracy":      float(accuracy_score(y_test, y_pred)),
        "macro_f1":      float(f1_score(y_test, y_pred, average="macro")),
        "weighted_f1":   float(f1_score(y_test, y_pred, average="weighted")),
        "report":        classification_report(y_test, y_pred, target_names=LABEL_NAMES, digits=3),
        "confusion":     confusion_matrix(y_test, y_pred).tolist(),
    }


def save_model(model, path: str = MODEL_PATH_DEFAULT) -> None:
    import joblib
    os.makedirs(os.path.dirname(path) or ".", exist_ok=True)
    joblib.dump({"model": model, "feature_names": feature_names()}, path)


def load_model(path: str = MODEL_PATH_DEFAULT):
    """Returns (model, feature_names) or None if no model on disk."""
    import joblib
    if not os.path.exists(path):
        return None
    bundle = joblib.load(path)
    return bundle["model"], bundle["feature_names"]


def predict_one(model, warnings_json: dict, year: str | None) -> dict:
    """Predict label + probabilities for a single film."""
    feat, _ = _film_features(warnings_json, year)
    pred_idx = int(model.predict(feat.reshape(1, -1))[0])
    proba    = model.predict_proba(feat.reshape(1, -1))[0]
    return {
        "label_idx":     pred_idx,
        "label":         LABEL_NAMES[pred_idx],
        "probabilities": {name: float(p) for name, p in zip(LABEL_NAMES, proba)},
    }


def disagreement_report(model, db_path: str) -> list[dict]:
    """For every labelable film, compare prediction vs stored MPA.
    Returns rows where they disagree — these are candidates for re-seeding
    or human review."""
    import sqlite3
    conn = sqlite3.connect(db_path)
    rows = conn.execute(
        "SELECT m.tmdb_id, m.title, m.year, m.metadata_json, cw.warnings_json "
        "FROM movies m JOIN content_warnings cw USING(tmdb_id)"
    ).fetchall()
    conn.close()

    out: list[dict] = []
    for tmdb_id, title, year, mj, wj in rows:
        try:
            meta = json.loads(mj or "{}")
            warn = json.loads(wj or "{}")
        except json.JSONDecodeError:
            continue
        cert = (meta.get("us_certification") or "").strip().upper()
        if cert not in CERT_TO_LABEL:
            continue
        feat, avg_conf = _film_features(warn, year)
        if avg_conf < GHOST_CONFIDENCE_FLOOR:
            continue
        pred_idx = int(model.predict(feat.reshape(1, -1))[0])
        actual_idx = CERT_TO_LABEL[cert]
        if pred_idx != actual_idx:
            proba = model.predict_proba(feat.reshape(1, -1))[0]
            out.append({
                "tmdb_id":     tmdb_id,
                "title":       title,
                "year":        year,
                "mpa_actual":  cert,
                "predicted":   LABEL_NAMES[pred_idx],
                "confidence":  float(proba[pred_idx]),
            })
    out.sort(key=lambda d: -d["confidence"])
    return out