robinhood/backend/app/services/ai_context_builder.py

from __future__ import annotations

from statistics import mean
from typing import List, Dict, Optional, Iterable

from app.schemas.schemas import AIAnalysisRequest, PriceData, PositionMetrics
from app.streaming.live_store import live_store
from app.services.candlestick_patterns import candlestick_detector


BB_LENGTH = 20
BB_MULT = 2.0
RSI_FAST_LENGTH = 3
ZLSMA_LENGTH = 50
CHAND_LENGTH = 22
CHAND_MULT = 2.0
MAX_PATTERN_SIGNALS = 10


class AIContextBuilder:
    """Builds enriched AIAnalysisRequest payloads from live store data."""

    def __init__(self, default_symbol: str = "XAUUSD", default_timeframe: str = "1m") -> None:
        self.default_symbol = default_symbol
        self.default_timeframe = default_timeframe

    def build_request(self, symbol: Optional[str] = None, timeframe: Optional[str] = None, limit: int = 400) -> AIAnalysisRequest:
        sym = (symbol or self.default_symbol).upper().replace("/", "")
        tf = timeframe or self.default_timeframe
        bars = self._load_bars(sym, tf, limit)
        if not bars:
            raise ValueError(f"No live data available for {sym} {tf}")

        trimmed = bars[-limit:]
        price_data = [
            PriceData(
                time=int(row["time"]),
                open=float(row["open"]),
                high=float(row["high"]),
                low=float(row["low"]),
                close=float(row["close"]),
                volume=float(row.get("volume", 0.0)),
            )
            for row in trimmed
        ]

        indicators = self._build_indicators(trimmed)
        current_price = price_data[-1].close if price_data else float(trimmed[-1]["close"])

        return AIAnalysisRequest(
            price_data=price_data,
            indicators=indicators,
            current_price=current_price,
            symbol=self._format_symbol(sym),
            timeframe=tf,
        )

    def build_metrics(
        self,
        symbol: str,
        timeframe: str,
        price_data: Iterable[PriceData] | Iterable[Dict[str, float]]
    ) -> PositionMetrics:
        rows = list(price_data)
        if not rows:
            raise ValueError("No price data available for metrics")

        def _get(row, key: str) -> float:
            if hasattr(row, key):
                return float(getattr(row, key))
            return float(row[key])

        closes = [float(_get(row, "close")) for row in rows]
        highs = [float(_get(row, "high")) for row in rows]
        lows = [float(_get(row, "low")) for row in rows]

        last = rows[-1]
        prev = rows[-2] if len(rows) > 1 else None

        change = None
        change_pct = None
        if prev is not None:
            prev_close = _get(prev, "close")
            last_close = _get(last, "close")
            change = last_close - prev_close
            if prev_close:
                change_pct = (change / prev_close) * 100

        recent_window = rows[-120:] if len(rows) > 120 else rows
        support_levels = sorted({round(_get(row, "low"), 2) for row in recent_window})[:4]
        resistance_levels = sorted({round(_get(row, "high"), 2) for row in recent_window}, reverse=True)[:4]

        atr14 = self._atr(highs, lows, closes, 14)
        rsi14 = self._rsi(closes, 14)
        ema21 = self._ema(closes, 21)
        sma55 = self._sma(closes, 55)
        sma100 = self._sma(closes, 100)
        sma200 = self._sma(closes, 200)
        volatility = self._volatility(closes, 30)
        momentum = self._momentum(closes, 12)

        current_price = float(_get(last, "close"))
        previous_close = float(_get(prev, "close")) if prev is not None else None

        rsi_fast = self._rsi(closes, RSI_FAST_LENGTH)
        rsi_fast_prev = self._rsi(closes[:-1], RSI_FAST_LENGTH) if len(closes) > RSI_FAST_LENGTH + 1 else None
        bb_basis, bb_upper, bb_lower = self._bollinger_bands(closes, BB_LENGTH, BB_MULT)
        bb_prev = self._bollinger_bands(closes[:-1], BB_LENGTH, BB_MULT) if len(closes) > BB_LENGTH else (None, None, None)
        bb_signal = None
        if (
            bb_basis is not None
            and bb_upper is not None
            and bb_lower is not None
            and rsi_fast is not None
            and rsi_fast_prev is not None
            and bb_prev[0] is not None
        ):
            close_prev = closes[-2]
            bb_upper_prev = bb_prev[1]
            bb_lower_prev = bb_prev[2]
            if (
                rsi_fast_prev < 30
                and close_prev < bb_lower_prev
                and rsi_fast > 30
                and closes[-1] > bb_lower
                and rsi_fast < 50
                and closes[-1] < bb_basis
            ):
                bb_signal = "LONG"
            elif (
                rsi_fast_prev > 70
                and close_prev > bb_upper_prev
                and rsi_fast < 70
                and closes[-1] < bb_upper
                and rsi_fast > 50
                and closes[-1] > bb_basis
            ):
                bb_signal = "SHORT"

        zlsma = self._zlsma(closes, ZLSMA_LENGTH)
        chandelier_long_stop, chandelier_short_stop = self._chandelier_exit(
            highs, lows, closes, CHAND_LENGTH, CHAND_MULT
        )
        chandelier_signal = None
        if chandelier_long_stop is not None or chandelier_short_stop is not None:
            if chandelier_long_stop is not None and chandelier_short_stop is not None:
                if current_price > chandelier_short_stop and (zlsma is None or current_price >= zlsma):
                    chandelier_signal = "LONG"
                elif current_price < chandelier_long_stop and (zlsma is None or current_price <= zlsma):
                    chandelier_signal = "SHORT"
                else:
                    chandelier_signal = "NEUTRAL"
            elif chandelier_long_stop is not None:
                chandelier_signal = "LONG" if current_price > chandelier_long_stop else "SHORT"
            else:
                chandelier_signal = "SHORT" if current_price < chandelier_short_stop else "LONG"

        symbol_fmt = self._format_symbol(symbol)
        timestamp = int(_get(last, "time"))
        pattern_signals = candlestick_detector.analyze(rows)
        recent_pattern_signals = pattern_signals[-MAX_PATTERN_SIGNALS:]

        return PositionMetrics(
            symbol=symbol_fmt,
            timeframe=timeframe,
            timestamp=timestamp,
            current_price=current_price,
            previous_close=previous_close,
            change=round(change, 4) if change is not None else None,
            change_percent=round(change_pct, 4) if change_pct is not None else None,
            high=round(max(_get(row, "high") for row in recent_window), 4) if recent_window else None,
            low=round(min(_get(row, "low") for row in recent_window), 4) if recent_window else None,
            atr14=round(atr14, 4) if atr14 is not None else None,
            rsi14=round(rsi14, 2) if rsi14 is not None else None,
            rsi3=round(rsi_fast, 2) if rsi_fast is not None else None,
            ema21=round(ema21, 4) if ema21 is not None else None,
            sma55=round(sma55, 4) if sma55 is not None else None,
            sma100=round(sma100, 4) if sma100 is not None else None,
            sma200=round(sma200, 4) if sma200 is not None else None,
            bb_basis=round(bb_basis, 4) if bb_basis is not None else None,
            bb_upper=round(bb_upper, 4) if bb_upper is not None else None,
            bb_lower=round(bb_lower, 4) if bb_lower is not None else None,
            bb_signal=bb_signal,
            zlsma=round(zlsma, 4) if zlsma is not None else None,
            chandelier_long_stop=round(chandelier_long_stop, 4) if chandelier_long_stop is not None else None,
            chandelier_short_stop=round(chandelier_short_stop, 4) if chandelier_short_stop is not None else None,
            chandelier_signal=chandelier_signal,
            volatility30=round(volatility * 100, 2) if volatility is not None else None,
            momentum12=round(momentum, 4) if momentum is not None else None,
            support_levels=support_levels,
            resistance_levels=resistance_levels,
            pattern_signals=recent_pattern_signals,
            bars_analyzed=len(rows),
        )

    def _load_bars(self, symbol: str, timeframe: str, limit: int) -> List[Dict[str, float]]:
        bars = live_store.get_history(symbol, timeframe)
        if not bars or len(bars) < limit:
            try:
                live_store.load_historical_data(symbol, timeframe, days_back=30)
                bars = live_store.get_history(symbol, timeframe)
            except Exception:
                pass
        return bars[-limit:] if bars else []

    def _build_indicators(self, bars: List[Dict[str, float]]) -> List[Dict[str, float]]:
        closes = [float(b["close"]) for b in bars]
        highs = [float(b["high"]) for b in bars]
        lows = [float(b["low"]) for b in bars]
        indicators: List[Dict[str, float]] = []

        for window in (8, 21, 55, 100, 200):
            val = self._sma(closes, window)
            if val is not None:
                indicators.append({"name": f"SMA_{window}", "value": round(val, 4)})

        ema21 = self._ema(closes, 21)
        if ema21 is not None:
            indicators.append({"name": "EMA_21", "value": round(ema21, 4)})

        rsi14 = self._rsi(closes, 14)
        if rsi14 is not None:
            indicators.append({"name": "RSI_14", "value": round(rsi14, 2)})

        atr14 = self._atr(highs, lows, closes, 14)
        if atr14 is not None:
            indicators.append({"name": "ATR_14", "value": round(atr14, 4)})

        volatility = self._volatility(closes, 30)
        if volatility is not None:
            indicators.append({"name": "VOLATILITY_30", "value": round(volatility * 100, 2), "unit": "%"})

        momentum = self._momentum(closes, 12)
        if momentum is not None:
            indicators.append({"name": "MOMENTUM_12", "value": round(momentum, 4)})

        rsi_fast = self._rsi(closes, RSI_FAST_LENGTH)
        if rsi_fast is not None:
            indicators.append({"name": f"RSI_{RSI_FAST_LENGTH}", "value": round(rsi_fast, 2)})

        bb_basis, bb_upper, bb_lower = self._bollinger_bands(closes, BB_LENGTH, BB_MULT)
        if bb_basis is not None:
            indicators.append({"name": f"BB_{BB_LENGTH}_BASIS", "value": round(bb_basis, 4)})
            indicators.append({"name": f"BB_{BB_LENGTH}_UPPER", "value": round(bb_upper, 4)})
            indicators.append({"name": f"BB_{BB_LENGTH}_LOWER", "value": round(bb_lower, 4)})

        zlsma = self._zlsma(closes, ZLSMA_LENGTH)
        if zlsma is not None:
            indicators.append({"name": f"ZLSMA_{ZLSMA_LENGTH}", "value": round(zlsma, 4)})

        chandelier_long_stop, chandelier_short_stop = self._chandelier_exit(
            highs, lows, closes, CHAND_LENGTH, CHAND_MULT
        )
        if chandelier_long_stop is not None and chandelier_short_stop is not None:
            indicators.append({"name": f"CHAND_{CHAND_LENGTH}_LONG", "value": round(chandelier_long_stop, 4)})
            indicators.append({"name": f"CHAND_{CHAND_LENGTH}_SHORT", "value": round(chandelier_short_stop, 4)})

        return indicators

    @staticmethod
    def _sma(values: List[float], window: int) -> Optional[float]:
        if len(values) < window:
            return None
        return mean(values[-window:])

    @staticmethod
    def _ema(values: List[float], window: int) -> Optional[float]:
        if len(values) < window:
            return None
        k = 2 / (window + 1)
        ema = mean(values[:window])
        for price in values[window:]:
            ema = price * k + ema * (1 - k)
        return ema

    @staticmethod
    def _rsi(values: List[float], window: int = 14) -> Optional[float]:
        if len(values) <= window:
            return None
        gains = []
        losses = []
        for i in range(1, window + 1):
            change = values[-i] - values[-i - 1]
            if change >= 0:
                gains.append(change)
            else:
                losses.append(abs(change))
        avg_gain = mean(gains) if gains else 0
        avg_loss = mean(losses) if losses else 0
        if avg_loss == 0:
            return 100.0
        rs = avg_gain / avg_loss if avg_loss else 0
        return 100 - (100 / (1 + rs))

    @staticmethod
    def _atr(highs: List[float], lows: List[float], closes: List[float], window: int = 14) -> Optional[float]:
        if len(closes) <= window:
            return None
        true_ranges = []
        for i in range(-window + 1, 0):
            high = highs[i]
            low = lows[i]
            prev_close = closes[i - 1]
            tr = max(high - low, abs(high - prev_close), abs(low - prev_close))
            true_ranges.append(tr)
        return mean(true_ranges) if true_ranges else None

    @staticmethod
    def _volatility(values: List[float], window: int) -> Optional[float]:
        if len(values) < window:
            return None
        subset = values[-window:]
        avg = mean(subset)
        variance = mean([(p - avg) ** 2 for p in subset])
        return (variance ** 0.5) / avg if avg else None

    @staticmethod
    def _momentum(values: List[float], lookback: int = 12) -> Optional[float]:
        if len(values) <= lookback:
            return None
        return values[-1] - values[-lookback - 1]

    @staticmethod
    def _bollinger_bands(values: List[float], length: int, multiplier: float) -> tuple[Optional[float], Optional[float], Optional[float]]:
        if len(values) < length:
            return (None, None, None)
        window = values[-length:]
        basis = mean(window)
        variance = mean([(price - basis) ** 2 for price in window])
        deviation = variance ** 0.5
        upper = basis + multiplier * deviation
        lower = basis - multiplier * deviation
        return (basis, upper, lower)

    @staticmethod
    def _linreg(values: List[float], length: int) -> Optional[float]:
        if len(values) < length:
            return None
        window = values[-length:]
        x = list(range(length))
        sum_x = sum(x)
        sum_y = sum(window)
        sum_x2 = sum(i * i for i in x)
        sum_xy = sum(i * y for i, y in zip(x, window))
        denominator = length * sum_x2 - sum_x ** 2
        if denominator == 0:
            return window[-1]
        slope = (length * sum_xy - sum_x * sum_y) / denominator
        intercept = (sum_y - slope * sum_x) / length
        return intercept + slope * (length - 1)

    @classmethod
    def _zlsma(cls, values: List[float], length: int) -> Optional[float]:
        if len(values) < length:
            return None

        lsma_series: List[float] = []
        for idx in range(length, len(values) + 1):
            segment = values[idx - length : idx]
            lsma_val = cls._linreg(segment, length)
            if lsma_val is not None:
                lsma_series.append(lsma_val)

        if not lsma_series:
            return None

        lsma_last = lsma_series[-1]
        if len(lsma_series) < length:
            return lsma_last

        lsma2_series: List[float] = []
        for idx in range(length, len(lsma_series) + 1):
            segment = lsma_series[idx - length : idx]
            lsma2_val = cls._linreg(segment, length)
            if lsma2_val is not None:
                lsma2_series.append(lsma2_val)

        if not lsma2_series:
            return lsma_last

        lsma2_last = lsma2_series[-1]
        return lsma_last + (lsma_last - lsma2_last)

    @classmethod
    def _chandelier_exit(
        cls, highs: List[float], lows: List[float], closes: List[float], length: int, multiplier: float
    ) -> tuple[Optional[float], Optional[float]]:
        if len(closes) <= length:
            return (None, None)

        recent_high = max(highs[-length:])
        recent_low = min(lows[-length:])
        atr = cls._atr(highs, lows, closes, length)
        if atr is None:
            return (None, None)

        long_stop = recent_high - multiplier * atr
        short_stop = recent_low + multiplier * atr
        return (long_stop, short_stop)

    @staticmethod
    def _format_symbol(symbol: str) -> str:
        if len(symbol) == 6 and symbol.isalpha():
            return f"{symbol[:3]}/{symbol[3:]}"
        return symbol
ai_context_builder = AIContextBuilder()