import pandas as pd
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.naive_bayes import MultinomialNB
from sklearn.pipeline import make_pipeline
import pickle
import os

# 1. Dataset of common foods and their typical shelf life (in days)
#    This acts as our "Knowledge Base"
data = [
    # Dairy
    ("Milk", 7), ("Yogurt", 14), ("Cheese", 30), ("Butter", 60), ("Cream", 10),
    # Vegetables
    ("Spinach", 5), ("Lettuce", 5), ("Tomato", 7), ("Carrot", 21), ("Potato", 30), ("Onion", 30),
    # Fruits
    ("Apple", 21), ("Banana", 5), ("Orange", 14), ("Grapes", 7), ("Strawberry", 4),
    # Meat
    ("Chicken", 2), ("Beef", 3), ("Pork", 3), ("Fish", 2), ("Ham", 5),
    # Grains
    ("Bread", 5), ("Rice", 365), ("Pasta", 365), ("Cereal", 180),
    # Misc
    ("Eggs", 21), ("Juice", 10), ("Sauce", 90), ("Canned Beans", 700)
]

MODEL_PATH = "expiry_model.pkl"

def train_model():
    print("Training Expiry Prediction Model...")
    df = pd.DataFrame(data, columns=["item", "days"])
    
    # Simple NLP pipeline: Text -> Vector -> Classifier
    model = make_pipeline(TfidfVectorizer(), MultinomialNB())
    model.fit(df["item"], df["days"])
    
    with open(MODEL_PATH, "wb") as f:
        pickle.dump(model, f)
    print("Model saved to", MODEL_PATH)
    return model

def load_model():
    if os.path.exists(MODEL_PATH):
        with open(MODEL_PATH, "rb") as f:
            return pickle.load(f)
    else:
        return train_model()

def predict_days(item_name):
    model = load_model()
    try:
        # Predict
        predicted_days = model.predict([item_name])[0]
        return int(predicted_days)
    except Exception as e:
        print(f"Prediction Error: {e}")
        return 7 # Default fallback