feature_engineering.py

"""Feature engineering for stock price prediction."""

import numpy as np
import pandas as pd


def add_technical_indicators(df):
    """Add technical indicators as features for the prediction model.

    Args:
        df: DataFrame with OHLCV columns.

    Returns:
        DataFrame with additional feature columns.
    """
    df = df.copy()

    # --- Price-based features ---
    # Simple Moving Averages
    df["sma_5"] = df["close"].rolling(window=5).mean()
    df["sma_10"] = df["close"].rolling(window=10).mean()
    df["sma_20"] = df["close"].rolling(window=20).mean()
    df["sma_50"] = df["close"].rolling(window=50).mean()

    # Exponential Moving Averages
    df["ema_12"] = df["close"].ewm(span=12, adjust=False).mean()
    df["ema_26"] = df["close"].ewm(span=26, adjust=False).mean()

    # MACD
    df["macd"] = df["ema_12"] - df["ema_26"]
    df["macd_signal"] = df["macd"].ewm(span=9, adjust=False).mean()
    df["macd_histogram"] = df["macd"] - df["macd_signal"]

    # --- Momentum features ---
    # Relative Strength Index (RSI)
    df["rsi_14"] = _compute_rsi(df["close"], 14)

    # Rate of Change
    df["roc_5"] = df["close"].pct_change(periods=5)
    df["roc_10"] = df["close"].pct_change(periods=10)

    # Momentum
    df["momentum_10"] = df["close"] - df["close"].shift(10)

    # --- Volatility features ---
    # Bollinger Bands
    bb_mid = df["close"].rolling(window=20).mean()
    bb_std = df["close"].rolling(window=20).std()
    df["bb_upper"] = bb_mid + 2 * bb_std
    df["bb_lower"] = bb_mid - 2 * bb_std
    df["bb_width"] = df["bb_upper"] - df["bb_lower"]
    # Position within bands (0 = lower band, 1 = upper band)
    df["bb_position"] = (df["close"] - df["bb_lower"]) / df["bb_width"]

    # Average True Range
    df["atr_14"] = _compute_atr(df, 14)

    # Historical volatility (20-day rolling std of returns)
    df["volatility_20"] = df["close"].pct_change().rolling(window=20).std()

    # --- Volume features ---
    df["volume_sma_20"] = df["volume"].rolling(window=20).mean()
    df["volume_ratio"] = df["volume"] / df["volume_sma_20"]

    # On-Balance Volume (simplified)
    df["obv"] = _compute_obv(df)

    # --- Lag features ---
    for lag in [1, 2, 3, 5]:
        df[f"return_lag_{lag}"] = df["close"].pct_change(periods=lag)

    # --- Date features ---
    df["day_of_week"] = df.index.dayofweek
    df["month"] = df.index.month

    # --- Target variable ---
    # Next day's close price (what we're predicting)
    df["target"] = df["close"].shift(-1)

    # Also create a directional target (up/down)
    df["target_direction"] = (df["target"] > df["close"]).astype(int)

    # Drop rows with NaN (from rolling calculations and target shift)
    df = df.dropna()

    return df


def _compute_rsi(series, period):
    """Compute Relative Strength Index."""
    delta = series.diff()
    gain = delta.where(delta > 0, 0.0)
    loss = -delta.where(delta < 0, 0.0)

    avg_gain = gain.rolling(window=period, min_periods=period).mean()
    avg_loss = loss.rolling(window=period, min_periods=period).mean()

    rs = avg_gain / avg_loss
    rsi = 100 - (100 / (1 + rs))
    return rsi


def _compute_atr(df, period):
    """Compute Average True Range."""
    high_low = df["high"] - df["low"]
    high_close = (df["high"] - df["close"].shift(1)).abs()
    low_close = (df["low"] - df["close"].shift(1)).abs()

    true_range = pd.concat([high_low, high_close, low_close], axis=1).max(axis=1)
    return true_range.rolling(window=period).mean()


def _compute_obv(df):
    """Compute On-Balance Volume."""
    obv = pd.Series(0, index=df.index, dtype=float)
    for i in range(1, len(df)):
        if df["close"].iloc[i] > df["close"].iloc[i - 1]:
            obv.iloc[i] = obv.iloc[i - 1] + df["volume"].iloc[i]
        elif df["close"].iloc[i] < df["close"].iloc[i - 1]:
            obv.iloc[i] = obv.iloc[i - 1] - df["volume"].iloc[i]
        else:
            obv.iloc[i] = obv.iloc[i - 1]
    return obv


def get_feature_columns(df):
    """Return list of feature column names (excludes target and raw OHLCV)."""
    exclude = {"open", "high", "low", "close", "volume", "target", "target_direction"}
    return [col for col in df.columns if col not in exclude]