src/meta_encoder.py

"""Metadata encoder: turn each row's product-side fields into a fixed-length vector.

Output layout (concatenated):
    [ TF-IDF on features_text (META_TFIDF_DIM)       # textual product attributes
    | numeric vector (META_NUM_DIM)                  # price + ratings + flags
    ]
Total raw dim = META_TFIDF_DIM + META_NUM_DIM.

The MLP inside the model maps this to META_HIDDEN_DIM. We keep encoding as a
separate, plain-numpy step so it can be fit once on train and reused without
touching the model.
"""
import logging
import pickle
from dataclasses import dataclass
from pathlib import Path
from typing import Optional, Sequence

import numpy as np
import pandas as pd
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.preprocessing import StandardScaler

from . import config as cfg


logger = logging.getLogger(__name__)


# Columns we expect on the input df (produced by preprocess.join_and_clean).
# Missing columns are handled gracefully.
FEATURES_TEXT_COL = "features_text"
CATEGORIES_TEXT_COL = "categories_text"
PRICE_COL = "price"
AVG_RATING_COL = "average_rating"
RATING_NUMBER_COL = "rating_number"


def _safe_text(series: pd.Series) -> pd.Series:
    return series.fillna("").astype(str)


def _safe_numeric(series: pd.Series) -> pd.Series:
    return pd.to_numeric(series, errors="coerce")


@dataclass
class MetaEncoder:
    """Fit on train df, then transform any df (train/val/test/new) to a (N, D) matrix."""

    tfidf_dim: int = cfg.META_TFIDF_DIM
    num_dim: int = cfg.META_NUM_DIM  # price, avg_rating, log_rating_number, price_missing_flag

    # Filled by fit()
    tfidf: Optional[TfidfVectorizer] = None
    scaler: Optional[StandardScaler] = None
    price_median_: Optional[float] = None
    avg_rating_median_: Optional[float] = None
    rating_number_median_: Optional[float] = None

    @property
    def total_dim(self) -> int:
        return self.tfidf_dim + self.num_dim

    # ------------------------------------------------------------------ fit
    def fit(self, df: pd.DataFrame) -> "MetaEncoder":
        # 1. TF-IDF on features_text + categories_text concatenated. categories
        # gives the high-level product class (e.g. 'Dresses > Casual'), which
        # is a strong prior for which aspects matter.
        feat_text = _safe_text(df.get(FEATURES_TEXT_COL, pd.Series([""] * len(df))))
        cat_text = _safe_text(df.get(CATEGORIES_TEXT_COL, pd.Series([""] * len(df))))
        combined_text = (feat_text + " " + cat_text).str.strip()

        self.tfidf = TfidfVectorizer(
            max_features=self.tfidf_dim,
            ngram_range=(1, 2),
            min_df=2,
            max_df=0.95,
            stop_words="english",
            sublinear_tf=True,
        )
        # If the corpus is too small to have any vocab, TF-IDF will raise.
        # Fall back to a single zero-dim by fitting on a synthetic vocab.
        try:
            self.tfidf.fit(combined_text)
        except ValueError:
            logger.warning("Meta TF-IDF found no usable vocabulary; using zero features.")
            self.tfidf = TfidfVectorizer(max_features=1)
            self.tfidf.fit(["placeholder"])

        # 2. Numeric features: price, avg_rating, log(rating_number+1), price_missing_flag
        self.price_median_ = float(_safe_numeric(df.get(PRICE_COL, pd.Series([np.nan]))).median())
        if not np.isfinite(self.price_median_):
            self.price_median_ = 0.0
        self.avg_rating_median_ = float(_safe_numeric(df.get(AVG_RATING_COL, pd.Series([np.nan]))).median())
        if not np.isfinite(self.avg_rating_median_):
            self.avg_rating_median_ = 4.0
        self.rating_number_median_ = float(_safe_numeric(df.get(RATING_NUMBER_COL, pd.Series([np.nan]))).median())
        if not np.isfinite(self.rating_number_median_):
            self.rating_number_median_ = 0.0

        num_mat = self._build_numeric(df)
        self.scaler = StandardScaler()
        self.scaler.fit(num_mat)
        return self

    # -------------------------------------------------------------- transform
    def transform(self, df: pd.DataFrame) -> np.ndarray:
        if self.tfidf is None or self.scaler is None:
            raise RuntimeError("MetaEncoder.fit() must be called before transform().")

        feat_text = _safe_text(df.get(FEATURES_TEXT_COL, pd.Series([""] * len(df))))
        cat_text = _safe_text(df.get(CATEGORIES_TEXT_COL, pd.Series([""] * len(df))))
        combined_text = (feat_text + " " + cat_text).str.strip()

        text_mat = self.tfidf.transform(combined_text).toarray().astype(np.float32)
        # Pad / truncate to declared tfidf_dim so the model input size is stable.
        if text_mat.shape[1] < self.tfidf_dim:
            pad = np.zeros((text_mat.shape[0], self.tfidf_dim - text_mat.shape[1]), dtype=np.float32)
            text_mat = np.hstack([text_mat, pad])
        elif text_mat.shape[1] > self.tfidf_dim:
            text_mat = text_mat[:, : self.tfidf_dim]

        num_mat = self.scaler.transform(self._build_numeric(df)).astype(np.float32)

        return np.hstack([text_mat, num_mat]).astype(np.float32)

    # ------------------------------------------------------------------ helpers
    def _build_numeric(self, df: pd.DataFrame) -> np.ndarray:
        price = _safe_numeric(df.get(PRICE_COL, pd.Series([np.nan] * len(df))))
        price_missing = price.isna().astype(np.float32).values
        price = price.fillna(self.price_median_).astype(np.float32).values

        avg = _safe_numeric(df.get(AVG_RATING_COL, pd.Series([np.nan] * len(df))))
        avg = avg.fillna(self.avg_rating_median_).astype(np.float32).values

        rnum = _safe_numeric(df.get(RATING_NUMBER_COL, pd.Series([np.nan] * len(df))))
        rnum = rnum.fillna(self.rating_number_median_).astype(np.float32).values
        log_rnum = np.log1p(np.maximum(rnum, 0.0))

        return np.vstack([price, avg, log_rnum, price_missing]).T  # (N, 4)

    # ------------------------------------------------------------------ io
    def save(self, path: Path = None) -> Path:
        if path is None:
            path = cfg.META_ENCODER_PATH
        path = Path(path)
        path.parent.mkdir(parents=True, exist_ok=True)
        with open(path, "wb") as f:
            pickle.dump(self, f)
        logger.info("Saved MetaEncoder to %s", path)
        return path

    @classmethod
    def load(cls, path: Path = None) -> "MetaEncoder":
        if path is None:
            path = cfg.META_ENCODER_PATH
        with open(path, "rb") as f:
            obj = pickle.load(f)
        if not isinstance(obj, cls):
            raise TypeError(f"Loaded object is not a MetaEncoder: {type(obj)}")
        return obj


def fit_and_save(train_df: pd.DataFrame, path: Path = None) -> MetaEncoder:
    enc = MetaEncoder().fit(train_df)
    enc.save(path)
    return enc