model.py

"""
Machine Learning Model for Predictive Maintenance
Uses Random Forest Classifier for failure prediction
"""

import pandas as pd
import numpy as np
import pickle
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import (
    accuracy_score, precision_score, recall_score, f1_score,
    confusion_matrix, classification_report, roc_auc_score
)
import warnings
warnings.filterwarnings('ignore')

class PredictiveMaintenanceModel:
    def __init__(self):
        """Initialize the model"""
        self.model = RandomForestClassifier(
            n_estimators=100,
            max_depth=10,
            min_samples_split=5,
            min_samples_leaf=2,
            random_state=42,
            class_weight='balanced'  # Handle class imbalance
        )
        self.is_trained = False
        self.feature_importance_ = None
        
    def train(self, X_train, y_train):
        """Train the model"""
        print("Training Random Forest Classifier...")
        self.model.fit(X_train, y_train)
        self.is_trained = True
        
        # Get feature importance
        self.feature_importance_ = pd.DataFrame({
            'feature': X_train.columns,
            'importance': self.model.feature_importances_
        }).sort_values('importance', ascending=False)
        
        print("Model training complete!")
        return self
    
    def predict(self, X):
        """Make predictions"""
        if not self.is_trained:
            raise ValueError("Model must be trained first!")
        return self.model.predict(X)
    
    def predict_proba(self, X):
        """Get prediction probabilities"""
        if not self.is_trained:
            raise ValueError("Model must be trained first!")
        return self.model.predict_proba(X)
    
    def evaluate(self, X_test, y_test):
        """Evaluate the model"""
        if not self.is_trained:
            raise ValueError("Model must be trained first!")
        
        # Predictions
        y_pred = self.predict(X_test)
        y_pred_proba = self.predict_proba(X_test)[:, 1]
        
        # Metrics
        accuracy = accuracy_score(y_test, y_pred)
        precision = precision_score(y_test, y_pred, zero_division=0)
        recall = recall_score(y_test, y_pred, zero_division=0)
        f1 = f1_score(y_test, y_pred, zero_division=0)
        roc_auc = roc_auc_score(y_test, y_pred_proba)
        
        # Confusion matrix
        cm = confusion_matrix(y_test, y_pred)
        
        # Classification report
        report = classification_report(y_test, y_pred, zero_division=0)
        
        results = {
            'accuracy': accuracy,
            'precision': precision,
            'recall': recall,
            'f1_score': f1,
            'roc_auc': roc_auc,
            'confusion_matrix': cm,
            'classification_report': report,
            'y_pred': y_pred,
            'y_pred_proba': y_pred_proba
        }
        
        print("="*60)
        print("MODEL EVALUATION RESULTS")
        print("="*60)
        print(f"Accuracy:  {accuracy:.4f}")
        print(f"Precision: {precision:.4f}")
        print(f"Recall:    {recall:.4f}")
        print(f"F1-Score:  {f1:.4f}")
        print(f"ROC-AUC:   {roc_auc:.4f}")
        print("\nConfusion Matrix:")
        print(cm)
        print("\nClassification Report:")
        print(report)
        
        return results
    
    def predict_maintenance(self, X, tool_wear_values=None):
        """
        Predict maintenance needs and time to failure
        
        Args:
            X: Feature matrix
            tool_wear_values: Array of tool wear values (optional)
        
        Returns:
            DataFrame with predictions and maintenance recommendations
        """
        if not self.is_trained:
            raise ValueError("Model must be trained first!")
        
        # Get failure predictions
        failure_pred = self.predict(X)
        failure_proba = self.predict_proba(X)[:, 1]
        
        # Estimate time to failure based on tool wear
        # Average tool wear at failure is around 100-150 minutes based on EDA
        avg_tool_wear_at_failure = 120  # minutes
        
        if tool_wear_values is None:
            # Try to get tool wear from features if available
            if 'Tool wear [min]' in X.columns:
                tool_wear_values = X['Tool wear [min]'].values
            else:
                tool_wear_values = np.zeros(len(X))
        else:
            # Convert to numpy array if it's a list or scalar
            tool_wear_values = np.array(tool_wear_values)
            # If it's a scalar, make it an array
            if tool_wear_values.ndim == 0:
                tool_wear_values = np.array([tool_wear_values])
        
        # Ensure tool_wear_values is 1D array with same length as predictions
        if len(tool_wear_values) != len(failure_pred):
            # If single value provided, repeat it for all predictions
            if len(tool_wear_values) == 1:
                tool_wear_values = np.repeat(tool_wear_values, len(failure_pred))
            else:
                raise ValueError(f"tool_wear_values length ({len(tool_wear_values)}) doesn't match predictions length ({len(failure_pred)})")
        
        # Calculate estimated time to failure
        time_to_failure = np.maximum(0, avg_tool_wear_at_failure - tool_wear_values)
        
        # Determine maintenance urgency
        maintenance_status = []
        maintenance_urgency = []
        
        for i in range(len(failure_pred)):
            # Check if tool wear has exceeded the average failure threshold
            tool_wear_exceeded = tool_wear_values[i] >= avg_tool_wear_at_failure
            tool_wear_severely_exceeded = tool_wear_values[i] >= (avg_tool_wear_at_failure * 1.5)  # 150% of threshold (180 min)
            
            # Combined risk assessment: consider both ML prediction AND tool wear
            if failure_pred[i] == 1:
                # Model predicts failure - highest priority
                maintenance_status.append("IMMEDIATE MAINTENANCE REQUIRED")
                maintenance_urgency.append("CRITICAL")
            elif tool_wear_severely_exceeded and failure_proba[i] > 0.2:
                # Severely exceeded tool wear AND significant failure probability
                maintenance_status.append("IMMEDIATE MAINTENANCE REQUIRED - Tool wear severely exceeded threshold")
                maintenance_urgency.append("CRITICAL")
            elif tool_wear_severely_exceeded:
                # Severely exceeded tool wear but low failure probability - still HIGH priority
                maintenance_status.append("URGENT MAINTENANCE NEEDED - Tool wear severely exceeded (monitor closely)")
                maintenance_urgency.append("HIGH")
            elif tool_wear_exceeded and failure_proba[i] > 0.3:
                # Tool wear exceeded AND moderate failure probability
                maintenance_status.append("IMMEDIATE MAINTENANCE REQUIRED - Tool wear exceeded threshold")
                maintenance_urgency.append("CRITICAL")
            elif tool_wear_exceeded:
                # Tool wear exceeded but low failure probability - HIGH priority
                maintenance_status.append("URGENT MAINTENANCE NEEDED - Tool wear exceeded threshold")
                maintenance_urgency.append("HIGH")
            elif failure_proba[i] > 0.7:
                # High failure probability regardless of tool wear
                maintenance_status.append("Maintenance needed soon")
                maintenance_urgency.append("HIGH")
            elif failure_proba[i] > 0.4:
                # Moderate failure probability
                maintenance_status.append("Schedule maintenance")
                maintenance_urgency.append("MEDIUM")
            elif time_to_failure[i] < 20:
                # Less than 20 minutes remaining
                maintenance_status.append("Monitor closely - Maintenance needed within 20 minutes")
                maintenance_urgency.append("MEDIUM")
            elif time_to_failure[i] < 60:
                # Less than 1 hour remaining
                maintenance_status.append("Monitor closely - Maintenance needed within 1 hour")
                maintenance_urgency.append("MEDIUM")
            else:
                # Low risk
                maintenance_status.append("No immediate maintenance needed")
                maintenance_urgency.append("LOW")
        
        results_df = pd.DataFrame({
            'Failure_Predicted': failure_pred,
            'Failure_Probability': failure_proba,
            'Time_to_Failure_Minutes': time_to_failure,
            'Maintenance_Status': maintenance_status,
            'Maintenance_Urgency': maintenance_urgency
        })
        
        return results_df
    
    
    def get_feature_importance(self):
        """Get feature importance"""
        if self.feature_importance_ is None:
            raise ValueError("Model must be trained first!")
        return self.feature_importance_
    
    def save_model(self, filepath='predictive_maintenance_model.pkl'):
        """Save the trained model"""
        if not self.is_trained:
            raise ValueError("Model must be trained first!")
        
        with open(filepath, 'wb') as f:
            pickle.dump(self.model, f)
        print(f"Model saved to {filepath}")
    
    def load_model(self, filepath='predictive_maintenance_model.pkl'):
        """Load a trained model"""
        with open(filepath, 'rb') as f:
            self.model = pickle.load(f)
        self.is_trained = True
        print(f"Model loaded from {filepath}")