fithia2/libs/backtest/selector.py

"""Build, rank and filter Candidate objects from raw Parquet row dicts."""
from __future__ import annotations

import datetime as dt
from typing import Any
from zoneinfo import ZoneInfo

from libs.backtest.domain import Candidate, SignalConfig, UniverseConfig
from libs.common.logging import get_logger

logger = get_logger(__name__)
_UTC = ZoneInfo("UTC")


def build_candidate(row: dict[str, Any]) -> Candidate | None:
    """Build a Candidate from a raw Parquet row dict.

    Returns None (logged as skip) if:
    - event_timestamp is null/missing
    - entry_price_est is null/zero
    - execution_date is null/missing
    """
    event_id = row.get("event_id", "")

    # Strict: no silent substitution for event_timestamp
    raw_ts = row.get("event_timestamp")
    if raw_ts is None:
        logger.warning("skip_candidate_no_timestamp", event_id=event_id)
        return None

    # Normalise to timezone-aware datetime
    if isinstance(raw_ts, str):
        try:
            event_timestamp = dt.datetime.fromisoformat(raw_ts)
        except ValueError:
            logger.warning("skip_candidate_bad_timestamp", event_id=event_id, raw=raw_ts)
            return None
    elif isinstance(raw_ts, dt.datetime):
        event_timestamp = raw_ts
    else:
        logger.warning("skip_candidate_unknown_timestamp_type", event_id=event_id)
        return None

    if event_timestamp.tzinfo is None:
        event_timestamp = event_timestamp.replace(tzinfo=_UTC)

    # entry_price_est (mapped from Parquet entry_price)
    entry_price_est = row.get("entry_price") or row.get("entry_price_est")
    if not entry_price_est:
        logger.warning("skip_candidate_no_entry_price", event_id=event_id)
        return None
    entry_price_est = float(entry_price_est)
    if entry_price_est <= 0:
        logger.warning("skip_candidate_zero_entry_price", event_id=event_id)
        return None

    # execution_date (mapped from Parquet entry_date)
    raw_exec_date = row.get("execution_date") or row.get("entry_date")
    if raw_exec_date is None:
        logger.warning("skip_candidate_no_exec_date", event_id=event_id)
        return None
    if isinstance(raw_exec_date, str):
        execution_date = dt.date.fromisoformat(raw_exec_date)
    elif isinstance(raw_exec_date, dt.date):
        execution_date = raw_exec_date
    else:
        logger.warning("skip_candidate_bad_exec_date", event_id=event_id)
        return None

    # reaction_date
    raw_react = row.get("reaction_date")
    if isinstance(raw_react, str):
        reaction_date = dt.date.fromisoformat(raw_react)
    elif isinstance(raw_react, dt.date):
        reaction_date = raw_react
    else:
        reaction_date = execution_date  # fallback: same as execution

    score = float(row.get("score", 0.0))
    avg_dollar_volume = float(row.get("avg_dollar_volume", 0.0))
    atr_14_raw = row.get("atr_14")
    atr_14 = float(atr_14_raw) if atr_14_raw is not None else None

    # Classify score bucket
    score_bucket = _classify_score_bucket(score)

    return Candidate(
        event_id=event_id,
        symbol=str(row.get("symbol", row.get("ticker", ""))),
        issuer_id=row.get("issuer_id"),
        score=score,
        sector=str(row.get("sector") or "UNKNOWN"),
        event_type=str(row.get("event_type", "")),
        event_timestamp=event_timestamp,
        filing_time_bucket=str(row.get("filing_time_bucket", "unknown")),
        reaction_date=reaction_date,
        execution_date=execution_date,
        entry_price_est=entry_price_est,
        avg_dollar_volume=avg_dollar_volume,
        atr_14=atr_14,
        score_bucket=score_bucket,
        features={k: v for k, v in row.items() if k not in _RESERVED_KEYS},
    )


_RESERVED_KEYS = {
    "event_id", "symbol", "ticker", "issuer_id", "score", "sector",
    "event_type", "event_timestamp", "filing_time_bucket", "reaction_date",
    "entry_date", "execution_date", "entry_price", "entry_price_est",
    "avg_dollar_volume", "atr_14", "score_bucket",
}


def _classify_score_bucket(score: float) -> str:
    if score >= 0.8:
        return "high"
    if score >= 0.6:
        return "medium_high"
    if score >= 0.4:
        return "medium"
    if score >= 0.2:
        return "medium_low"
    return "low"


def rank_candidates(candidates: list[Candidate]) -> list[Candidate]:
    """Sort by score DESC, avg_dollar_volume DESC, symbol ASC (stable, deterministic)."""
    return sorted(candidates, key=lambda c: (-c.score, -c.avg_dollar_volume, c.symbol))


def filter_by_universe(
    candidates: list[Candidate],
    config: UniverseConfig,
) -> list[Candidate]:
    """Apply universe filters: min_price, min_avg_dollar_volume, exchange."""
    filtered = []
    for c in candidates:
        if c.entry_price_est < config.min_price:
            continue
        if c.avg_dollar_volume < config.min_avg_dollar_volume:
            continue
        filtered.append(c)
    return filtered


def filter_by_score(
    candidates: list[Candidate],
    score_threshold: float,
) -> list[Candidate]:
    return [c for c in candidates if c.score >= score_threshold]


def truncate_candidates(
    candidates: list[Candidate],
    max_per_day: int,
) -> list[Candidate]:
    return candidates[:max_per_day]


def select_candidates(
    raw_rows: list[dict[str, Any]],
    universe_config: UniverseConfig,
    signal_config: SignalConfig,
) -> list[Candidate]:
    """Full selection pipeline: build → filter → rank → truncate."""
    candidates = []
    for row in raw_rows:
        c = build_candidate(row)
        if c is not None:
            candidates.append(c)

    candidates = filter_by_universe(candidates, universe_config)
    candidates = filter_by_score(candidates, signal_config.score_threshold)
    candidates = rank_candidates(candidates)
    candidates = truncate_candidates(candidates, signal_config.max_candidates_per_day)
    return candidates