data/preprocessor.py

"""
Feature engineering and dataset preparation for the recommendation pipeline.

Steps:
1. Encode user/movie IDs to contiguous 0-based integers
2. Compute user features (genre preferences, activity stats)
3. Compute item features (genre encoding, genome PCA embeddings, year)
4. Temporal train/val/test split
5. Generate negative samples for training
"""

import logging
import pickle
from pathlib import Path
from typing import Optional

import numpy as np
import pandas as pd
from sklearn.decomposition import TruncatedSVD
from sklearn.preprocessing import StandardScaler

from .loader import ALL_GENRES, load_all

logger = logging.getLogger(__name__)

GENRE_COLS = [
    f"g_{g.replace('-','_').replace('(','').replace(')','').replace(' ','_')}"
    for g in ALL_GENRES
]


class Preprocessor:
    """
    Encapsulates all feature-engineering logic.
    Call fit() on training data, then transform() on any split.
    Artifacts are saved to output_dir for serving-time use.
    """

    def __init__(self, genome_n_components: int = 32):
        self.genome_n_components = genome_n_components
        self.user_id_map: dict[int, int] = {}
        self.movie_id_map: dict[int, int] = {}
        self.user_id_unmap: dict[int, int] = {}
        self.movie_id_unmap: dict[int, int] = {}
        self.num_users: int = 0
        self.num_movies: int = 0
        self.genre_scaler = StandardScaler()
        self.genome_svd: Optional[TruncatedSVD] = None
        self.item_features: Optional[np.ndarray] = None  # (num_movies, feature_dim)
        self.user_features: Optional[np.ndarray] = None  # (num_users, feature_dim)
        self.movie_df: Optional[pd.DataFrame] = None
        self.fitted = False

    # ------------------------------------------------------------------
    # ID encoding
    # ------------------------------------------------------------------

    def _build_id_maps(self, ratings: pd.DataFrame, movies: pd.DataFrame) -> None:
        unique_users = sorted(ratings["userId"].unique())
        unique_movies = sorted(movies["movieId"].unique())

        self.user_id_map = {uid: idx for idx, uid in enumerate(unique_users)}
        self.movie_id_map = {mid: idx for idx, mid in enumerate(unique_movies)}
        self.user_id_unmap = {v: k for k, v in self.user_id_map.items()}
        self.movie_id_unmap = {v: k for k, v in self.movie_id_map.items()}
        self.num_users = len(unique_users)
        self.num_movies = len(unique_movies)
        logger.info(f"ID maps: {self.num_users:,} users, {self.num_movies:,} movies")

    def encode_ratings(self, ratings: pd.DataFrame) -> pd.DataFrame:
        df = ratings.copy()
        df["user_idx"] = df["userId"].map(self.user_id_map)
        df["movie_idx"] = df["movieId"].map(self.movie_id_map)
        # Drop rows where IDs aren't in training vocabulary
        df = df.dropna(subset=["user_idx", "movie_idx"])
        df["user_idx"] = df["user_idx"].astype(np.int32)
        df["movie_idx"] = df["movie_idx"].astype(np.int32)
        return df

    # ------------------------------------------------------------------
    # Item feature matrix
    # ------------------------------------------------------------------

    def _build_item_features(
        self, movies: pd.DataFrame, genome_scores: Optional[pd.DataFrame]
    ) -> np.ndarray:
        """Returns (num_movies, feature_dim) float32 array."""
        # Genre multi-hot — shape (num_movies, 20)
        genre_matrix = np.zeros((self.num_movies, len(GENRE_COLS)), dtype=np.float32)
        for _, row in movies.iterrows():
            idx = self.movie_id_map.get(row["movieId"])
            if idx is None:
                continue
            for i, col in enumerate(GENRE_COLS):
                genre_matrix[idx, i] = row.get(col, 0)

        # Normalised release year — shape (num_movies, 1)
        year_vec = np.zeros((self.num_movies, 1), dtype=np.float32)
        for _, row in movies.iterrows():
            idx = self.movie_id_map.get(row["movieId"])
            if idx is None:
                continue
            year_vec[idx, 0] = max(0, row.get("year", 2000) - 1900) / 100.0

        features = [genre_matrix, year_vec]

        # Genome tag embeddings via SVD — shape (num_movies, genome_n_components)
        if genome_scores is not None:
            genome_matrix = self._build_genome_matrix(movies, genome_scores)
            features.append(genome_matrix)

        item_feats = np.concatenate(features, axis=1)
        logger.info(f"Item feature dim: {item_feats.shape[1]}")
        return item_feats

    def _build_genome_matrix(
        self, movies: pd.DataFrame, genome_scores: pd.DataFrame
    ) -> np.ndarray:
        """Pivot genome scores to (num_movies, num_tags), then reduce with SVD."""
        pivot = genome_scores.pivot_table(
            index="movieId", columns="tagId", values="relevance", fill_value=0.0
        )
        # Only keep movies in our vocabulary
        valid_ids = [
            mid for mid in pivot.index if mid in self.movie_id_map
        ]
        pivot = pivot.loc[valid_ids]

        if self.genome_svd is None:
            n = min(self.genome_n_components, pivot.shape[1] - 1, pivot.shape[0] - 1)
            self.genome_svd = TruncatedSVD(n_components=n, random_state=42)
            reduced = self.genome_svd.fit_transform(pivot.values)
        else:
            reduced = self.genome_svd.transform(pivot.values)

        genome_matrix = np.zeros(
            (self.num_movies, self.genome_svd.n_components), dtype=np.float32
        )
        for orig_id, row in zip(valid_ids, reduced):
            idx = self.movie_id_map[orig_id]
            genome_matrix[idx] = row.astype(np.float32)

        logger.info(f"Genome SVD shape: {genome_matrix.shape}")
        return genome_matrix

    # ------------------------------------------------------------------
    # User feature matrix
    # ------------------------------------------------------------------

    def _build_user_features(
        self, ratings: pd.DataFrame, encoded: pd.DataFrame
    ) -> np.ndarray:
        """
        Returns (num_users, feature_dim) float32 array.
        Features: genre preferences (20), avg_rating (1), log_count (1) → 22-dim.

        Fully vectorised: uses np.bincount (compiled C) per genre column instead
        of a Python groupby loop, reducing runtime from minutes to seconds on 2M rows.
        """
        n_genres = len(GENRE_COLS)
        genre_pref = np.zeros((self.num_users, n_genres), dtype=np.float64)
        avg_rating = np.zeros((self.num_users, 1), dtype=np.float32)
        log_count = np.zeros((self.num_users, 1), dtype=np.float32)

        if self.item_features is not None:
            user_idxs = encoded["user_idx"].values.astype(np.int64)
            movie_idxs = encoded["movie_idx"].values.astype(np.int64)
            rating_vals = encoded["rating"].values.astype(np.float64)

            item_genre = self.item_features[:, :n_genres].astype(np.float64)

            # Weighted genre signal per interaction: [N, n_genres]
            genre_weighted = item_genre[movie_idxs] * rating_vals[:, None]

            # Accumulate per user with np.bincount (fast C loop, one call per genre)
            for g in range(n_genres):
                genre_pref[:, g] = np.bincount(
                    user_idxs, weights=genre_weighted[:, g], minlength=self.num_users
                )

            # Per-user interaction counts and rating sums
            counts = np.bincount(user_idxs, minlength=self.num_users).astype(np.float64)
            rating_sum = np.bincount(user_idxs, weights=rating_vals, minlength=self.num_users)

            mask = counts > 0
            genre_pref[mask] /= counts[mask, None]
            avg_rating[mask, 0] = (rating_sum[mask] / counts[mask] / 5.0).astype(np.float32)
            log_count[mask, 0] = (np.log1p(counts[mask]) / 10.0).astype(np.float32)

        user_feats = np.concatenate([genre_pref.astype(np.float32), avg_rating, log_count], axis=1)
        logger.info(f"User feature dim: {user_feats.shape[1]}")
        return user_feats

    # ------------------------------------------------------------------
    # Temporal split
    # ------------------------------------------------------------------

    def temporal_split(
        self,
        encoded: pd.DataFrame,
        val_frac: float = 0.10,
        test_frac: float = 0.10,
    ) -> tuple[pd.DataFrame, pd.DataFrame, pd.DataFrame]:
        """
        Split by timestamp: oldest interactions → train, newest → test.
        This avoids data leakage from the future.
        """
        n = len(encoded)
        test_start = int(n * (1 - test_frac))
        val_start = int(n * (1 - val_frac - test_frac))

        train = encoded.iloc[:val_start].copy()
        val = encoded.iloc[val_start:test_start].copy()
        test = encoded.iloc[test_start:].copy()

        logger.info(
            f"Split — train: {len(train):,}, val: {len(val):,}, test: {len(test):,}"
        )
        return train, val, test

    # ------------------------------------------------------------------
    # Negative sampling
    # ------------------------------------------------------------------

    def generate_negatives(
        self,
        encoded: pd.DataFrame,
        n_neg: int = 4,
        seed: int = 42,
    ) -> pd.DataFrame:
        """
        For each positive (user, movie) pair, sample n_neg random negatives.
        Returns DataFrame with columns: user_idx, pos_movie_idx, neg_movie_idx.

        Fully vectorised: samples negatives in one numpy call, then fixes the rare
        case where a sampled negative equals the positive (collision rate ≈ 0.5% on
        MovieLens given ~145 ratings per user vs 27K movies).  Full seen-set filtering
        is skipped intentionally — the sparsity of the dataset means almost all random
        items are true negatives, which is standard practice for large-scale BPR.
        """
        rng = np.random.default_rng(seed)

        user_idxs = encoded["user_idx"].values.astype(np.int32)
        pos_idxs = encoded["movie_idx"].values.astype(np.int32)
        N = len(user_idxs)

        # Sample n_neg negatives per interaction in one vectorised call
        neg_samples = rng.integers(0, self.num_movies, size=(N, n_neg), dtype=np.int32)

        # Fix the trivial collision: neg == pos (rare but easy to handle vectorially)
        for k in range(n_neg):
            collision = neg_samples[:, k] == pos_idxs
            neg_samples[collision, k] = (pos_idxs[collision] + 1) % self.num_movies

        # Flatten: each positive gets n_neg rows
        uid_rep = np.repeat(user_idxs, n_neg)
        pos_rep = np.repeat(pos_idxs, n_neg)
        neg_flat = neg_samples.ravel()

        df = pd.DataFrame({
            "user_idx": uid_rep,
            "pos_movie_idx": pos_rep,
            "neg_movie_idx": neg_flat,
        })
        logger.info(f"Generated {len(df):,} training triplets")
        return df

    # ------------------------------------------------------------------
    # Fit / transform / save
    # ------------------------------------------------------------------

    def fit(
        self,
        ratings: pd.DataFrame,
        movies: pd.DataFrame,
        genome_scores: Optional[pd.DataFrame] = None,
    ) -> "Preprocessor":
        self.movie_df = movies
        self._build_id_maps(ratings, movies)
        encoded = self.encode_ratings(ratings)

        self.item_features = self._build_item_features(movies, genome_scores)
        self.user_features = self._build_user_features(ratings, encoded)
        self.fitted = True
        return self

    def fit_transform(
        self,
        ratings: pd.DataFrame,
        movies: pd.DataFrame,
        genome_scores: Optional[pd.DataFrame] = None,
    ) -> pd.DataFrame:
        self.fit(ratings, movies, genome_scores)
        return self.encode_ratings(ratings)

    def save(self, output_dir: str | Path) -> None:
        output_dir = Path(output_dir)
        output_dir.mkdir(parents=True, exist_ok=True)

        # Temporarily strip movie_df before pickling.
        # Pandas DataFrames with list/object columns trigger a NDArrayBacked
        # __setstate__ incompatibility across pandas 2.x minor versions.
        # We save movie metadata as a plain CSV instead — always portable.
        movie_df_backup = self.movie_df
        self.movie_df = None

        with open(output_dir / "preprocessor.pkl", "wb") as f:
            pickle.dump(self, f)

        self.movie_df = movie_df_backup  # restore in-memory object

        # Save movie metadata as version-safe CSV
        if movie_df_backup is not None:
            meta_cols = [c for c in ["movieId", "title", "genres", "year"] if c in movie_df_backup.columns]
            movie_df_backup[meta_cols].to_csv(output_dir / "movie_meta.csv", index=False)

        np.save(output_dir / "item_features.npy", self.item_features)
        np.save(output_dir / "user_features.npy", self.user_features)
        logger.info(f"Preprocessor saved to {output_dir}")

    @classmethod
    def load(cls, output_dir: str | Path) -> "Preprocessor":
        output_dir = Path(output_dir).resolve()  # always absolute — safe regardless of CWD
        with open(output_dir / "preprocessor.pkl", "rb") as f:
            obj = pickle.load(f)

        # Reload movie_df from CSV (bypasses pandas pickle version issues)
        meta_path = output_dir / "movie_meta.csv"
        if meta_path.exists():
            df = pd.read_csv(meta_path)
            df["genre_list"] = df["genres"].apply(
                lambda g: [] if str(g) == "(no genres listed)" else str(g).split("|")
            )
            obj.movie_df = df
        else:
            obj.movie_df = None

        # Also reload the numpy feature matrices in case they were updated
        item_path = output_dir / "item_features.npy"
        user_path = output_dir / "user_features.npy"
        if item_path.exists():
            obj.item_features = np.load(str(item_path))
        if user_path.exists():
            obj.user_features = np.load(str(user_path))

        return obj


# ------------------------------------------------------------------
# Convenience entry-point
# ------------------------------------------------------------------

def preprocess_all(
    data_dir: str | Path,
    output_dir: str | Path,
    sample_frac: Optional[float] = None,
    genome_n_components: int = 32,
) -> tuple["Preprocessor", dict]:
    """
    Full preprocessing pipeline: load → encode → features → split → save.
    Returns (preprocessor, splits) where splits = {train, val, test, negatives}.
    """
    data = load_all(data_dir, sample_frac=sample_frac)

    prep = Preprocessor(genome_n_components=genome_n_components)
    encoded = prep.fit_transform(
        data["ratings"], data["movies"], data["genome_scores"]
    )
    train, val, test = prep.temporal_split(encoded)
    negatives = prep.generate_negatives(train, n_neg=4)

    prep.save(output_dir)

    return prep, {
        "train": train,
        "val": val,
        "test": test,
        "negatives": negatives,
        "movies": data["movies"],
    }