predictor.py

"""
AI Predictor for System Health
Uses machine learning to predict potential system failures
"""

import pandas as pd
import numpy as np
from sklearn.ensemble import RandomForestClassifier
from sklearn.preprocessing import StandardScaler
import joblib
import os


class SystemHealthPredictor:
    def __init__(self, model_path='model/health_predictor.pkl'):
        self.model_path = model_path
        self.model = None
        self.scaler = StandardScaler()
        os.makedirs('model', exist_ok=True)

    def create_features(self, df):

        df['cpu_ma_5'] = df['cpu_percent'].rolling(window=5, min_periods=1).mean()
        df['memory_ma_5'] = df['memory_percent'].rolling(window=5, min_periods=1).mean()


        df['cpu_change'] = df['cpu_percent'].diff().fillna(0)
        df['memory_change'] = df['memory_percent'].diff().fillna(0)

        df['stress_score'] = (df['cpu_percent'] + df['memory_percent']) / 2

        return df

    def create_labels(self, df):

        at_risk = (
                (df['cpu_percent'] > 80) |
                (df['memory_percent'] > 85) |
                (df['disk_percent'] > 90) |
                (df['cpu_change'] > 10) |
                (df['memory_change'] > 10)
        ).astype(int)

        return at_risk

    def train(self, csv_file='data/metrics.csv'):
        """
        Train the prediction model
        """
        print("Loading training data...")
        df = pd.read_csv(csv_file)

        if len(df) < 20:
            print("Warning: Need at least 20 data points for training")
            print("Run collector.py for longer to gather more data")
            return False

        print(f"Loaded {len(df)} data points")

        df = self.create_features(df)

        df['at_risk'] = self.create_labels(df)


        feature_columns = [
            'cpu_percent', 'memory_percent', 'disk_percent',
            'cpu_ma_5', 'memory_ma_5',
            'cpu_change', 'memory_change',
            'stress_score'
        ]

        X = df[feature_columns].fillna(0)
        y = df['at_risk']

        X_scaled = self.scaler.fit_transform(X)


        print("\nTraining AI model...")
        self.model = RandomForestClassifier(
            n_estimators=100,
            max_depth=10,
            random_state=42
        )
        self.model.fit(X_scaled, y)

        feature_importance = pd.DataFrame({
            'feature': feature_columns,
            'importance': self.model.feature_importances_
        }).sort_values('importance', ascending=False)

        print("\nFeature Importance:")
        print(feature_importance.to_string(index=False))

        accuracy = self.model.score(X_scaled, y)
        print(f"\nModel Training Accuracy: {accuracy * 100:.1f}%")


        joblib.dump({
            'model': self.model,
            'scaler': self.scaler
        }, self.model_path)
        print(f"\nModel saved to: {self.model_path}")

        return True

    def load_model(self):
        """Load trained model"""
        if os.path.exists(self.model_path):
            saved = joblib.load(self.model_path)
            self.model = saved['model']
            self.scaler = saved['scaler']
            return True
        return False

    def predict(self, current_metrics, recent_history):
        if self.model is None:
            if not self.load_model():
                return {'error': 'No trained model found. Run training first.'}

        df = pd.DataFrame(recent_history + [current_metrics])
        df = self.create_features(df)


        current = df.iloc[-1:][
            ['cpu_percent', 'memory_percent', 'disk_percent',
             'cpu_ma_5', 'memory_ma_5',
             'cpu_change', 'memory_change',
             'stress_score']
        ].fillna(0)


        current_scaled = self.scaler.transform(current)
        prediction = self.model.predict(current_scaled)[0]
        probability = self.model.predict_proba(current_scaled)[0]

        try:
            confidence = probability[prediction] * 100
            risk_prob = probability[1] * 100 if len(probability) > 1 else probability[0] * 100
        except IndexError:

            confidence = max(probability) * 100
            risk_prob = probability[0] * 100

        return {
            'at_risk': bool(prediction),
            'confidence': round(confidence, 1),
            'risk_probability': round(risk_prob, 1),
            'status': 'WARNING' if prediction == 1 else 'HEALTHY'
        }


if __name__ == "__main__":
    # Train the model
    predictor = SystemHealthPredictor()

    if os.path.exists('data/metrics.csv'):
        success = predictor.train('data/metrics.csv')

        if success:
            print("\n✓ Model trained successfully!")
            print("\nNext steps:")
            print("1. Run monitor.py to see real-time predictions")
            print("2. Or integrate this into your dashboard")
    else:
        print("Error: No metrics data found!")
        print("Please run collector.py first to gather training data")