engine.py

# -*- coding: utf-8 -*-
"""Core backtest engine: monthly-rebalanced, cost-aware, look-ahead-safe by construction.

A *strategy* is any callable ``strategy(view) -> dict[symbol, weight]`` that is called once per
rebalance date. ``view`` is a :class:`DataView` exposing ONLY price history strictly BEFORE the
rebalance date — the same information a live trader would have at the open. Weights are held
between rebalances; any residual (1 - sum of weights) sits in the cash symbol if one is given.

Costs are charged on turnover: ``cost_per_side * sum(|weight change|)`` at every rebalance.
"""
from __future__ import annotations

import math

import numpy as np
import pandas as pd


class DataView:
    """A strategy's window onto the data: everything strictly before ``date``, nothing after.

    This is the framework's first honesty guardrail. Strategies receive a view, not the raw
    frame, so using tomorrow's price to pick today's portfolio requires deliberately breaking
    the API — and ``backtest_lab.audit.audit_no_lookahead`` exists to catch exactly that.
    """

    def __init__(self, prices: pd.DataFrame, date: pd.Timestamp):
        self._prices = prices
        self.date = pd.Timestamp(date)

    @property
    def prices(self) -> pd.DataFrame:
        """Price history strictly before the rebalance date."""
        return self._prices[self._prices.index < self.date]

    @property
    def returns(self) -> pd.DataFrame:
        """Daily simple returns strictly before the rebalance date."""
        return self.prices.pct_change().fillna(0.0)


def rebalance_dates(index: pd.DatetimeIndex, start: str | pd.Timestamp) -> list[pd.Timestamp]:
    """First trading day of each month, on or after ``start``."""
    first_of_month = ~pd.Index(pd.PeriodIndex(index, freq="M")).duplicated()
    return [d for d in index[first_of_month] if d >= pd.Timestamp(start)]


class BacktestResult:
    """Bundle of everything a run produces (daily, net of costs)."""

    def __init__(self, returns: pd.Series, weights: pd.DataFrame, gross: pd.Series,
                 turnover: pd.Series):
        self.returns = returns          # daily net strategy returns
        self.weights = weights          # daily weight matrix (incl. cash column if any)
        self.gross = gross              # daily gross risk exposure (ex-cash)
        self.turnover = turnover        # |weight change| summed, per day

    @property
    def equity(self) -> pd.Series:
        return (1.0 + self.returns).cumprod()


def run(prices: pd.DataFrame, strategy, start: str = "2007-01-01",
        cost_per_side: float = 0.0020, cash: str | None = None) -> BacktestResult:
    """Run ``strategy`` over ``prices`` with monthly rebalancing.

    Parameters
    ----------
    prices : wide DataFrame of adjusted closes (columns = symbols, rows = trading days).
    strategy : callable(view) -> dict[symbol, weight]; weights must be >= 0 and sum to <= 1.
        Returning ``{}`` at a rebalance date means "go 100% flat" (all cash), NOT "hold the
        previous weights" — an explicit empty decision resets the book.
    start : first rebalance on/after this date (history before it is signal warm-up).
    cost_per_side : proportional cost charged on every unit of turnover (0.0020 = 20 bps).
    cash : optional symbol that absorbs residual weight (e.g. a T-bill ETF). Must be a column
           of ``prices``. If None, residual weight simply earns zero and is NOT traded.
           NOTE: when set, the cash symbol is a *real, tradeable* instrument — rotating capital
           into/out of it counts toward turnover and is charged ``cost_per_side`` like any leg
           (a T-bill ETF really costs a spread to buy). So a de-risking move with ``cash`` set
           shows more turnover than the same move with ``cash=None``; that is the modeled cost
           of actually buying the cash sleeve, not a discrepancy.
    """
    if cash is not None and cash not in prices.columns:
        raise ValueError(f"cash symbol {cash!r} not in prices columns")
    idx = prices.index
    rets = prices.pct_change().fillna(0.0)
    rebal = set(rebalance_dates(idx, start))
    risk_cols = [c for c in prices.columns if c != cash]

    weights = pd.DataFrame(0.0, index=idx, columns=list(prices.columns))
    current: dict[str, float] = {}
    for d in idx:
        if d in rebal:
            w = strategy(DataView(prices, d)) or {}
            bad = [s for s in w if s not in prices.columns]
            if bad:
                raise ValueError(f"strategy returned unknown symbols {bad}")
            # Reject NaN/inf FIRST: `NaN < -1e-12` and `NaN > 1` are both False, so a non-finite
            # weight would slip past the range checks below and — via pandas' skipna — silently
            # become a fake flat 0%-return curve instead of an error. Errors must never pass silently.
            nonfinite = [s for s, v in w.items() if not math.isfinite(v)]
            if nonfinite:
                raise ValueError(f"strategy returned non-finite weight(s) for {nonfinite}")
            if any(v < -1e-12 for v in w.values()):
                raise ValueError("strategy returned a negative weight (long-only engine)")
            total = float(sum(w.values()))
            if total > 1.0 + 1e-9:
                raise ValueError(f"strategy weights sum to {total:.4f} > 1 (unlevered engine)")
            current = {s: float(v) for s, v in w.items() if s != cash}
        for s in risk_cols:
            weights.at[d, s] = current.get(s, 0.0)
        if cash is not None:
            weights.at[d, cash] = max(0.0, 1.0 - sum(current.values()))

    mask = idx >= pd.Timestamp(start)
    weights = weights[mask]
    rets = rets[mask]
    # shift(1) leaves the first row NaN; fill with 0 so the INITIAL entry is charged too —
    # a backtest that gets its first positions for free is quietly overstating returns.
    turnover = (weights - weights.shift(1).fillna(0.0)).abs().sum(axis=1)
    net = (weights * rets).sum(axis=1) - turnover * cost_per_side
    gross = weights[risk_cols].sum(axis=1)
    return BacktestResult(net, weights, gross, turnover)