fithia2/libs/backtest/splits.py

"""Walk-forward and analysis split utilities for the backtester."""
from __future__ import annotations

import datetime as dt
from typing import Any

from libs.backtest.calendar import get_trading_days


class WalkForwardWindow:
    """A single walk-forward window with train and test date ranges."""

    def __init__(
        self,
        window_index: int,
        train_start: dt.date,
        train_end: dt.date,
        test_start: dt.date,
        test_end: dt.date,
    ) -> None:
        self.window_index = window_index
        self.train_start = train_start
        self.train_end = train_end
        self.test_start = test_start
        self.test_end = test_end

    def __repr__(self) -> str:
        return (
            f"WalkForwardWindow(idx={self.window_index}, "
            f"train={self.train_start}→{self.train_end}, "
            f"test={self.test_start}→{self.test_end})"
        )


class RobustnessWindow:
    """A single rolling validation window for one horizon."""

    def __init__(
        self,
        window_index: int,
        horizon_days: int,
        start: dt.date,
        end: dt.date,
    ) -> None:
        self.window_index = window_index
        self.horizon_days = horizon_days
        self.start = start
        self.end = end

    def __repr__(self) -> str:
        return (
            f"RobustnessWindow(idx={self.window_index}, "
            f"horizon={self.horizon_days}, "
            f"window={self.start}→{self.end})"
        )


def generate_walk_forward_windows(
    all_dates: list[dt.date],
    train_days: int = 252,
    test_days: int = 63,
    step_days: int | None = None,
) -> list[WalkForwardWindow]:
    """Generate walk-forward windows over a list of trading dates.

    Args:
        all_dates: Sorted list of trading days (ascending).
        train_days: Number of trading days in each train window.
        test_days: Number of trading days in each test window.
        step_days: Number of days to advance between windows (defaults to test_days).

    Returns:
        List of WalkForwardWindow objects.
    """
    if step_days is None:
        step_days = test_days

    windows = []
    idx = 0
    window_index = 0
    while idx + train_days + test_days <= len(all_dates):
        train_slice = all_dates[idx : idx + train_days]
        test_slice = all_dates[idx + train_days : idx + train_days + test_days]
        windows.append(
            WalkForwardWindow(
                window_index=window_index,
                train_start=train_slice[0],
                train_end=train_slice[-1],
                test_start=test_slice[0],
                test_end=test_slice[-1],
            )
        )
        idx += step_days
        window_index += 1
    return windows


def generate_robustness_windows(
    all_dates: list[dt.date],
    horizons_days: list[int],
    step_days: int = 21,
) -> dict[int, list[RobustnessWindow]]:
    """Generate rolling windows for each requested horizon."""
    windows_by_horizon: dict[int, list[RobustnessWindow]] = {}
    if step_days <= 0:
        return windows_by_horizon
    for horizon_days in horizons_days:
        if horizon_days <= 0:
            continue
        windows: list[RobustnessWindow] = []
        idx = 0
        window_index = 0
        while idx + horizon_days <= len(all_dates):
            window_slice = all_dates[idx : idx + horizon_days]
            windows.append(
                RobustnessWindow(
                    window_index=window_index,
                    horizon_days=horizon_days,
                    start=window_slice[0],
                    end=window_slice[-1],
                )
            )
            idx += step_days
            window_index += 1
        if windows:
            windows_by_horizon[horizon_days] = windows
    return windows_by_horizon


def split_by_year(
    dates: list[dt.date],
) -> dict[int, list[dt.date]]:
    """Group trading dates by calendar year."""
    result: dict[int, list[dt.date]] = {}
    for d in dates:
        result.setdefault(d.year, []).append(d)
    return result


def split_by_regime(
    dates: list[dt.date],
    regime_map: dict[dt.date, str],
    default_regime: str = "unknown",
) -> dict[str, list[dt.date]]:
    """Group trading dates by market regime label.

    Args:
        dates: Sorted list of trading dates.
        regime_map: Mapping of date → regime label (e.g. "bull", "bear", "sideways").
        default_regime: Label to use when no regime data is available.

    Returns:
        Dict mapping regime label to list of dates.
    """
    result: dict[str, list[dt.date]] = {}
    for d in dates:
        regime = regime_map.get(d, default_regime)
        result.setdefault(regime, []).append(d)
    return result


def get_date_range_for_split(
    snapshot_manifest: dict[str, Any],
    split_name: str,
) -> tuple[dt.date | None, dt.date | None]:
    """Extract start/end dates for a named split from a snapshot manifest."""
    split_info = snapshot_manifest.get("splits", {}).get(split_name, {})
    start_str = split_info.get("start_date")
    end_str = split_info.get("end_date")
    start = dt.date.fromisoformat(start_str) if start_str else None
    end = dt.date.fromisoformat(end_str) if end_str else None
    return start, end