training/train_global.py

"""
Train a Global Multivariate Machine Learning model for NutriLoop AI.
Predicts quantity based on restaurant_id, item_name, region (lat/lon), and time elements.
Replaces the old univariate Prophet models.
"""
import json
import os
import sys
from datetime import datetime, timedelta
from pathlib import Path

import holidays
import joblib
import pandas as pd
from sklearn.pipeline import Pipeline
from sklearn.compose import ColumnTransformer
from sklearn.preprocessing import OrdinalEncoder
from sklearn.ensemble import HistGradientBoostingRegressor
from sklearn.metrics import mean_absolute_error
from supabase import create_client
from dotenv import load_dotenv

# Ensure models directory exists
MODELS_DIR = Path(__file__).parent.parent / "models"
MODELS_DIR.mkdir(exist_ok=True)

# Allow direct execution from the project root and load local env vars.
project_root = Path(__file__).resolve().parent.parent
if str(project_root) not in sys.path:
    sys.path.insert(0, str(project_root))

load_dotenv()


def get_india_holidays(years=range(2020, 2028)) -> set:
    """Return a set of holiday dates for India."""
    in_holidays = holidays.India(years=years)
    return set(in_holidays.keys())


def extract_features(df: pd.DataFrame, holiday_dates: set) -> pd.DataFrame:
    """Extract temporal features from a DataFrame containing 'sale_date'."""
    df = df.copy()
    df['sale_date'] = pd.to_datetime(df['sale_date'])
    df['day_of_week'] = df['sale_date'].dt.dayofweek
    df['day_of_year'] = df['sale_date'].dt.dayofyear
    df['month'] = df['sale_date'].dt.month
    df['year'] = df['sale_date'].dt.year
    df['is_holiday'] = df['sale_date'].dt.date.isin(holiday_dates).astype(int)
    return df


def train_global_model():
    """
    Main training loop for the Global Multivariate Model.
    """
    print("[NutriLoop] Starting Global Model training")

    supabase_url = os.environ.get("SUPABASE_URL")
    supabase_key = os.environ.get("SUPABASE_KEY")

    if not supabase_url or not supabase_key:
        print("[NutriLoop] ERROR: SUPABASE_URL and SUPABASE_KEY must be set")
        return False, {}

    client = create_client(supabase_url, supabase_key)

    print("[NutriLoop] Fetching sales_logs from Supabase")
    try:
        response_sales = client.table("sales_logs").select("*").execute()
        response_restaurants = client.table("restaurants").select("*").execute()
    except Exception as e:
        print(f"[NutriLoop] ERROR: Could not load data from Supabase: {e}")
        return False, {}

    sales_df = pd.DataFrame(response_sales.data)
    rests_df = pd.DataFrame(response_restaurants.data)

    if sales_df.empty:
        print("[NutriLoop] No data in sales_logs.")
        return False, {}

    print(f"[NutriLoop] Merging {len(sales_df)} sales logs with {len(rests_df)} restaurants")
    
    # Merge sales with restaurant geographic data
    # Fallback missing data handles appropriately
    if not rests_df.empty:
        # Avoid column conflicts, keep only what we need from restaurants
        rests_df = rests_df[['restaurant_id', 'latitude', 'longitude', 'cuisine_type', 'avg_daily_quantity']]
        df = pd.merge(sales_df, rests_df, on="restaurant_id", how="left")
    else:
        df = sales_df
        df['latitude'] = 0.0
        df['longitude'] = 0.0
        df['cuisine_type'] = "Unknown"
        df['avg_daily_quantity'] = 0.0

    df["quantity"] = pd.to_numeric(df["quantity"], errors="coerce").fillna(0)

    # Feature Engineering
    holiday_dates = get_india_holidays()
    df = extract_features(df, holiday_dates)
    
    # Fill any null latitude/longitudes with global averages just in case
    df['latitude'] = df['latitude'].fillna(0.0)
    df['longitude'] = df['longitude'].fillna(0.0)
    df['cuisine_type'] = df['cuisine_type'].fillna('Unknown')
    df['avg_daily_quantity'] = df['avg_daily_quantity'].fillna(0.0)
    
    # Sort chronologically for valid temporal holdout
    df.sort_values("sale_date", inplace=True)
    
    # Establish a 14-day holdout validation set
    holdout_days = 14
    cutoff_date = df["sale_date"].max() - timedelta(days=holdout_days)
    
    train_df = df[df["sale_date"] < cutoff_date].copy()
    holdout_df = df[df["sale_date"] >= cutoff_date].copy()
    
    if len(train_df) < 50:
        print("[NutriLoop] Insufficient training data.")
        return False, {}

    print(f"[NutriLoop] Training on : {len(train_df)} rows, Validation: {len(holdout_df)} rows")

    # Define the Model Pipeline
    categorical_features = ["restaurant_id", "item_name", "cuisine_type"]
    numeric_features = ["day_of_week", "day_of_year", "month", "year", "is_holiday", "latitude", "longitude", "avg_daily_quantity"]
    all_features = categorical_features + numeric_features

    X_train = train_df[all_features]
    y_train = train_df["quantity"]
    
    X_val = holdout_df[all_features]
    y_val = holdout_df["quantity"]

    # Preprocessor encodes strings to integers for HGBR native categorical support
    preprocessor = ColumnTransformer(
        transformers=[
            (
                "cat", 
                OrdinalEncoder(handle_unknown="use_encoded_value", unknown_value=-1), 
                categorical_features
            )
        ],
        remainder="passthrough", # Keep numeric as is
        verbose_feature_names_out=False
    )
    
    model = Pipeline([
        ("preprocessor", preprocessor),
        ("regressor", HistGradientBoostingRegressor(
            max_iter=150,
            categorical_features=list(range(len(categorical_features))),
            loss='absolute_error', # MAE focus
            random_state=42
        ))
    ])

    print("[NutriLoop] Starting Fit...")
    model.fit(X_train, y_train)

    # Evaluation
    mae = 0.0
    if len(X_val) > 0:
        y_pred = model.predict(X_val)
        mae = float(mean_absolute_error(y_val, y_pred))

    print(f"[NutriLoop] Model Trained! Global Validation MAE: {mae:.2f}")

    # Artifact generation
    model_path = MODELS_DIR / "global_model.pkl"
    joblib.dump(model, model_path)
    print(f"[NutriLoop] Saved Model: {model_path}")

    model_registry = {
        "global_model": {
            "trained_at": datetime.now().isoformat(),
            "mae": round(mae, 4),
            "rows_used": len(df),
            "features": all_features,
            "algorithms": "HistGradientBoostingRegressor"
        }
    }
    
    registry_path = MODELS_DIR / "model_registry.json"
    with open(registry_path, "w") as f:
        json.dump(model_registry, f, indent=2)

    # Log to Supabase
    try:
        client.table("retrain_log").insert({
            "model_version": datetime.now().isoformat(),
            "rows_used": len(df),
            "mae_score": mae,
            "status": "success",
        }).execute()
    except Exception as e:
        print(f"[NutriLoop] Failed to log to Supabase retrain_log: {e}")

    return True, model_registry


if __name__ == "__main__":
    success, registry = train_global_model()
    if success:
        print(f"\n[NutriLoop Summary] Global Multivariate Model Trained | Valid MAE: {registry['global_model']['mae']}")