"""Example: Volatility regime tracker.

Adapts signal strength based on ATR (volatility). In high-volatility
regimes, uses mean reversion (large moves tend to reverse). In low
volatility, follows momentum (breakouts tend to continue).
"""

from __future__ import annotations

from vanta.tracker import TrackerBase
from vanta.indicators import ATR, SMA, EMA


class VolatilityRegimeTracker(TrackerBase):
    """Switches between mean-reversion and momentum based on volatility."""

    def __init__(self, atr_period: int = 14, fast_period: int = 5, slow_period: int = 20):
        super().__init__()
        self.atr = ATR(period=atr_period)
        self.atr_sma = SMA(period=atr_period)  # average ATR for regime detection
        self.fast_ema = EMA(period=fast_period)
        self.slow_ema = EMA(period=slow_period)
        self.sma = SMA(period=slow_period)
        self._last_close: float | None = None

    def tick(self, data: dict) -> None:
        super().tick(data)
        for candle in data.get("candles_1m", []):
            high = float(candle.get("high", 0))
            low = float(candle.get("low", 0))
            close = float(candle.get("close", 0))

            self.atr.update(high, low, close)
            if self.atr.ready:
                self.atr_sma.update(self.atr.value)
            self.fast_ema.update(close)
            self.slow_ema.update(close)
            self.sma.update(close)
            self._last_close = close

    def predict(self, symbol: str, horizon_seconds: int, step_seconds: int) -> dict:
        if not all([self.atr.ready, self.atr_sma.ready, self.fast_ema.ready,
                     self.slow_ema.ready, self.sma.ready, self._last_close]):
            return {"expected_return": 0.0}

        # Regime: high volatility = ATR above its average
        high_vol = self.atr.value > self.atr_sma.value

        if high_vol:
            # Mean reversion in high vol
            deviation = (self._last_close - self.sma.value) / max(abs(self.sma.value), 1e-9)
            return {"expected_return": -0.5 * deviation}
        else:
            # Momentum in low vol
            spread = (self.fast_ema.value - self.slow_ema.value) / max(abs(self.slow_ema.value), 1e-9)
            return {"expected_return": 0.3 * spread}
