ml_engine.py

import pandas as pd
import numpy as np
import warnings
warnings.filterwarnings("ignore")
import matplotlib.pyplot as plt
import seaborn as sns
import os
from groq import Groq
from sklearn.model_selection import train_test_split, GridSearchCV
from sklearn.compose import ColumnTransformer
from sklearn.preprocessing import OneHotEncoder, StandardScaler, LabelEncoder
from sklearn.impute import SimpleImputer
from sklearn.metrics import (
    classification_report, confusion_matrix, accuracy_score,
    r2_score, mean_absolute_error, mean_squared_error, mean_absolute_percentage_error
)
from sklearn.ensemble import RandomForestClassifier, RandomForestRegressor
from sklearn.svm import SVC
from sklearn.linear_model import Ridge
from xgboost import XGBClassifier, XGBRegressor
from imblearn.over_sampling import SMOTE
from imblearn.pipeline import Pipeline
import joblib
import io
import base64

class PredictIQEngine:
    def __init__(self, groq_api_key=None):
        self.groq_api_key = groq_api_key or os.getenv("GROQ_API_KEY")
        if self.groq_api_key:
            self.client = Groq(api_key=self.groq_api_key)
        else:
            self.client = None
        
        self.df = None
        self.target_column = None
        self.problem_type = None
        self.best_pipeline = None
        self.best_score = -np.inf
        self.X_test = None
        self.y_test = None
        self.results = {}

    def load_data(self, df):
        self.df = df
        return self.df.columns.tolist()

    def get_suggested_drops(self):
        """Identify potential ID columns or high-cardinality non-predictive columns."""
        suggested_drops = []
        for col in self.df.columns:
            # Drop columns with 'id', 'uuid', 'index', 'serial' in name
            if any(k in col.lower() for k in ['id', 'uuid', 'index', 'serial', 'roll_no']):
                suggested_drops.append(col)
            # Drop if all values are unique and non-numeric
            elif self.df[col].nunique() == len(self.df) and self.df[col].dtype == 'object':
                suggested_drops.append(col)
        return suggested_drops

    def preprocess_dates(self, df):
        """Automatically detect and extract features from date columns."""
        cols_to_drop = []
        for col in df.columns:
            if df[col].dtype == 'object':
                try:
                    # Attempt to convert to datetime
                    date_series = pd.to_datetime(df[col], errors='coerce')
                    if date_series.notnull().mean() > 0.8: # If 80% looks like dates
                        df[f"{col}_year"] = date_series.dt.year
                        df[f"{col}_month"] = date_series.dt.month
                        df[f"{col}_day"] = date_series.dt.day
                        df[f"{col}_dayofweek"] = date_series.dt.dayofweek
                        cols_to_drop.append(col)
                except:
                    continue
        return df.drop(columns=cols_to_drop)

    def detect_target(self):
        keywords = ['target', 'label', 'placement', 'result', 'output', 'grade', 'price', 'salary', 'status']
        for col in self.df.columns:
            if any(k in col.lower() for k in keywords):
                self.target_column = col
                break
        if self.target_column is None:
            self.target_column = self.df.columns[-1]
        return self.target_column

    def identify_problem(self, target_col):
        self.target_column = target_col
        y = self.df[self.target_column]
        
        # Smart detection: Categorical vs Numeric
        if y.dtype == 'object' or y.dtype.name == 'category':
            self.problem_type = "classification"
        elif y.nunique() < 10: # Low unique numeric values usually imply classification
            self.problem_type = "classification"
        else:
            self.problem_type = "regression"
        return self.problem_type

    def get_eda_plots(self):
        plots = {}
        df = self.df.copy()
        
        # Limit to top 15 columns for performance
        numeric_df = df.select_dtypes(include=['int64', 'float64'])
        numeric_cols = numeric_df.columns[:15]

        # Correlation Matrix
        if len(numeric_cols) > 1:
            plt.figure(figsize=(10, 8))
            corr = numeric_df[numeric_cols].corr()
            mask = np.triu(np.ones_like(corr, dtype=bool))
            sns.heatmap(corr, mask=mask, cmap="coolwarm", annot=len(numeric_cols) < 10, fmt=".2f")
            plt.title("Feature Correlation")
            plt.tight_layout()
            plots['correlation_matrix'] = self._plt_to_base64()

        # Distributions (Top 4 most varied numeric)
        varied_cols = numeric_df[numeric_cols].std().sort_values(ascending=False).index[:4]
        for col in varied_cols:
            plt.figure(figsize=(6, 4))
            sns.histplot(df[col], kde=True, color='#6366f1')
            plt.title(f"{col} Distribution")
            plots[f'dist_{col}'] = self._plt_to_base64()

        return plots

    def _plt_to_base64(self):
        buf = io.BytesIO()
        plt.savefig(buf, format='png', bbox_inches='tight', transparent=True)
        plt.close()
        return base64.b64encode(buf.getvalue()).decode('utf-8')

    def train(self, target_col, columns_to_drop=[]):
        self.target_column = target_col
        self.identify_problem(target_col)
        
        # Reset for fresh training session
        self.best_score = -np.inf
        self.best_pipeline = None
        
        # Drop rows where target is NaN 
        df_clean = self.df.dropna(subset=[self.target_column])
        
        # Prepare Data
        X = df_clean.drop(columns=[self.target_column] + columns_to_drop)
        X = self.preprocess_dates(X)
        y = df_clean[self.target_column]

        label_encoder = None
        if self.problem_type == "classification" and y.dtype == 'object':
            label_encoder = LabelEncoder()
            y = label_encoder.fit_transform(y)

        X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
        self.X_test, self.y_test = X_test, y_test
        
        # Build Preprocessor dynamically based on filtered X
        preprocessor = self.build_preprocessor(X)

        if self.problem_type == "classification":
            apply_smote = self.check_imbalance(y_train)
            models = {
                "Random Forest": (RandomForestClassifier(), {"model__n_estimators": [100]}),
                "SVM": (SVC(probability=True), {"model__C": [1]}),
                "XGBoost": (XGBClassifier(eval_metric='logloss'), {"model__n_estimators": [100]})
            }
            scoring = "accuracy"
        else:
            apply_smote = False
            models = {
                "Random Forest Regressor": (RandomForestRegressor(), {"model__n_estimators": [100]}),
                "Ridge": (Ridge(), {"model__alpha": [1]}),
                "XGBoost Regressor": (XGBRegressor(), {"model__n_estimators": [100]})
            }
            scoring = "r2"

        best_name = ""
        for name, (model, params) in models.items():
            steps = [("preprocessor", preprocessor)]
            if self.problem_type == "classification" and apply_smote:
                steps.append(("smote", SMOTE(k_neighbors=1, random_state=42)))
            steps.append(("model", model))
            pipeline = Pipeline(steps)
            grid = GridSearchCV(pipeline, params, cv=3, scoring=scoring)
            grid.fit(X_train, y_train)
            score = grid.score(X_test, y_test)
            if score > self.best_score:
                self.best_score = score
                self.best_pipeline = grid.best_estimator_
                best_name = name

        self.results['best_model'] = best_name
        self.results['best_score'] = self.best_score
        return self.results

    def build_preprocessor(self, X):
        numeric_features = X.select_dtypes(include=['int64', 'float64']).columns
        categorical_features = X.select_dtypes(include=['object']).columns
        
        numeric_transformer = Pipeline([
            ("imputer", SimpleImputer(strategy="median")),
            ("scaler", StandardScaler())
        ])
        
        categorical_transformer = Pipeline([
            ("imputer", SimpleImputer(strategy="most_frequent")),
            ("onehot", OneHotEncoder(handle_unknown="ignore"))
        ])
        
        return ColumnTransformer([
            ("num", numeric_transformer, numeric_features),
            ("cat", categorical_transformer, categorical_features)
        ])

    def check_imbalance(self, y_train):
        if self.problem_type != "classification":
            return False
        # Calculate balance ratio
        from collections import Counter
        counts = Counter(y_train)
        if len(counts) < 2: return False
        min_c = min(counts.values())
        max_c = max(counts.values())
        return (min_c / max_c) < 0.6

    def evaluate(self):
        if self.best_pipeline is None:
            raise ValueError("No model has been trained yet or training failed.")
            
        preds = self.best_pipeline.predict(self.X_test)
        eval_results = {}
        
        if self.problem_type == "classification":
            eval_results['report'] = classification_report(self.y_test, preds, output_dict=True)
            eval_results['report_text'] = classification_report(self.y_test, preds)
            cm = confusion_matrix(self.y_test, preds)
            plt.figure(figsize=(6, 4))
            sns.heatmap(cm, annot=True, fmt='d', cmap='Blues')
            plt.title("Confusion Matrix")
            eval_results['confusion_matrix'] = self._plt_to_base64()
            
            summary = f"""
            Target Variable: {self.target_column}
            Problem Type: Classification
            Best Model Selected: {self.results['best_model']}
            Accuracy Score: {self.best_score:.4f}
            
            Classification Detailed Report:
            {eval_results['report_text']}
            """
        else:
            r2 = r2_score(self.y_test, preds)
            mape = mean_absolute_percentage_error(self.y_test, preds)
            rmse = np.sqrt(mean_squared_error(self.y_test, preds))
            eval_results['metrics'] = {"R2": r2, "MAPE": mape, "RMSE": rmse}
            
            summary = f"""
            Target Variable: {self.target_column}
            Problem Type: Regression
            Best Model Selected: {self.results['best_model']}
            R2 Score: {r2:.4f}
            Mean Absolute Percentage Error: {mape:.4f}
            Root Mean Squared Error: {rmse:.4f}
            """

        self.results['summary'] = summary
        if self.client:
            eval_results['ai_insights'] = self.generate_ai_insights(summary)
        else:
            eval_results['ai_insights'] = "AI Insights unavailable (No Groq API Key)."
        
        return eval_results

    def generate_ai_insights(self, summary_text):
        # Gather feature names for context
        features = self.df.drop(columns=[self.target_column]).columns.tolist()
        
        prompt = f"""
        You are a Senior Data Scientist analyzing a Machine Learning project.
        
        PROJECT DATA CONTEXT:
        - Target Column: '{self.target_column}'
        - Features analyzed: {', '.join(features)}
        
        MODEL PERFORMANCE SUMMARY:
        {summary_text}
        
        Please provide a professional narrative analysis:
        1. **Model Performance**: Interpret the {self.results['best_model']} results. Is it reliable for predicting '{self.target_column}'?
        2. **Feature Importance**: Based on the features listed, which ones likely drive the prediction for '{self.target_column}'?
        3. **Impact of this Prediction in real world**: How can these insights be used in a real-world scenario ?
        4. **Next Steps**: One technical recommendation to improve this specific model.
        
        Keep the tone professional, insightful, and specific to the target column provided.
        """
        try:
            response = self.client.chat.completions.create(
                model="llama-3.3-70b-versatile",
                messages=[
                    {"role": "system", "content": "You are a professional data scientist."},
                    {"role": "user", "content": prompt}
                ]
            )
            return response.choices[0].message.content
        except Exception as e:
            return f"Error generating AI insights: {str(e)}"

    def save_model(self, path="best_model.pkl"):
        joblib.dump(self.best_pipeline, path)
        return path