"""
Weight predictor classes for food delivery platforms.
"""
import numpy as np


class WeightPredictor:
    """Per-item-type weight predictor using TF-IDF + Ridge regression."""
    def __init__(self, tfidf, model):
        self.tfidf = tfidf
        self.model = model

    def predict(self, texts):
        """Predict weights for a list of texts."""
        X = self.tfidf.transform(texts)
        return np.expm1(self.model.predict(X))

    def predict_single(self, text):
        """Predict weight for a single text."""
        return self.predict([text])[0]


class UnifiedWeightPredictor:
    """Unified predictor that routes to per-type models."""
    def __init__(self, predictors, default_type="grocery"):
        self.predictors = predictors
        self.default_type = default_type

    def predict(self, text, item_type=None):
        """
        Predict weight from text description.
        item_type should be one of: menu_item, grocery, non_food
        """
        if item_type is None:
            if text.startswith("[MENU_ITEM]"):
                item_type = "menu_item"
            elif text.startswith("[GROCERY]"):
                item_type = "grocery"
            elif text.startswith("[NON_FOOD]"):
                item_type = "non_food"
            else:
                item_type = self.default_type
        predictor = self.predictors.get(item_type, self.predictors.get(self.default_type))
        return predictor.predict([text])[0]