src/models/train.py

import os
import warnings
# Suppress deprecation and future warnings from third-party libraries (like MLflow and statsmodels)
warnings.filterwarnings("ignore", category=FutureWarning)
warnings.filterwarnings("ignore")

import numpy as np
import pandas as pd
import mlflow
import mlflow.xgboost
import mlflow.statsmodels
import xgboost as xgb
from statsmodels.tsa.arima.model import ARIMA
from sklearn.metrics import mean_squared_error, mean_absolute_error
import logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

COMMODITY_MAP = {
    "Samba": 0,
    "Kekulu": 1,
    "Big Onion": 2,
    "Potato": 3,
    "Dried Chilli": 4,
    "Coconut": 5
}

def load_real_data():
    """Loads cleaned daily real price data."""
    csv_path = "data/processed/clean_prices.csv"
    if not os.path.exists(csv_path):
        raise FileNotFoundError(f"Clean data not found at {csv_path}. Please run clean_data.py first.")
    
    df = pd.read_csv(csv_path)
    df['Date'] = pd.to_datetime(df['Date'])
    df.set_index('Date', inplace=True)
    return df

def train_arima(train, test):
    """Trains an ARIMA baseline model and logs to MLflow."""
    logging.info("Training ARIMA model...")
    with mlflow.start_run(run_name="ARIMA_Baseline"):
        # Define hyperparams dynamically to prevent crashing on short data
        ar_order = min(5, max(1, len(train) - 2))
        order = (ar_order, 1, 0)
        mlflow.log_param("order", str(order))
        
        # Train
        model = ARIMA(train, order=order)
        fitted_model = model.fit()
        
        # Predict
        predictions = fitted_model.forecast(steps=len(test))
        
        # Metrics
        rmse = np.sqrt(mean_squared_error(test, predictions))
        mae = mean_absolute_error(test, predictions)
        
        mlflow.log_metric("rmse", rmse)
        mlflow.log_metric("mae", mae)
        
        # Log model artifact
        mlflow.statsmodels.log_model(fitted_model, name="arima_model")
        logging.info(f"ARIMA RMSE: {rmse:.2f}, MAE: {mae:.2f}")

def train_linear_regression(X_train, X_test, y_train, y_test):
    """Trains a baseline Linear Regression model on lag features for comparison and logs to MLflow."""
    logging.info("Training Linear Regression baseline...")
    from sklearn.linear_model import LinearRegression
    with mlflow.start_run(run_name="LinearRegression_Baseline"):
        # Fit Linear Regression model
        model = LinearRegression()
        model.fit(X_train, y_train)
        
        # Predict on test set
        predictions = model.predict(X_test)
        
        # Calculate evaluation metrics
        rmse = np.sqrt(mean_squared_error(y_test, predictions))
        mae = mean_absolute_error(y_test, predictions)
        
        # Log metrics to MLflow
        mlflow.log_metric("rmse", rmse)
        mlflow.log_metric("mae", mae)
        
        # Log model artifact
        mlflow.sklearn.log_model(model, name="linear_regression_model")
        logging.info(f"Linear Regression RMSE: {rmse:.2f}, MAE: {mae:.2f}")

def train_xgboost(df):
    """Trains an XGBoost candidate model with lag features and logs to MLflow."""
    logging.info("Training XGBoost model...")
    
    # Feature Engineering (Dynamic lag based on data size)
    min_days = df['Commodity'].value_counts().min()
    max_lags = max(1, min(7, min_days - 3))
    if min_days <= 3:
        logging.warning("Data volume extremely low. Setting lag to 1. XGBoost model might not be robust.")
        max_lags = 1

    df_list = []
    for commodity in df['Commodity'].unique():
        df_c = df[df['Commodity'] == commodity].copy()
        for i in range(1, max_lags + 1):
            df_c[f'Lag_{i}'] = df_c['Price'].shift(i)
        df_list.append(df_c)
        
    df_xgb = pd.concat(df_list)
    df_xgb.dropna(inplace=True)
    
    if len(df_xgb) < 2:
        logging.error("Not enough data to train XGBoost after lagging. Please gather more days of data.")
        return

    # Add Commodity_ID
    df_xgb['Commodity_ID'] = df_xgb['Commodity'].map(COMMODITY_MAP)
    
    feature_cols = ['Commodity_ID'] + [f'Lag_{i}' for i in range(1, max_lags + 1)]
    X = df_xgb[feature_cols]
    y = df_xgb['Price']
    
    # Train-test split (Dynamic based on data size)
    test_size = max(1, int(len(X) * 0.2))
    train_size = len(X) - test_size
    X_train, X_test = X.iloc[:train_size], X.iloc[train_size:]
    y_train, y_test = y.iloc[:train_size], y.iloc[train_size:]
    
    # Train Linear Regression baseline for academic performance comparison
    train_linear_regression(X_train, X_test, y_train, y_test)
    
    with mlflow.start_run(run_name="XGBoost_Candidate"):
        params = {"objective": "reg:squarederror", "n_estimators": 100, "learning_rate": 0.1, "max_depth": 5}
        mlflow.log_params(params)
        
        model = xgb.XGBRegressor(**params)
        model.fit(X_train, y_train)
        
        predictions = model.predict(X_test)
        
        rmse = np.sqrt(mean_squared_error(y_test, predictions))
        mae = mean_absolute_error(y_test, predictions)
        
        mlflow.log_metric("rmse", rmse)
        mlflow.log_metric("mae", mae)
        
        # Log model artifact
        mlflow.xgboost.log_model(model, name="xgboost_model")
        logging.info(f"XGBoost RMSE: {rmse:.2f}, MAE: {mae:.2f}")

if __name__ == "__main__":
    # Use environment variable for MLflow tracking URI, or default to local folder
    tracking_uri = os.environ.get("MLFLOW_TRACKING_URI", "file:./mlruns")
    mlflow.set_tracking_uri(tracking_uri)
    
    # If using local tracking, ensure the mlruns directory exists
    if tracking_uri.startswith("file:"):
        os.makedirs(tracking_uri.replace("file:", ""), exist_ok=True)
        
    # Configure DagsHub credentials automatically if URI is provided
    dagshub_token = os.environ.get("DAGSHUB_USER_TOKEN")
    if dagshub_token and "dagshub.com" in tracking_uri:
        # Extract username from URI (e.g. https://dagshub.com/username/repo.mlflow)
        username = tracking_uri.split('/')[-2]
        os.environ["MLFLOW_TRACKING_USERNAME"] = username
        os.environ["MLFLOW_TRACKING_PASSWORD"] = dagshub_token
        logging.info(f"Configured DagsHub MLflow remote for user: {username}")
    
    mlflow.set_experiment("SL_Commodity_Forecasting")
    
    df = load_real_data()
    
    # Split for ARIMA (Baseline uses only Samba)
    df_samba = df[df['Commodity'] == 'Samba']
    if len(df_samba) < 3:
        logging.error("Not enough data for Samba to train ARIMA. Exiting.")
    else:
        test_size = max(1, int(len(df_samba) * 0.2))
        train_size = len(df_samba) - test_size
        train, test = df_samba.iloc[:train_size]['Price'], df_samba.iloc[train_size:]['Price']
    
    train_arima(train, test)
    train_xgboost(df)
    
    logging.info("Training complete. Run 'mlflow ui' in the terminal to view your experiment dashboards!")