regime.py

import pandas as pd
import numpy as np
from typing import Dict, Any
from config import ATR_PERIOD, STRUCTURE_LOOKBACK, VOLATILITY_EXPANSION_MULTIPLIER


def compute_atr(df: pd.DataFrame, period: int = ATR_PERIOD) -> pd.Series:
    high = df["high"]
    low = df["low"]
    close = df["close"]
    prev_close = close.shift(1)
    tr = pd.concat(
        [
            high - low,
            (high - prev_close).abs(),
            (low - prev_close).abs(),
        ],
        axis=1,
    ).max(axis=1)
    atr = tr.ewm(span=period, adjust=False).mean()
    return atr


def detect_structure(df: pd.DataFrame, lookback: int = STRUCTURE_LOOKBACK) -> pd.Series:
    highs = df["high"].rolling(lookback).max()
    lows = df["low"].rolling(lookback).min()
    prev_highs = highs.shift(lookback // 2)
    prev_lows = lows.shift(lookback // 2)
    hh = highs > prev_highs
    hl = lows > prev_lows
    lh = highs < prev_highs
    ll = lows < prev_lows
    structure = pd.Series(0, index=df.index)
    structure[hh & hl] = 1
    structure[lh & ll] = -1
    return structure


def compute_volatility_regime(df: pd.DataFrame, period: int = ATR_PERIOD) -> pd.Series:
    atr = compute_atr(df, period)
    atr_ma = atr.rolling(period * 2).mean()
    ratio = atr / atr_ma.replace(0, np.nan)
    return ratio


def detect_regime(df: pd.DataFrame) -> Dict[str, Any]:
    atr = compute_atr(df, ATR_PERIOD)
    structure = detect_structure(df, STRUCTURE_LOOKBACK)
    vol_ratio = compute_volatility_regime(df, ATR_PERIOD)

    last_atr = atr.iloc[-1]
    last_structure = structure.iloc[-1]
    last_vol_ratio = vol_ratio.iloc[-1]

    recent_structure = structure.iloc[-5:]
    bullish_count = (recent_structure == 1).sum()
    bearish_count = (recent_structure == -1).sum()

    if bullish_count >= 3:
        trend = "bullish"
    elif bearish_count >= 3:
        trend = "bearish"
    else:
        trend = "ranging"

    volatility_expanding = last_vol_ratio > VOLATILITY_EXPANSION_MULTIPLIER

    if trend == "bullish" and not volatility_expanding:
        regime_score = 1.0
    elif trend == "bullish" and volatility_expanding:
        regime_score = 0.6
    elif trend == "ranging":
        regime_score = 0.3
    elif trend == "bearish" and not volatility_expanding:
        regime_score = 0.2
    else:
        regime_score = 0.1

    close = df["close"].iloc[-1]
    atr_pct = last_atr / close if close > 0 else 0.0

    return {
        "atr": last_atr,
        "atr_pct": atr_pct,
        "structure": last_structure,
        "trend": trend,
        "vol_ratio": last_vol_ratio,
        "volatility_expanding": volatility_expanding,
        "regime_score": regime_score,
        "atr_series": atr,
        "structure_series": structure,
        "vol_ratio_series": vol_ratio,
    }