ml_engine.py

import pandas as pd
import numpy as np
import xgboost as xgb
from sklearn.ensemble import IsolationForest
from sklearn.preprocessing import StandardScaler
import shap

class MLEngine:
    def __init__(self, data_path='data/interval_data.csv'):
        self.data_path = data_path
        self.df = None
        self.forecast_model = None
        self.anomaly_model = None
        self.scaler = StandardScaler()
        self.explainer = None
        
    def load_data(self):
        try:
            self.df = pd.read_csv(self.data_path)
            self.df['Timestamp'] = pd.to_datetime(self.df['Timestamp'])
            self.df = self.df.sort_values(by=['Meter_ID_Hash', 'Timestamp'])
        except FileNotFoundError:
            print(f"Data not found at {self.data_path}. Run data_generator.py first.")
            
    def prepare_forecast_features(self, df):
        """Creates time-based features for forecasting."""
        df_feat = df.copy()
        df_feat['Hour'] = df_feat['Timestamp'].dt.hour
        df_feat['Minute'] = df_feat['Timestamp'].dt.minute
        df_feat['DayOfWeek'] = df_feat['Timestamp'].dt.dayofweek
        # Lag features
        df_feat['Lag_1'] = df_feat.groupby('Meter_ID_Hash')['Consumption_kWh'].shift(1)
        df_feat['Lag_4'] = df_feat.groupby('Meter_ID_Hash')['Consumption_kWh'].shift(4) # 1 hour ago
        df_feat = df_feat.dropna()
        return df_feat
        
    def train_forecast_model(self):
        """Trains an XGBoost model for demand forecasting (Quantile Regression)."""
        if self.df is None:
            self.load_data()
            
        print("Training Forecast Model (XGBoost)...")
        df_feat = self.prepare_forecast_features(self.df)
        
        features = ['Hour', 'Minute', 'DayOfWeek', 'Lag_1', 'Lag_4']
        X = df_feat[features]
        y = df_feat['Consumption_kWh']
        
        # We use objective='reg:quantileerror' with quantile_alpha=0.5 (Median) and upper/lower if needed.
        # For simplicity, we'll train a standard regressor to predict the expected demand.
        self.forecast_model = xgb.XGBRegressor(
            objective='reg:squarederror', 
            n_estimators=50,
            max_depth=4,
            learning_rate=0.1
        )
        self.forecast_model.fit(X, y)
        print("Forecast Model trained.")
        
    def train_anomaly_model(self):
        """Trains an Isolation Forest ensemble for anomaly detection."""
        if self.df is None:
            self.load_data()
            
        print("Training Anomaly Model (Isolation Forest)...")
        # Features for anomaly detection: Consumption and Grid_Stress
        features = ['Consumption_kWh', 'Grid_Stress']
        X = self.df[features].dropna()
        
        X_scaled = self.scaler.fit_transform(X)
        
        self.anomaly_model = IsolationForest(
            n_estimators=100, 
            contamination=0.05, # Expect 5% anomalies
            random_state=42
        )
        self.anomaly_model.fit(X_scaled)
        
        # Initialize SHAP explainer
        # TreeExplainer works for IsolationForest in SHAP, but can be slow. We'll use a simpler approximation if needed, 
        # or just train it so it's ready.
        try:
            # SHAP works natively with sklearn IF, using TreeExplainer
            self.explainer = shap.TreeExplainer(self.anomaly_model)
        except Exception as e:
            print(f"SHAP TreeExplainer failed to initialize: {e}")
            self.explainer = None
            
        print("Anomaly Model trained.")
        
    def get_forecast(self, meter_id_hash=None, steps=96):
        """Returns 24h forecast (96 x 15min) for a given meter or aggregated."""
        if self.forecast_model is None:
            self.train_forecast_model()
            
        # Simplified forecast: predict for the next 24h based on the last available data
        if meter_id_hash:
            meter_data = self.df[self.df['Meter_ID_Hash'] == meter_id_hash].copy()
        else:
            # Random meter if none provided
            meter_id_hash = self.df['Meter_ID_Hash'].iloc[0]
            meter_data = self.df[self.df['Meter_ID_Hash'] == meter_id_hash].copy()
            
        if len(meter_data) == 0:
            return pd.DataFrame()
            
        last_record = meter_data.iloc[-1]
        last_time = last_record['Timestamp']
        
        future_times = [last_time + pd.Timedelta(minutes=15 * i) for i in range(1, steps + 1)]
        future_df = pd.DataFrame({'Timestamp': future_times})
        future_df['Hour'] = future_df['Timestamp'].dt.hour
        future_df['Minute'] = future_df['Timestamp'].dt.minute
        future_df['DayOfWeek'] = future_df['Timestamp'].dt.dayofweek
        
        # Naive lag assumption for future (use the last known value)
        future_df['Lag_1'] = last_record['Consumption_kWh']
        future_df['Lag_4'] = meter_data.iloc[-4]['Consumption_kWh'] if len(meter_data) >=4 else last_record['Consumption_kWh']
        
        features = ['Hour', 'Minute', 'DayOfWeek', 'Lag_1', 'Lag_4']
        preds = self.forecast_model.predict(future_df[features])
        future_df['Forecast_kWh'] = preds
        future_df['Meter_ID_Hash'] = meter_id_hash
        
        return future_df[['Timestamp', 'Meter_ID_Hash', 'Forecast_kWh']]
        
    def detect_anomalies(self, recent_data_only=True):
        """Runs the anomaly detector and returns flagged meters along with SHAP values."""
        if self.anomaly_model is None:
            self.train_anomaly_model()
            
        if recent_data_only:
            # Let's say "recent" is the last 4 hours (16 records) of each meter
            df_eval = self.df.groupby('Meter_ID_Hash').tail(16).copy()
        else:
            df_eval = self.df.copy()
            
        features = ['Consumption_kWh', 'Grid_Stress']
        X = df_eval[features].dropna()
        if len(X) == 0:
            return []
            
        X_scaled = self.scaler.transform(X)
        
        preds = self.anomaly_model.predict(X_scaled)
        # IsolationForest: -1 for anomaly, 1 for normal
        df_eval['Is_Flagged'] = np.where(preds == -1, 1, 0)
        
        anomalies = df_eval[df_eval['Is_Flagged'] == 1].copy()
        
        # Calculate SHAP for anomalies to generate evidence
        results = []
        if len(anomalies) > 0 and self.explainer is not None:
            X_anom = anomalies[features]
            X_anom_scaled = self.scaler.transform(X_anom)
            shap_values = self.explainer.shap_values(X_anom_scaled)
            
            # Create a summary for each anomaly
            for i, (_, row) in enumerate(anomalies.iterrows()):
                evidence = {
                    'Timestamp': str(row['Timestamp']),
                    'Meter_ID_Hash': row['Meter_ID_Hash'],
                    'Consumption_kWh': float(row['Consumption_kWh']),
                    'Grid_Stress': float(row['Grid_Stress']),
                    'SHAP_Consumption': float(shap_values[i][0]),
                    'SHAP_Grid_Stress': float(shap_values[i][1])
                }
                results.append(evidence)
                
        # Deduplicate by meter to show only the latest anomaly per meter
        unique_meters = {}
        for r in reversed(results): # reversed so the latest timestamp overwrites
            unique_meters[r['Meter_ID_Hash']] = r
            
        return list(unique_meters.values())

if __name__ == "__main__":
    engine = MLEngine()
    engine.load_data()
    engine.train_forecast_model()
    engine.train_anomaly_model()
    anoms = engine.detect_anomalies()
    print(f"Detected {len(anoms)} recent anomalies.")