Machine Learning Trading Strategy

import numpy as np
import pandas as pd
from sklearn.ensemble import RandomForestClassifier, GradientBoostingClassifier
from sklearn.linear_model import LogisticRegression
from sklearn.preprocessing import StandardScaler


def validate_data(df: pd.DataFrame, min_rows: int = 300) -> bool:
    """Check that OHLCV data meets minimum quality for ML training.

    Args:
        df: DataFrame with DatetimeIndex.
        min_rows: Minimum number of rows required.

    Returns:
        True if data is usable.
    """
    required = {"open", "high", "low", "close", "volume"}
    if not required.issubset(df.columns):
        return False
    if len(df) < min_rows:
        return False
    if df["close"].isnull().mean() > 0.2:
        return False
    return True


def build_features(df: pd.DataFrame) -> pd.DataFrame:
    """Build a machine-learning feature matrix from OHLCV data.

    All features are guarded against division-by-zero and sanitized
    (inf replaced with NaN) so downstream code never sees inf values.

    Args:
        df: DataFrame containing open, high, low, close, and volume columns.

    Returns:
        DataFrame with feature columns prefixed by 'f_'.
    """
    c = df["close"]
    v = df["volume"]
    ret = c.pct_change(fill_method=None)

    features = pd.DataFrame(index=df.index)
    features["f_ret_5d"] = c.pct_change(5, fill_method=None)
    features["f_ret_20d"] = c.pct_change(20, fill_method=None)
    features["f_vol_20d"] = ret.rolling(20).std()
    features["f_ma_ratio"] = c / c.rolling(20).mean()
    features["f_volume_ratio"] = v / v.rolling(20).mean()

    # RSI(14) — guard: loss=0 in zero-volatility periods produces inf
    delta = c.diff()
    gain = delta.clip(lower=0).rolling(14).mean()
    loss = (-delta.clip(upper=0)).rolling(14).mean()
    rs = gain / loss.replace(0, np.nan)
    features["f_rsi_14"] = 100 - (100 / (1 + rs))

    # Bollinger Band position — guard: bb_upper == bb_lower when std=0
    ma20 = c.rolling(20).mean()
    std20 = c.rolling(20).std()
    bb_upper = ma20 + 2 * std20
    bb_lower = ma20 - 2 * std20
    bb_range = (bb_upper - bb_lower).replace(0, np.nan)
    features["f_bb_position"] = (c - bb_lower) / bb_range

    # Intraday features
    features["f_high_low_ratio"] = (df["high"] - df["low"]) / c
    features["f_close_open_ratio"] = (c - df["open"]) / df["open"]
    features["f_skew_20d"] = ret.rolling(20).skew()

    # Sanitize: replace all inf with NaN (NaN handled by walk-forward)
    features = features.replace([np.inf, -np.inf], np.nan)
    return features


def walk_forward_predict(
    features: pd.DataFrame,
    labels: pd.Series,
    min_train_size: int = 252,
    retrain_freq: int = 20,
    model_type: str = "random_forest",
    window_type: str = "expanding",
    sliding_size: int = 504,
    prediction_horizon: int = 5,
) -> pd.Series:
    """Walk-forward training and prediction to avoid future data leakage.

    Args:
        features: Feature matrix aligned with labels by row index.
        labels: Binary labels (0/1), representing the direction of future N-day returns.
        min_train_size: Minimum training-set size in trading days.
        retrain_freq: Retrain the model every N days.
        model_type: One of "random_forest" / "gradient_boosting" / "ridge".
        window_type: "expanding" uses all history; "sliding" uses a fixed lookback.
        sliding_size: Lookback window size when window_type is "sliding".
        prediction_horizon: Number of bars each target label looks ahead.

    Returns:
        Predicted signal series with range [-1.0, 1.0], no NaN values.
    """
    predictions = pd.Series(0.0, index=features.index)
    model = None
    scaler = None

    if prediction_horizon < 1:
        raise ValueError("prediction_horizon must be >= 1")

    for i in range(min_train_size, len(features)):
        # Retrain every retrain_freq days
        if model is None or (i - min_train_size) % retrain_freq == 0:
            # A label at row t is observable only once t + horizon <= i.
            train_stop = max(0, i - prediction_horizon + 1)
            start = (
                max(0, train_stop - sliding_size)
                if window_type == "sliding"
                else 0
            )
            X_train = features.iloc[start:train_stop].values
            y_train = labels.iloc[start:train_stop].values

            # Drop rows with NaN
            valid = ~(np.isnan(X_train).any(axis=1) | np.isnan(y_train))
            X_train = X_train[valid]
            y_train = y_train[valid]

            # Too few rows, or a single class (a sustained one-directional
            # trend, which crashes fit()/predict_proba() on every model type
            # below): skip the retrain. The previous model keeps serving; before
            # the first model there is no prediction.
            if len(X_train) >= 50 and len(np.unique(y_train)) >= 2:
                # Standardization: fit only on training set
                scaler = StandardScaler()
                X_train = scaler.fit_transform(X_train)

                # Build the model
                if model_type == "random_forest":
                    model = RandomForestClassifier(
                        n_estimators=100, max_depth=5, random_state=42,
                    )
                elif model_type == "gradient_boosting":
                    model = GradientBoostingClassifier(
                        n_estimators=100, max_depth=3, learning_rate=0.05,
                        random_state=42,
                    )
                elif model_type == "ridge":
                    model = LogisticRegression(penalty="l2", C=1.0, random_state=42)
                else:
                    raise ValueError(f"Unsupported model_type: {model_type}")

                model.fit(X_train, y_train)

        if model is None:
            continue

        # Predict today
        X_today = features.iloc[i : i + 1].values
        if np.isnan(X_today).any():
            predictions.iloc[i] = 0.0
            continue

        X_today = scaler.transform(X_today)

        if hasattr(model, "predict_proba"):
            prob = model.predict_proba(X_today)[0, 1]
            predictions.iloc[i] = prob * 2 - 1  # [0,1] -> [-1,1]
        else:
            predictions.iloc[i] = float(model.predict(X_today)[0])

    # Output contract: no NaN, clipped to [-1, 1]
    predictions = predictions.fillna(0.0).clip(-1.0, 1.0)
    return predictions


class SignalEngine:
    """Complete ML strategy with built-in data validation and safety."""

    def generate(self, data_map: dict) -> dict:
        """Generate signals for each symbol.

        Args:
            data_map: code -> OHLCV DataFrame.

        Returns:
            code -> signal Series in [-1.0, 1.0].
        """
        signals = {}
        for code, df in data_map.items():
            if not validate_data(df):
                print(f"[WARN] {code}: data quality insufficient, skipping")
                continue

            features = build_features(df)
            prediction_horizon = 5
            future_returns = (
                df["close"].shift(-prediction_horizon) / df["close"] - 1
            )
            labels = (future_returns > 0).astype(float).where(future_returns.notna())
            signal = walk_forward_predict(
                features,
                labels,
                prediction_horizon=prediction_horizon,
            )
            signals[code] = signal

        return signals