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_FEATURE_TFIDF_DIM)     # product attributes
    | TF-IDF on categories_text (META_CATEGORY_TFIDF_DIM)   # product taxonomy
    | 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."""

    feature_tfidf_dim: int = cfg.META_FEATURE_TFIDF_DIM
    category_tfidf_dim: int = cfg.META_CATEGORY_TFIDF_DIM
    num_dim: int = cfg.META_NUM_DIM  # price, avg_rating, log_rating_number, price_missing_flag

    # Filled by fit()
    feature_tfidf: Optional[TfidfVectorizer] = None
    category_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 tfidf_dim(self) -> int:
        return self.feature_tfidf_dim + self.category_tfidf_dim

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

    # ------------------------------------------------------------------ fit
    def fit(self, df: pd.DataFrame) -> "MetaEncoder":
        # 1. Fit text metadata as separate semantic sources. Keeping features
        # and categories apart gives cross-attention tokens real meaning.
        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))))

        self.feature_tfidf = TfidfVectorizer(
            max_features=self.feature_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.feature_tfidf.fit(feat_text.str.strip())
        except ValueError:
            logger.warning("Meta TF-IDF found no usable vocabulary; using zero features.")
            self.feature_tfidf = TfidfVectorizer(max_features=1)
            self.feature_tfidf.fit(["placeholder"])

        self.category_tfidf = TfidfVectorizer(
            max_features=self.category_tfidf_dim,
            ngram_range=(1, 2),
            min_df=2,
            max_df=0.98,
            stop_words="english",
            sublinear_tf=True,
        )
        try:
            self.category_tfidf.fit(cat_text.str.strip())
        except ValueError:
            logger.warning("Category TF-IDF found no usable vocabulary; using zero features.")
            self.category_tfidf = TfidfVectorizer(max_features=1)
            self.category_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.feature_tfidf is None or self.scaler is None:
            raise RuntimeError("MetaEncoder.fit() must be called before transform().")

        structured = self.transform_structured(df)
        return np.hstack([
            structured["features"],
            structured["categories"],
            structured["numeric"],
        ]).astype(np.float32)

    def transform_structured(self, df: pd.DataFrame) -> dict:
        """Return source-separated metadata matrices.

        Keys map directly to semantic meta tokens used by the upgraded model:
        features -> product attribute text, categories -> taxonomy text,
        numeric -> price/rating signals.
        """
        if self.feature_tfidf is None or self.category_tfidf is None or self.scaler is None:
            raise RuntimeError("MetaEncoder.fit() must be called before transform_structured().")

        feat_text = _safe_text(df.get(FEATURES_TEXT_COL, pd.Series([""] * len(df)))).str.strip()
        cat_text = _safe_text(df.get(CATEGORIES_TEXT_COL, pd.Series([""] * len(df)))).str.strip()
        feat_mat = self.feature_tfidf.transform(feat_text).toarray().astype(np.float32)
        cat_mat = self.category_tfidf.transform(cat_text).toarray().astype(np.float32)
        feat_mat = self._pad_or_truncate(feat_mat, self.feature_tfidf_dim)
        cat_mat = self._pad_or_truncate(cat_mat, self.category_tfidf_dim)
        num_mat = self.scaler.transform(self._build_numeric(df)).astype(np.float32)
        return {"features": feat_mat, "categories": cat_mat, "numeric": num_mat}

    @staticmethod
    def _pad_or_truncate(mat: np.ndarray, dim: int) -> np.ndarray:
        if mat.shape[1] < dim:
            pad = np.zeros((mat.shape[0], dim - mat.shape[1]), dtype=np.float32)
            return np.hstack([mat, pad])
        if mat.shape[1] > dim:
            return mat[:, :dim]
        return mat

    # ------------------------------------------------------------------ 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