fithia2/libs/intraday/screener.py

"""Two-phase intraday data loading pipeline.

Phase 1: Fetch daily bars for the full universe → identify candidate ticker-days.
Phase 2: Fetch 5-min intraday bars for candidates (cache-first).

This approach reduces API calls by ~97% vs brute-force (fetching intraday for all stocks every day).
"""
from __future__ import annotations

import asyncio
import sys
from collections import defaultdict
from typing import Callable

import yaml

from libs.intraday.cache import IntradayCache
from libs.intraday.domain import UniverseParams
from libs.oracle_client import OracleClient
from libs.oracle_client.price import PriceService
from libs.oracle_client.screener import ScreenerService


# ── Universe Shortcuts ─────────────────────────────────────────────────────

_UNIVERSE_YAML_MAP = {
    "midlarge": "configs/symbols_midlarge_snapshot_exact.yaml",
    "largecap": "configs/symbols.yaml",
    "midcap": "configs/symbols_midcap.yaml",
}


async def resolve_universe(params: UniverseParams, client: OracleClient) -> list[str]:
    """Resolve the list of ticker symbols based on universe config.

    Supported sources:
    - 'sp500': S&P 500 constituents from Oracle index endpoint
    - 'nasdaq100': Nasdaq 100 constituents from Oracle index endpoint
    - 'midlarge': 971-ticker YAML fallback file
    - 'largecap': 741-ticker YAML fallback file
    - 'midcap': 802-ticker YAML file
    - 'yaml': custom YAML file (requires symbols_file param)
    - 'screener': live screener query

    Returns sorted, deduplicated list of uppercase ticker symbols.
    """
    source = params.source.lower()

    # Index-based sources
    if source in ("sp500", "nasdaq100"):
        data = await client.get(f"/api/v1/stocks/index/{source}")
        tickers = [c["symbol"].upper() for c in data.get("constituents", []) if c.get("symbol")]
        return sorted(set(tickers))

    # YAML file shortcuts
    if source in _UNIVERSE_YAML_MAP:
        yaml_path = _UNIVERSE_YAML_MAP[source]
        return _load_yaml_symbols(yaml_path)

    # Custom YAML file
    if source == "yaml":
        if not params.symbols_file:
            raise ValueError("source='yaml' requires symbols_file to be set")
        return _load_yaml_symbols(params.symbols_file)

    # Live screener query
    if source == "screener":
        screener = ScreenerService(client)
        stocks = await screener.search_all_stocks(
            market_cap_min=params.market_cap_min,
            min_avg_volume=params.avg_volume_min,
            price_min=params.min_price if params.min_price > 0 else None,
        )
        return sorted({s.symbol.upper() for s in stocks if s.symbol})

    raise ValueError(f"Unknown universe source: {source!r}. "
                     "Use: sp500, nasdaq100, midlarge, largecap, midcap, yaml, screener")


def _load_yaml_symbols(path: str) -> list[str]:
    """Load ticker list from a YAML symbols file."""
    with open(path) as f:
        data = yaml.safe_load(f)
    if isinstance(data, list):
        return sorted({str(s).upper() for s in data if s})
    if isinstance(data, dict):
        symbols: list[str] = []
        for key in ("symbols", "existing_only", "screener", "existing-only"):
            if key in data:
                val = data[key]
                if isinstance(val, list):
                    symbols.extend(str(s).upper() for s in val if s)
        if symbols:
            return sorted(set(symbols))
        # Flat dict of symbol -> metadata
        return sorted({str(k).upper() for k in data.keys() if not k.startswith("_")})
    raise ValueError(f"Unexpected YAML format in {path}")


# ── Phase 1: Daily Bar Bulk Fetch ──────────────────────────────────────────


async def fetch_daily_bars_bulk(
    tickers: list[str],
    start_date: str,
    end_date: str,
    client: OracleClient,
    concurrency: int = 20,
    progress_callback: Callable[[int, int], None] | None = None,
) -> dict[str, list[dict]]:
    """Fetch daily OHLCV bars for all tickers in parallel.

    Returns {ticker: [bar_dict, ...]}.
    Tickers with no data or errors are silently omitted.
    """
    svc = PriceService(client)
    semaphore = asyncio.Semaphore(concurrency)
    results: dict[str, list[dict]] = {}
    lock = asyncio.Lock()
    completed = 0

    async def fetch_one(ticker: str) -> None:
        nonlocal completed
        async with semaphore:
            try:
                resp = await svc.get_daily_bars(ticker, start=start_date, end=end_date)
                if resp.bars:
                    bars = [
                        {
                            "date": b.date,
                            "open": b.open,
                            "high": b.high,
                            "low": b.low,
                            "close": b.close,
                            "volume": b.volume,
                        }
                        for b in resp.bars
                    ]
                    async with lock:
                        results[ticker] = bars
            except Exception:
                pass
            finally:
                async with lock:
                    completed += 1
                if progress_callback:
                    progress_callback(completed, len(tickers))

    tasks = [asyncio.create_task(fetch_one(t)) for t in tickers]
    await asyncio.gather(*tasks)
    return results


# ── Phase 2: Pre-Screening ─────────────────────────────────────────────────


def pre_screen_candidates(
    daily_bars: dict[str, list[dict]],
    trading_days: list[str],
    threshold: float = 0.015,
    max_per_day: int = 30,
) -> dict[str, list[str]]:
    """Identify candidate ticker-days using daily bar heuristics.

    A ticker-day is a candidate if (high - open) / open >= threshold.
    This indicates the stock rose significantly above its opening price at some
    point during the day — a necessary condition for being a morning gainer.

    Returns:
        {date: [ticker1, ticker2, ...]} — ranked by (high-open)/open descending,
        capped at max_per_day per date.

    Note: This is a conservative filter. If threshold is too high, some actual
    morning gainers may be missed. 1.5% default catches most meaningful movers.
    """
    # Build per-ticker lookup: {ticker -> {date -> bar}}
    ticker_date_bar: dict[str, dict[str, dict]] = {}
    for ticker, bars in daily_bars.items():
        date_map: dict[str, dict] = {}
        for b in bars:
            d = b["date"][:10]
            date_map[d] = b
        ticker_date_bar[ticker] = date_map

    candidates: dict[str, list[str]] = defaultdict(list)
    day_scores: dict[str, dict[str, float]] = defaultdict(dict)

    for ticker, date_map in ticker_date_bar.items():
        for day in trading_days:
            b = date_map.get(day)
            if not b:
                continue
            open_p = b.get("open", 0)
            high_p = b.get("high", 0)
            if open_p <= 0:
                continue
            score = (high_p - open_p) / open_p
            if score >= threshold:
                candidates[day].append(ticker)
                day_scores[day][ticker] = score

    # Rank and cap per day
    result: dict[str, list[str]] = {}
    for day, tickers in candidates.items():
        ranked = sorted(tickers, key=lambda t: day_scores[day].get(t, 0), reverse=True)
        result[day] = ranked[:max_per_day]

    return result


# ── ORB Pre-Screening ─────────────────────────────────────────────────────


def orb_pre_screen_candidates(
    daily_bars: dict[str, list[dict]],
    trading_days: list[str],
    enrichment: dict[str, dict[str, dict]],
    min_price: float = 10.0,
    min_atr: float = 0.50,
    min_avg_dollar_vol: float = 25_000_000.0,
    max_per_day: int | None = 50,
) -> dict[str, list[str]]:
    """Identify ORB candidate ticker-days using liquidity and volatility filters.

    All filters use PRIOR-day information only (no lookahead):
    - price >= min_price (from today's daily bar open — known at market open)
    - ATR(14) >= min_atr (from prior bars via enrichment)
    - avg_dollar_vol_30d >= min_avg_dollar_vol (from prior bars via enrichment)

    Ranking uses ATR/price (expected daily range as % of price) — a lookahead-free
    proxy for intraday volatility. Higher ATR/price = more likely to produce
    meaningful ORB breakouts.

    Returns:
        {date: [ticker1, ticker2, ...]} — ranked by ATR/price descending,
        capped at max_per_day per date when max_per_day is set.
    """
    from collections import defaultdict

    # Build {ticker -> {date -> bar}} for quick lookup
    ticker_date_bar: dict[str, dict[str, dict]] = {}
    for ticker, bars in daily_bars.items():
        date_map: dict[str, dict] = {}
        for b in bars:
            d = b["date"][:10]
            date_map[d] = b
        ticker_date_bar[ticker] = date_map

    candidates: dict[str, list[str]] = defaultdict(list)
    day_scores: dict[str, dict[str, float]] = defaultdict(dict)

    for ticker, date_map in ticker_date_bar.items():
        for day in trading_days:
            b = date_map.get(day)
            if not b:
                continue

            open_p = b.get("open", 0)
            if open_p < min_price:
                continue

            # All filters from enrichment (computed from PRIOR bars — no lookahead)
            ticker_enrich = enrichment.get(ticker, {}).get(day, {})
            atr = ticker_enrich.get("atr_14")
            avg_dollar_vol = ticker_enrich.get("avg_dollar_vol_30d")

            if atr is None or atr < min_atr:
                continue
            if avg_dollar_vol is None or avg_dollar_vol < min_avg_dollar_vol:
                continue

            # Score: ATR / price = expected daily range % (lookahead-free)
            atr_pct = atr / open_p

            candidates[day].append(ticker)
            day_scores[day][ticker] = atr_pct

    # Rank by ATR/price descending and optionally cap
    result: dict[str, list[str]] = {}
    for day, tickers in candidates.items():
        ranked = sorted(tickers, key=lambda t: day_scores[day].get(t, 0), reverse=True)
        result[day] = ranked if max_per_day is None else ranked[:max_per_day]

    return result


# ── Phase 2: Intraday Bar Bulk Fetch ──────────────────────────────────────


async def fetch_intraday_bulk(
    candidates: dict[str, list[str]],
    client: OracleClient,
    cache: IntradayCache | None,
    interval: str = "5min",
    concurrency: int = 8,
    progress_callback: Callable[[int, int, int, int], None] | None = None,
) -> dict[str, dict[str, list[dict]]]:
    """Fetch 5-min intraday bars for candidate ticker-days (cache-first).

    Returns {date: {ticker: [bar_dict, ...]}}.

    Args:
        candidates: {date -> [ticker, ...]} from pre_screen_candidates().
        client: Oracle API client.
        cache: Intraday Parquet cache.
        interval: Candle interval (default '5min').
        concurrency: Max concurrent API calls (Semaphore). Keep <= 8 to respect rate limits.
        progress_callback: Called with (completed, total, cache_hits, api_calls).
    """
    # Oracle endpoint uses "5m" format; config uses "5min" format
    oracle_interval = interval.replace("min", "m")

    # Build flat list and separate cache hits from misses
    total = sum(len(tickers) for tickers in candidates.values())
    completed = 0
    cache_hits = 0
    api_calls = 0
    lock = asyncio.Lock()
    semaphore = asyncio.Semaphore(concurrency)

    result: dict[str, dict[str, list[dict]]] = defaultdict(dict)
    misses: dict[str, list[str]] = defaultdict(list)  # {day: [ticker, ...]}

    # Phase 1: resolve cache hits synchronously
    for day in sorted(candidates.keys()):
        for ticker in candidates[day]:
            cached = cache.get(ticker, day) if cache else None
            if cached is not None:
                result[day][ticker] = cached
                completed += 1
                cache_hits += 1
            else:
                misses[day].append(ticker)

    if progress_callback:
        progress_callback(completed, total, cache_hits, api_calls)

    # Phase 2: fetch cache misses via multi-ticker endpoint (batched by date, chunk ≤ 75)
    CHUNK = 75

    async def fetch_day_chunk(day: str, chunk: list[str]) -> None:
        nonlocal completed, api_calls
        fetched: dict[str, list[dict]] = {}
        async with semaphore:
            try:
                raw = await client.get(
                    "/api/v1/alpaca/intraday",
                    params={
                        "tickers": ",".join(chunk),
                        "interval": oracle_interval,
                        "start_date": day,
                        "end_date": day,
                    },
                )
                bars_by_ticker = raw.get("bars", {})
                for ticker in chunk:
                    fetched[ticker] = [
                        {
                            "timestamp": b.get("timestamp", ""),
                            "open": float(b.get("open", 0)),
                            "high": float(b.get("high", 0)),
                            "low": float(b.get("low", 0)),
                            "close": float(b.get("close", 0)),
                            "volume": float(b.get("volume", 0)),
                            "vwap": float(b.get("vwap", 0) or 0),
                        }
                        for b in bars_by_ticker.get(ticker, [])
                    ]
            except Exception:
                pass

        async with lock:
            for ticker in chunk:
                bars = fetched.get(ticker, [])
                if cache and bars:
                    cache.put(ticker, day, bars)
                if bars:
                    result[day][ticker] = bars
            completed += len(chunk)
            api_calls += 1
            _c, _t, _ch, _ac = completed, total, cache_hits, api_calls

        if progress_callback:
            progress_callback(_c, _t, _ch, _ac)

    tasks = [
        asyncio.create_task(fetch_day_chunk(day, chunk))
        for day, tickers in sorted(misses.items())
        for chunk in (
            tickers[i : i + CHUNK] for i in range(0, len(tickers), CHUNK)
        )
    ]
    await asyncio.gather(*tasks)
    return dict(result)