fithia2/apps/intraday_bt/sweep.py

"""Parameter grid search engine for intraday backtesting.

Data is fetched once; simulations run repeatedly with different params.
288 combinations × ~200 days ≈ 5 minutes total simulation time.
"""
from __future__ import annotations

import itertools
from pathlib import Path
from typing import Any

import yaml

from libs.intraday.domain import (
    IntradayConfig,
    ORBStrategyParams,
    StrategyParams,
    SweepResult,
)
from libs.intraday.metrics import compute_metrics
from libs.intraday.simulator import run_simulation


class SweepConfig:
    """Parsed sweep configuration."""

    def __init__(self, base_config: IntradayConfig, sweep_params: dict[str, list[Any]]) -> None:
        self.base_config = base_config
        self.sweep_params = sweep_params

    @property
    def total_combinations(self) -> int:
        total = 1
        for vals in self.sweep_params.values():
            total *= len(vals)
        return total


def load_sweep_config(sweep_path: str, base_config: IntradayConfig) -> SweepConfig:
    """Load a sweep YAML and merge with the base config."""
    with open(sweep_path) as f:
        raw = yaml.safe_load(f)

    sweep_params: dict[str, list[Any]] = {}
    for key, vals in raw.get("sweep", {}).items():
        if not isinstance(vals, list):
            vals = [vals]
        # Normalize None strings and null values
        normalized = [None if v in (None, "null", "none", "None") else v for v in vals]
        sweep_params[key] = normalized

    return SweepConfig(base_config=base_config, sweep_params=sweep_params)


def generate_combinations(sweep: SweepConfig) -> list[dict[str, Any]]:
    """Generate Cartesian product of all sweep parameters."""
    keys = sorted(sweep.sweep_params.keys())
    values = [sweep.sweep_params[k] for k in keys]
    combos = list(itertools.product(*values))
    return [dict(zip(keys, combo)) for combo in combos]


def apply_overrides(base_config: IntradayConfig, overrides: dict[str, Any]) -> IntradayConfig:
    """Apply parameter overrides to base config, returning a new config.

    Branches on strategy_mode: momentum overrides go to StrategyParams,
    ORB overrides go to ORBStrategyParams.
    """
    if base_config.strategy_mode == "orb":
        orb = base_config.orb_strategy or ORBStrategyParams()
        orb_dict = orb.model_dump()
        orb_fields = set(ORBStrategyParams.model_fields.keys())
        for key, val in overrides.items():
            if key in orb_fields:
                orb_dict[key] = val
        new_orb = ORBStrategyParams(**orb_dict)
        return base_config.model_copy(update={"orb_strategy": new_orb})

    # Momentum mode (default)
    strategy_dict = base_config.strategy.model_dump()
    strategy_fields = set(StrategyParams.model_fields.keys())
    for key, val in overrides.items():
        if key in strategy_fields:
            strategy_dict[key] = val
    new_strategy = StrategyParams(**strategy_dict)
    return base_config.model_copy(update={"strategy": new_strategy})


def run_sweep(
    sweep: SweepConfig,
    all_intraday: dict[str, dict[str, list[dict]]],
    trading_days: list[str],
    progress_callback: Any = None,
    enrichment: dict | None = None,
    momentum_enrichment: dict | None = None,
    vix_by_day: dict[str, float] | None = None,
    ticker_sectors: dict[str, str] | None = None,
    sector_proxy_intraday_by_day: dict[str, dict[str, list[dict]]] | None = None,
) -> list[SweepResult]:
    """Run simulation for each parameter combination.

    Data is pre-fetched and shared across all runs.
    Only the simulation (pure CPU computation) varies per combination.

    Args:
        sweep: SweepConfig with base config and param grid.
        all_intraday: Pre-loaded {date: {ticker: [bars]}} data.
        trading_days: List of dates.
        progress_callback: Optional callable(completed, total) for progress.
        enrichment: Pre-computed daily enrichment (required for ORB strategy).

    Returns:
        List of SweepResult sorted by Sharpe ratio descending.
    """
    combos = generate_combinations(sweep)
    results: list[SweepResult] = []
    is_orb = sweep.base_config.strategy_mode == "orb"

    for i, overrides in enumerate(combos):
        config = apply_overrides(sweep.base_config, overrides)

        if is_orb:
            from libs.intraday.orb_simulator import run_orb_simulation
            day_results = run_orb_simulation(
                all_intraday, trading_days, config.orb_strategy, enrichment or {},
                vix_by_day=vix_by_day,
            )
        else:
            day_results = run_simulation(
                all_intraday,
                trading_days,
                config.strategy,
                daily_enrichment=momentum_enrichment,
                vix_by_day=vix_by_day,
                ticker_sectors=ticker_sectors,
                sector_proxy_intraday_by_day=sector_proxy_intraday_by_day,
            )

        metrics = compute_metrics(day_results, config, run_id=f"sw{i:04d}")
        results.append(SweepResult(params=overrides, metrics=metrics))

        if progress_callback:
            progress_callback(i + 1, len(combos))

    # Sort by Sharpe descending (None treated as -inf)
    results.sort(
        key=lambda r: (r.metrics.sharpe_ratio or float("-inf"), r.metrics.total_return_pct or -999),
        reverse=True,
    )
    return results