ml_models.py

"""Model training, evaluation, and code generation."""

import io
import os

import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from sklearn.ensemble import (
    GradientBoostingClassifier,
    GradientBoostingRegressor,
    RandomForestClassifier,
    RandomForestRegressor,
)
from sklearn.linear_model import LinearRegression, LogisticRegression
from sklearn.metrics import (
    accuracy_score,
    confusion_matrix,
    f1_score,
    mean_absolute_error,
    mean_squared_error,
    precision_score,
    r2_score,
    recall_score,
)
from sklearn.model_selection import train_test_split
from sklearn.neighbors import KNeighborsClassifier, KNeighborsRegressor
from sklearn.preprocessing import LabelEncoder
from sklearn.svm import SVC, SVR
from sklearn.tree import DecisionTreeClassifier, DecisionTreeRegressor

from data_processor import TEMP_DIR

REGRESSION_MODELS = {
    "Linear Regression": LinearRegression,
    "Decision Tree Regressor": DecisionTreeRegressor,
    "Random Forest Regressor": RandomForestRegressor,
    "Gradient Boosting Regressor": GradientBoostingRegressor,
    "Support Vector Regressor (SVR)": SVR,
    "K-Nearest Neighbors Regressor": KNeighborsRegressor,
}

CLASSIFICATION_MODELS = {
    "Logistic Regression": LogisticRegression,
    "Decision Tree Classifier": DecisionTreeClassifier,
    "Random Forest Classifier": RandomForestClassifier,
    "Gradient Boosting Classifier": GradientBoostingClassifier,
    "Support Vector Classifier (SVC)": SVC,
    "K-Nearest Neighbors Classifier": KNeighborsClassifier,
}

MODEL_IMPORTS = {
    "Linear Regression": "from sklearn.linear_model import LinearRegression",
    "Decision Tree Regressor": "from sklearn.tree import DecisionTreeRegressor",
    "Random Forest Regressor": "from sklearn.ensemble import RandomForestRegressor",
    "Gradient Boosting Regressor": "from sklearn.ensemble import GradientBoostingRegressor",
    "Support Vector Regressor (SVR)": "from sklearn.svm import SVR",
    "K-Nearest Neighbors Regressor": "from sklearn.neighbors import KNeighborsRegressor",
    "Logistic Regression": "from sklearn.linear_model import LogisticRegression",
    "Decision Tree Classifier": "from sklearn.tree import DecisionTreeClassifier",
    "Random Forest Classifier": "from sklearn.ensemble import RandomForestClassifier",
    "Gradient Boosting Classifier": "from sklearn.ensemble import GradientBoostingClassifier",
    "Support Vector Classifier (SVC)": "from sklearn.svm import SVC",
    "K-Nearest Neighbors Classifier": "from sklearn.neighbors import KNeighborsClassifier",
}


def suggest_task(df: pd.DataFrame, target: str) -> str:
    """Suggest 'Classification' or 'Regression' based on the target column."""
    s = df[target]
    if not pd.api.types.is_numeric_dtype(s):
        return "Classification"
    if s.nunique() <= 10:
        return "Classification"
    return "Regression"


def _make_model(name: str, task: str):
    cls = (CLASSIFICATION_MODELS if task == "Classification" else REGRESSION_MODELS)[name]
    kwargs = {}
    if "Logistic" in name:
        kwargs["max_iter"] = 1000
    if "Random Forest" in name or "Gradient Boosting" in name:
        kwargs["random_state"] = 42
    if "Decision Tree" in name:
        kwargs["random_state"] = 42
    return cls(**kwargs)


def train_model(
    df: pd.DataFrame,
    target: str,
    model_name: str,
    task: str,
    test_size: float = 0.2,
) -> dict:
    """Train a model and return metrics, plot path, and predictions sample."""
    df = df.dropna(subset=[target])
    X = df.drop(columns=[target]).select_dtypes(include=np.number)
    if X.empty:
        raise ValueError("No numeric feature columns available. Run preprocessing first.")
    y = df[target]

    label_encoder = None
    if task == "Classification" and not pd.api.types.is_numeric_dtype(y):
        label_encoder = LabelEncoder()
        y = pd.Series(label_encoder.fit_transform(y), index=df.index)

    stratify = y if (task == "Classification" and y.value_counts().min() >= 2) else None
    X_train, X_test, y_train, y_test = train_test_split(
        X, y, test_size=test_size, random_state=42, stratify=stratify
    )

    model = _make_model(model_name, task)
    model.fit(X_train, y_train)
    y_pred = model.predict(X_test)

    if task == "Classification":
        avg = "binary" if pd.Series(y).nunique() == 2 else "weighted"
        metrics = {
            "Accuracy": round(accuracy_score(y_test, y_pred), 4),
            "Precision": round(precision_score(y_test, y_pred, average=avg, zero_division=0), 4),
            "Recall": round(recall_score(y_test, y_pred, average=avg, zero_division=0), 4),
            "F1 Score": round(f1_score(y_test, y_pred, average=avg, zero_division=0), 4),
        }
        plot_path = _plot_confusion(y_test, y_pred, model_name, label_encoder)
    else:
        metrics = {
            "R² Score": round(r2_score(y_test, y_pred), 4),
            "MAE": round(mean_absolute_error(y_test, y_pred), 4),
            "RMSE": round(float(np.sqrt(mean_squared_error(y_test, y_pred))), 4),
        }
        plot_path = _plot_regression(y_test, y_pred, model_name)

    importance_path = _plot_importance(model, X.columns, model_name)

    return {
        "model_name": model_name,
        "task": task,
        "target": target,
        "features": X.columns.tolist(),
        "n_train": len(X_train),
        "n_test": len(X_test),
        "metrics": metrics,
        "plot_path": plot_path,
        "importance_path": importance_path,
    }


def _plot_confusion(y_test, y_pred, model_name, label_encoder=None):
    cm = confusion_matrix(y_test, y_pred)
    labels = label_encoder.classes_ if label_encoder is not None else sorted(set(y_test))
    fig, ax = plt.subplots(figsize=(5.5, 4.5))
    im = ax.imshow(cm, cmap="Blues")
    ax.set_xticks(range(len(labels)), labels=[str(l) for l in labels], rotation=45, ha="right")
    ax.set_yticks(range(len(labels)), labels=[str(l) for l in labels])
    for i in range(cm.shape[0]):
        for j in range(cm.shape[1]):
            ax.text(j, i, cm[i, j], ha="center", va="center",
                    color="white" if cm[i, j] > cm.max() / 2 else "black")
    ax.set_xlabel("Predicted")
    ax.set_ylabel("Actual")
    ax.set_title(f"Confusion Matrix — {model_name}")
    fig.colorbar(im)
    fig.tight_layout()
    path = os.path.join(TEMP_DIR, "result_plot.png")
    fig.savefig(path, dpi=120)
    plt.close(fig)
    return path


def _plot_regression(y_test, y_pred, model_name):
    fig, ax = plt.subplots(figsize=(5.5, 4.5))
    ax.scatter(y_test, y_pred, alpha=0.5, edgecolors="none")
    lims = [min(y_test.min(), y_pred.min()), max(y_test.max(), y_pred.max())]
    ax.plot(lims, lims, "r--", label="Perfect prediction")
    ax.set_xlabel("Actual")
    ax.set_ylabel("Predicted")
    ax.set_title(f"Actual vs Predicted — {model_name}")
    ax.legend()
    fig.tight_layout()
    path = os.path.join(TEMP_DIR, "result_plot.png")
    fig.savefig(path, dpi=120)
    plt.close(fig)
    return path


def _plot_importance(model, feature_names, model_name):
    importances = None
    if hasattr(model, "feature_importances_"):
        importances = model.feature_importances_
    elif hasattr(model, "coef_"):
        coef = model.coef_
        importances = np.abs(coef).mean(axis=0) if coef.ndim > 1 else np.abs(coef)
    if importances is None:
        return None
    order = np.argsort(importances)[-15:]
    fig, ax = plt.subplots(figsize=(6, 4.5))
    ax.barh([feature_names[i] for i in order], importances[order], color="#4C72B0")
    ax.set_title(f"Feature Importance — {model_name}")
    fig.tight_layout()
    path = os.path.join(TEMP_DIR, "importance_plot.png")
    fig.savefig(path, dpi=120)
    plt.close(fig)
    return path


def generate_model_code(model_name: str, task: str, target: str, test_size: float) -> str:
    """Return equivalent standalone sklearn code."""
    cls = (CLASSIFICATION_MODELS if task == "Classification" else REGRESSION_MODELS)[model_name]
    kwargs = []
    if "Logistic" in model_name:
        kwargs.append("max_iter=1000")
    if any(k in model_name for k in ("Random Forest", "Gradient Boosting", "Decision Tree")):
        kwargs.append("random_state=42")
    ctor = f"{cls.__name__}({', '.join(kwargs)})"

    lines = [
        "import numpy as np",
        "import pandas as pd",
        "from sklearn.model_selection import train_test_split",
        MODEL_IMPORTS[model_name],
    ]
    if task == "Classification":
        lines += [
            "from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score",
            "from sklearn.preprocessing import LabelEncoder",
        ]
    else:
        lines.append(
            "from sklearn.metrics import r2_score, mean_absolute_error, mean_squared_error"
        )
    lines += [
        "",
        "df = pd.read_csv('cleaned_data.csv')",
        f"target = {target!r}",
        "X = df.drop(columns=[target]).select_dtypes(include=np.number)",
        "y = df[target]",
    ]
    if task == "Classification":
        lines += [
            "if not pd.api.types.is_numeric_dtype(y):",
            "    y = LabelEncoder().fit_transform(y)",
        ]
    lines += [
        "",
        "X_train, X_test, y_train, y_test = train_test_split(",
        f"    X, y, test_size={test_size}, random_state=42)",
        "",
        f"model = {ctor}",
        "model.fit(X_train, y_train)",
        "y_pred = model.predict(X_test)",
        "",
    ]
    if task == "Classification":
        lines += [
            "print('Accuracy :', accuracy_score(y_test, y_pred))",
            "print('Precision:', precision_score(y_test, y_pred, average='weighted'))",
            "print('Recall   :', recall_score(y_test, y_pred, average='weighted'))",
            "print('F1 Score :', f1_score(y_test, y_pred, average='weighted'))",
        ]
    else:
        lines += [
            "print('R2  :', r2_score(y_test, y_pred))",
            "print('MAE :', mean_absolute_error(y_test, y_pred))",
            "print('RMSE:', np.sqrt(mean_squared_error(y_test, y_pred)))",
        ]
    return "\n".join(lines)