model.py

"""Random Forest stock prediction model with GridSearchCV optimization."""

import numpy as np
import pandas as pd
from sklearn.ensemble import RandomForestRegressor
from sklearn.model_selection import GridSearchCV, TimeSeriesSplit, cross_val_score
from sklearn.metrics import (
    mean_absolute_error,
    mean_squared_error,
    r2_score,
    accuracy_score,
)
from sklearn.preprocessing import StandardScaler
import joblib

from feature_engineering import add_technical_indicators, get_feature_columns


class StockPredictor:
    """Random Forest-based stock price predictor with hyperparameter optimization."""

    def __init__(self):
        self.model = None
        self.scaler = StandardScaler()
        self.feature_columns = None
        self.is_fitted = False

    def prepare_data(self, df, test_size=0.2):
        """Prepare data for training with time-series aware splitting.

        Args:
            df: Raw OHLCV DataFrame.
            test_size: Fraction of data to use for testing.

        Returns:
            Tuple of (X_train, X_test, y_train, y_test, featured_df).
        """
        featured_df = add_technical_indicators(df)
        self.feature_columns = get_feature_columns(featured_df)

        X = featured_df[self.feature_columns].values
        y = featured_df["target"].values

        # Time-series split: use last test_size fraction as test set
        split_idx = int(len(X) * (1 - test_size))
        X_train, X_test = X[:split_idx], X[split_idx:]
        y_train, y_test = y[:split_idx], y[split_idx:]

        # Scale features
        X_train = self.scaler.fit_transform(X_train)
        X_test = self.scaler.transform(X_test)

        return X_train, X_test, y_train, y_test, featured_df

    def train_baseline(self, X_train, y_train):
        """Train a baseline Random Forest model without tuning.

        Args:
            X_train: Training features.
            y_train: Training targets.

        Returns:
            Fitted RandomForestRegressor.
        """
        self.model = RandomForestRegressor(
            n_estimators=100,
            random_state=42,
            n_jobs=-1,
        )
        self.model.fit(X_train, y_train)
        self.is_fitted = True
        return self.model

    def train_optimized(self, X_train, y_train, verbose=1):
        """Train an optimized model using GridSearchCV with TimeSeriesSplit.

        Args:
            X_train: Training features.
            y_train: Training targets.
            verbose: Verbosity level for GridSearchCV.

        Returns:
            Dict with best_params, best_score, and the fitted model.
        """
        param_grid = {
            "n_estimators": [100, 200, 300],
            "max_depth": [10, 20, 30, None],
            "min_samples_split": [2, 5, 10],
            "min_samples_leaf": [1, 2, 4],
            "max_features": ["sqrt", "log2"],
        }

        tscv = TimeSeriesSplit(n_splits=5)

        grid_search = GridSearchCV(
            estimator=RandomForestRegressor(random_state=42, n_jobs=-1),
            param_grid=param_grid,
            cv=tscv,
            scoring="neg_mean_squared_error",
            verbose=verbose,
            n_jobs=-1,
        )

        grid_search.fit(X_train, y_train)

        self.model = grid_search.best_estimator_
        self.is_fitted = True

        return {
            "best_params": grid_search.best_params_,
            "best_score": -grid_search.best_score_,  # Convert back to positive MSE
            "model": self.model,
        }

    def evaluate(self, X_test, y_test):
        """Evaluate the model on test data.

        Args:
            X_test: Test features.
            y_test: Test targets (actual next-day prices).

        Returns:
            Dict with MAE, RMSE, R2, MAPE, and directional accuracy.
        """
        if not self.is_fitted:
            raise RuntimeError("Model must be trained before evaluation.")

        predictions = self.model.predict(X_test)

        mae = mean_absolute_error(y_test, predictions)
        rmse = np.sqrt(mean_squared_error(y_test, predictions))
        r2 = r2_score(y_test, predictions)
        mape = np.mean(np.abs((y_test - predictions) / y_test)) * 100

        # Directional accuracy: did we predict the right direction of movement?
        if len(y_test) > 1:
            actual_direction = (np.diff(y_test) > 0).astype(int)
            pred_direction = (np.diff(predictions) > 0).astype(int)
            dir_accuracy = accuracy_score(actual_direction, pred_direction)
        else:
            dir_accuracy = 0.0

        return {
            "mae": mae,
            "rmse": rmse,
            "r2": r2,
            "mape": mape,
            "directional_accuracy": dir_accuracy,
            "predictions": predictions,
        }

    def cross_validate(self, X_train, y_train, n_splits=5):
        """Perform time-series cross-validation.

        Args:
            X_train: Training features.
            y_train: Training targets.
            n_splits: Number of CV splits.

        Returns:
            Dict with mean and std of cross-validation scores.
        """
        if not self.is_fitted:
            raise RuntimeError("Model must be trained before cross-validation.")

        tscv = TimeSeriesSplit(n_splits=n_splits)
        scores = cross_val_score(
            self.model, X_train, y_train,
            cv=tscv, scoring="neg_mean_squared_error"
        )

        return {
            "mean_mse": -scores.mean(),
            "std_mse": scores.std(),
            "scores": [-s for s in scores],
        }

    def get_feature_importance(self, top_n=15):
        """Get top feature importances from the trained model.

        Args:
            top_n: Number of top features to return.

        Returns:
            List of (feature_name, importance) tuples sorted by importance.
        """
        if not self.is_fitted or self.feature_columns is None:
            raise RuntimeError("Model must be trained first.")

        importances = self.model.feature_importances_
        feature_imp = list(zip(self.feature_columns, importances))
        feature_imp.sort(key=lambda x: x[1], reverse=True)
        return feature_imp[:top_n]

    def save_model(self, filepath):
        """Save the trained model and scaler to disk."""
        joblib.dump({
            "model": self.model,
            "scaler": self.scaler,
            "feature_columns": self.feature_columns,
        }, filepath)

    def load_model(self, filepath):
        """Load a trained model and scaler from disk."""
        data = joblib.load(filepath)
        self.model = data["model"]
        self.scaler = data["scaler"]
        self.feature_columns = data["feature_columns"]
        self.is_fitted = True