utils/helpers.py

"""Shared helpers used by the app, live_runner and the feature pipeline."""
import math

import pandas as pd

MINUTES_PER_CANDLE = {"1m": 1, "5m": 5, "10m": 10, "15m": 15, "30m": 30,
                      "1h": 60, "4h": 240, "1d": 1440}

# SMA_20/Bollinger need 20 candles; the rest is safety margin.
INDICATOR_WARMUP_CANDLES = 30

# History handed to a model for a single live prediction. Comfortably above
# every model's minimum (TimesFM's 32 is the largest) even on weekday-only
# forex calendars, where ~5/7 of the requested days carry candles.
PREDICTION_HISTORY_CANDLES = 200


def infer_step(index: pd.DatetimeIndex):
    """Median candle spacing, used to project forecasts past the last candle."""
    if len(index) < 2:
        return pd.Timedelta(minutes=1)
    diffs = index.to_series().diff().dropna()
    if diffs.empty:
        return pd.Timedelta(minutes=1)
    return diffs.median()


def future_index(last_timestamp, step, n):
    """`n` timestamps spaced `step` apart, starting after last_timestamp."""
    return [last_timestamp + step * (i + 1) for i in range(n)]


def history_days_for_candles(timeframe: str, candles: int, minimum_days: int = 1) -> int:
    """Calendar days that cover `candles` candles of this timeframe. Markets
    that close (forex, weekends) yield fewer candles than this, so callers
    should treat the result as a lower bound and clamp it against the data
    source's own history limit."""
    minutes = MINUTES_PER_CANDLE.get(timeframe, 60) * max(int(candles), 0)
    return max(minimum_days, math.ceil(minutes / (24 * 60)))


def feature_warmup_days(timeframe: str, minimum_days: int = 2) -> int:
    """Extra days of history the indicators need before the tested period."""
    return history_days_for_candles(timeframe, INDICATOR_WARMUP_CANDLES,
                                    minimum_days=minimum_days)


def drop_forming_candle(df: pd.DataFrame, features_df: pd.DataFrame = None):
    """Drop the newest candle while it is still forming. Yahoo returns the
    in-progress candle, whose Close keeps moving, so grading a prediction
    against it compares against an unsettled value. features_df is trimmed
    with it to keep the two positionally aligned."""
    if df is None or len(df) == 0:
        return df, features_df
    last = df.index[-1]
    now = pd.Timestamp.now(tz="UTC")
    if last.tzinfo is None:
        now = now.tz_localize(None)
    if now - last >= infer_step(df.index):
        return df, features_df
    df = df.iloc[:-1]
    if features_df is not None:
        features_df = features_df.iloc[:-1].reset_index(drop=True)
    return df, features_df