preprocessing.py

"""
preprocessing.py

Helpers for turning cached price history into model-ready return windows.
"""

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd

from .loader import read_etf, read_stock


STOCK_ETF_MAP = {
    "AMZN": "XLY",
    "HD": "XLY",
    "NKE": "XLY",
    "CL": "XLP",
    "EL": "XLP",
    "KO": "XLP",
    "PEP": "XLP",
    "APA": "XLE",
    "OXY": "XLE",
    "WFC": "XLF",
    "GS": "XLF",
    "BLK": "XLF",
    "PFE": "XLV",
    "HUM": "XLV",
    "FDX": "XLI",
    "GD": "XLI",
    "IBM": "XLK",
    "TER": "XLK",
    "ECL": "XLB",
    "IP": "XLB",
    "DTE": "XLU",
    "WEC": "XLU",
    "SYY": "XLP",
    "TSN": "XLP",
    "CCL": "XLY",
    "EBAY": "XLY",
    "TROW": "XLF",
    "CERN": "XLV",
}


def get_etf(stock):
    """Return the benchmark ETF associated with a stock ticker."""
    return STOCK_ETF_MAP.get(stock)


def clip_returns(values):
    """Clip extreme returns so training is less dominated by outliers."""
    return np.clip(values, -0.15, 0.15)


def interleaved_log_returns(df):
    """Compute alternating overnight/intraday log returns from adjusted prices."""
    log_close = np.log(df["AdjClose"].values)
    log_open = np.log(df["AdjOpen"].values)

    # Interleaving opens and closes lets np.diff produce the two half-day legs.
    interleaved = np.empty(2 * len(log_close))
    for i in range(len(log_close)):
        interleaved[2 * i] = log_open[i]
        interleaved[2 * i + 1] = log_close[i]

    return np.diff(interleaved)


def excess_returns(stock, plotcheck=False, plotsloc=""):
    """Build clipped stock-minus-sector-ETF excess returns for one symbol."""
    etf = get_etf(stock)
    if etf is None:
        raise ValueError(f"No sector ETF mapping configured for stock '{stock}'.")

    s_df = read_stock(stock).reset_index()
    e_df = read_etf(etf).reset_index()

    s_df["date"] = pd.to_datetime(s_df["date"])
    e_df["date"] = pd.to_datetime(e_df["date"])

    df = pd.merge(s_df, e_df, on="date", suffixes=("_s", "_e"))
    df = df[df["date"] < "2022-01-01"]

    s_ret = clip_returns(
        interleaved_log_returns(
            df.rename(
                columns={
                    "AdjOpen_s": "AdjOpen",
                    "AdjClose_s": "AdjClose",
                }
            )
        )
    )
    e_ret = clip_returns(
        interleaved_log_returns(
            df.rename(
                columns={
                    "AdjOpen_e": "AdjOpen",
                    "AdjClose_e": "AdjClose",
                }
            )
        )
    )

    excess = s_ret - e_ret
    dates = np.repeat(df["date"].values, 2)[1:]

    if plotcheck:
        plt.figure()
        plt.plot(df["date"], df["AdjClose_s"])
        plt.title(f"{stock} price")
        plt.show()

        plt.figure()
        plt.plot(s_ret, label="stock")
        plt.plot(e_ret, label="etf")
        plt.plot(excess, label="excess")
        plt.legend()
        plt.savefig(plotsloc + stock + "_excess.png")
        plt.show()

    return excess, dates


def raw_returns(stock, plotcheck=False, plotsloc=""):
    """Build clipped raw returns for a stock or ETF without sector adjustment."""
    try:
        s_df = read_stock(stock).reset_index()
    except Exception:
        s_df = read_etf(stock).reset_index()

    s_df["date"] = pd.to_datetime(s_df["date"])
    s_df = s_df[s_df["date"] < "2022-01-01"]

    returns = clip_returns(interleaved_log_returns(s_df))
    dates = np.repeat(s_df["date"].values, 2)[1:]

    if plotcheck:
        plt.figure()
        plt.plot(s_df["date"], s_df["AdjClose"])
        plt.title(f"{stock} price")
        plt.show()

        plt.figure()
        plt.plot(returns, label="returns")
        plt.legend()
        plt.savefig(plotsloc + stock + "_raw.png")
        plt.show()

    return returns, dates


def beta_excessreturns_interleaved(stock, window=60, plotcheck=False, plotsloc=""):
    """
    Like excess_returns(), but the ETF leg is multiplied by a rolling beta
    estimated separately for intraday and overnight returns.
    """
    etf = get_etf(stock)
    if etf is None:
        raise ValueError(f"No sector ETF mapping configured for stock '{stock}'.")

    s_df = read_stock(stock).reset_index()
    e_df = read_etf(etf).reset_index()

    s_df["date"] = pd.to_datetime(s_df["date"])
    e_df["date"] = pd.to_datetime(e_df["date"])

    df = pd.merge(s_df, e_df, on="date", suffixes=("_s", "_e"))
    df = df[df["date"] < "2022-01-01"].reset_index(drop=True)
    n = len(df)

    s_ret = interleaved_log_returns(df.rename(columns={
        "AdjOpen_s": "AdjOpen",
        "AdjClose_s": "AdjClose",
    }))
    e_ret = interleaved_log_returns(df.rename(columns={
        "AdjOpen_e": "AdjOpen",
        "AdjClose_e": "AdjClose",
    }))

    # even idx = intraday returns (length n), odd idx = overnight (length n-1)
    s_day,   e_day   = pd.Series(s_ret[0::2]), pd.Series(e_ret[0::2])
    s_night, e_night = pd.Series(s_ret[1::2]), pd.Series(e_ret[1::2])

    beta_day = (s_day.rolling(window).cov(e_day) / e_day.rolling(window).var()).shift(1).clip(-2, 2)
    beta_night = (s_night.rolling(window).cov(e_night) / e_night.rolling(window).var()).shift(1).clip(-2, 2)

    day_excess = clip_returns((s_day - beta_day * e_day).values)
    night_excess = clip_returns((s_night - beta_night * e_night).values)

    beta_excessret = np.empty(2 * n - 1)
    beta_excessret[0::2] = day_excess
    beta_excessret[1::2] = night_excess

    beta_day_used = np.full(2 * n - 1, np.nan)
    beta_night_used = np.full(2 * n - 1, np.nan)
    beta_day_used[0::2] = beta_day.values
    beta_night_used[1::2] = beta_night.values

    dates_interleaved = np.repeat(df["date"].values, 2)[1:]

    mask = ~np.isnan(beta_excessret)
    beta_excessret = beta_excessret[mask]
    beta_day_used = beta_day_used[mask]
    beta_night_used = beta_night_used[mask]
    dates_interleaved = dates_interleaved[mask]

    if plotcheck:
        plt.figure()
        plt.plot(df["date"], df["AdjClose_s"])
        plt.title(f"{stock} price")
        plt.show()

        plt.figure()
        plt.plot(beta_day, alpha=0.8, label="day beta")
        plt.axhline(1.0, linestyle="--")
        plt.title(f"Rolling intraday beta {stock} vs {etf}")
        plt.legend()
        plt.savefig(plotsloc + stock + "-day-beta.png")
        plt.show()

        plt.figure()
        plt.plot(beta_night, alpha=0.8, label="night beta")
        plt.axhline(1.0, linestyle="--")
        plt.title(f"Rolling overnight beta {stock} vs {etf}")
        plt.legend()
        plt.savefig(plotsloc + stock + "-night-beta.png")
        plt.show()

        plt.figure()
        plt.plot(beta_excessret, alpha=0.8, label="beta-adjusted excess")
        plt.title(f"Separate interleaved beta-adjusted excess returns {stock}")
        plt.legend()
        plt.savefig(plotsloc + stock + "-separate-beta-excessret.png")
        plt.show()

    return beta_excessret, beta_day_used, beta_night_used, dates_interleaved


def make_windows(data, l=10, pred=1, h=1):
    """Slice a 1D return series into overlapping condition-plus-target windows."""
    n = int((len(data) - l - pred) / h) + 1
    windows = np.zeros((n, l + pred))

    idx = 0
    for i in range(n):
        windows[i] = data[idx:idx + l + pred]
        idx += h

    return windows


def split_series_train_val_test(series, dates, tr=0.8, vl=0.1, l=10, pred=1, h=1):
    """Split the raw series first, then window each split independently."""
    total = len(series)
    n_train = int(tr * total)
    n_val = int(vl * total)

    # Splitting before windowing avoids windows that leak across partitions.
    train = series[:n_train]
    val = series[n_train:n_train + n_val]
    test = series[n_train + n_val:]

    train_data = make_windows(train, l, pred, h)
    val_data = make_windows(val, l, pred, h)
    test_data = make_windows(test, l, pred, h)

    return train_data, val_data, test_data, dates


def split_train_val_test(stock, tr=0.8, vl=0.1, l=10, pred=1, h=1):
    """Load excess returns for a stock and split them into model windows."""
    series, dates = excess_returns(stock)
    return split_series_train_val_test(series, dates, tr=tr, vl=vl, l=l, pred=pred, h=h)


def split_train_val_testraw(stock, tr=0.8, vl=0.1, l=10, pred=1, h=1):
    """Load raw returns for a stock/ETF and split them into model windows."""
    series, dates = raw_returns(stock)
    return split_series_train_val_test(series, dates, tr=tr, vl=vl, l=l, pred=pred, h=h)