Machine Learning

# src/train.py
import pandas as pd
import numpy as np
from pathlib import Path
from sklearn.model_selection import TimeSeriesSplit
from sklearn.metrics import mean_absolute_error
import lightgbm as lgb
import joblib
import argparse

def create_features(df):
    df = df.copy()
    df['day_of_week'] = df['date'].dt.dayofweek
    df['week_of_year'] = df['date'].dt.isocalendar().week.astype(int)
    df['month'] = df['date'].dt.month

    # lag features
    for lag in [1, 7, 14, 30]:
        df[f'lag_{lag}'] = df['usage_qty'].shift(lag)

    # rolling features (shifted to avoid leakage)
    df['rmean_7'] = df['usage_qty'].rolling(7, min_periods=1).mean().shift(1)
    df['rmean_30'] = df['usage_qty'].rolling(30, min_periods=1).mean().shift(1)

    # inventory metric
    df['days_of_stock'] = df['on_hand_qty'] / (df['rmean_30'].replace(0, np.nan))
    df['days_of_stock'] = df['days_of_stock'].fillna(0)

    df = df.fillna(0)
    return df

def train_single_sku(df_sku, model_out_dir: Path):
    df_sku = df_sku.sort_values('date').reset_index(drop=True)
    feat = create_features(df_sku)
    # Create single-step next-day target
    feat['target_1'] = feat['usage_qty'].shift(-1)
    feat = feat.dropna(subset=['target_1']).reset_index(drop=True)

    ignore_cols = ['date', 'sku_id', 'usage_qty', 'target_1']
    feature_cols = [c for c in feat.columns if c not in ignore_cols]

    X = feat[feature_cols]
    y = feat['target_1']

    tscv = TimeSeriesSplit(n_splits=5)
    models = []
    maes = []
    fold = 0
    for train_idx, val_idx in tscv.split(X):
        fold += 1
        X_train, X_val = X.iloc[train_idx], X.iloc[val_idx]
        y_train, y_val = y.iloc[train_idx], y.iloc[val_idx]

        dtrain = lgb.Dataset(X_train, label=y_train)
        dval = lgb.Dataset(X_val, label=y_val)

        params = {
            'objective': 'regression',
            'metric': 'l1',
            'verbosity': -1,
            'boosting_type': 'gbdt',
            'learning_rate': 0.05,
            'num_leaves': 31,
        }
        model = lgb.train(params, dtrain, valid_sets=[dval], early_stopping_rounds=50, num_boost_round=2000, verbose_eval=False)
        y_pred = model.predict(X_val)
        mae = mean_absolute_error(y_val, y_pred)
        print(f"Fold {fold} MAE: {mae:.4f}")
        models.append(model)
        maes.append(mae)

    # Save ensemble (list of models) and feature columns
    model_out_dir.mkdir(parents=True, exist_ok=True)
    sku = df_sku['sku_id'].iloc[0]
    joblib.dump(models, model_out_dir / f"models_{sku}.pkl")
    joblib.dump(feature_cols, model_out_dir / f"feature_cols_{sku}.pkl")
    print("Saved models and feature list to", model_out_dir)
    print("Average MAE:", np.mean(maes))
    return model_out_dir / f"models_{sku}.pkl", model_out_dir / f"feature_cols_{sku}.pkl"

def main(args):
    data_path = Path(args.data)
    df = pd.read_csv(data_path, parse_dates=['date'])
    out_dir = Path(args.out)
    out_dir.mkdir(parents=True, exist_ok=True)

    # For POC keep first SKU only; loop across SKUs for full run
    sku = df['sku_id'].unique()[0]
    df_sku = df[df['sku_id'] == sku].copy()
    train_single_sku(df_sku, out_dir)

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--data", default="data/usage_example.csv")
    parser.add_argument("--out", default="models")
    args = parser.parse_args()
    main(args)
print("Saved:", OUT / "usage_example.csv")