backtest.py

"""Event-driven backtester for the momentum / breakout strategy.

Two P&L modes
-------------
'underlying' : trade the stock directly; share count is sized so that a stop-out
               loses `risk_per_trade` of equity.
'option'     : 'buy' one ATM call (long signal) or put (short signal) and mark it
               with Black-Scholes each bar. This approximates the leverage and
               time-decay of an options day trade -- it is an ESTIMATE: it assumes
               a fixed implied vol and does not model bid/ask spread or slippage.

Exits: an ATR stop-loss (optionally trailing), an ATR take-profit, a
max-hold-bars cap, and a hard end-of-day close. Day-trade rule: every position is
closed by the end of its trading day and none is opened on the last bar of a day.
"""
from __future__ import annotations

import math
from dataclasses import dataclass

import numpy as np
import pandas as pd

import config
import sizing
from options_pricing import bs_price
from strategy import generate_signals


@dataclass
class Trade:
    symbol: str
    direction: int            # 1 = long/calls, -1 = short/puts
    instrument: str           # 'shares', 'CALL' or 'PUT'
    entry_time: pd.Timestamp
    exit_time: pd.Timestamp
    entry_underlying: float   # stock price at entry
    exit_underlying: float    # stock price at exit
    qty: float                # shares, or option contracts
    entry_value: float        # dollars deployed (notional, or premium paid)
    pnl: float
    exit_reason: str          # 'stop', 'target', 'max_hold' or 'eod'

    @property
    def return_pct(self) -> float:
        return self.pnl / self.entry_value if self.entry_value else 0.0


def _fill_price(mid: float, side: str, option_params) -> float:
    """Apply bid/ask spread and slippage to a theoretical (mid) option price.

    A buy fills at the ask (mid marked up), a sell at the bid (marked down), by
    half the spread plus slippage. With both pct params at 0 this is a no-op, so
    a frictionless backtest reproduces the pre-cost numbers exactly.

    `limit_fill_frac` models limit-near-mid execution: it scales the friction
    down by that fraction (1.0 = fill at mid, paying nothing). Default 0.0 keeps
    the market-order behaviour.
    """
    edge = (option_params.bid_ask_spread_pct / 2.0 + option_params.slippage_pct)
    edge *= (1.0 - option_params.limit_fill_frac)
    factor = 1.0 + edge if side == "buy" else 1.0 - edge
    return max(mid * factor, 0.0)


def _lagged_entries(signal_arr: np.ndarray, day, lag: int) -> np.ndarray:
    """Shift each non-zero signal `lag` bars forward to model execution delay.

    A shifted entry that would land on a later trading day is dropped (mirrors
    the live entry cutoff). `lag == 0` returns the signals unchanged. No
    look-ahead: bar i+lag only ever reads the signal from the earlier bar i.
    """
    if lag <= 0:
        return signal_arr
    entries = np.zeros_like(signal_arr)
    n = len(signal_arr)
    for i in range(n):
        j = i + lag
        if signal_arr[i] != 0 and j < n and day[j] == day[i] and entries[j] == 0:
            entries[j] = signal_arr[i]
    return entries


def run_backtest(df: pd.DataFrame, symbol: str, mode: str = "underlying",
                 strategy_params=None, risk=None, option_params=None,
                 market_trend=None, signal_fn=None):
    """Run the strategy over `df`. Returns (list[Trade], stats dict).

    `market_trend` (optional) is a +1/-1 broad-market regime series (see
    market_regime.trend_series); when given, entries against the market trend
    (a long while the market is down, or vice versa) are skipped.

    `signal_fn` (optional) swaps the entry-signal generator while keeping the
    entire exit/cost/sizing engine identical. It must have the same contract as
    `strategy.generate_signals` (the default): `(df, strategy_params) -> frame`
    with `high/low/close/atr/signal` columns. This is how the signal-free
    control baselines in `controls.py` run through the same engine as the real
    strategy, so a claimed edge can be compared against entries that contain no
    information.
    """
    strategy_params = strategy_params or config.STRATEGY
    risk = risk or config.RISK
    option_params = option_params or config.OPTIONS
    if mode not in ("underlying", "option"):
        raise ValueError("mode must be 'underlying' or 'option'")

    sig = (signal_fn or generate_signals)(df, strategy_params)
    index = sig.index
    n = len(sig)
    if n == 0:
        return [], summarize(symbol, mode, [], risk.starting_equity)

    # Pull the hot columns into plain arrays -- indexing a DataFrame row by row
    # is far slower, which matters under the optimiser's thousands of sweeps.
    high = sig["high"].to_numpy(dtype=float)
    low = sig["low"].to_numpy(dtype=float)
    close = sig["close"].to_numpy(dtype=float)
    atr_arr = sig["atr"].to_numpy(dtype=float)
    signal_arr = sig["signal"].to_numpy()

    # Broad-market regime gate, aligned to this frame (NaN where unknown -> allow).
    market_arr = (market_trend.reindex(index, method="ffill").to_numpy(dtype=float)
                  if market_trend is not None else None)

    # Last bar of each trading day -- drives the end-of-day close-out rule.
    day = index.date
    is_eod = np.empty(n, dtype=bool)
    is_eod[-1] = True
    is_eod[:-1] = day[1:] != day[:-1]

    # Execution lag: a signal on bar i is acted on `lag` bars later, modelling
    # the gap between a scan printing a signal and the order actually filling.
    entry_signal = _lagged_entries(signal_arr, day, max(risk.execution_lag_bars, 0))

    trades: list[Trade] = []
    equity = risk.starting_equity

    # Drawdown circuit-breaker state (no-op when risk.max_drawdown_halt == 0).
    peak_equity = equity
    dd_halted = False
    dd_resume = (risk.drawdown_resume if risk.drawdown_resume > 0
                 else risk.max_drawdown_halt * 0.5)

    in_position = False
    pos: dict = {}

    for i in range(n):
        # --- manage an open position ---------------------------------------
        if in_position:
            direction = pos["direction"]

            # Trailing stop: ratchet the stop toward price by trail_atr_mult x the
            # entry ATR. pos['trail_ref'] only holds extremes from bars up to i-1,
            # so the stop checked on bar i never peeks at bar i -- no look-ahead.
            if risk.use_trailing_stop:
                trail_dist = risk.trail_atr_mult * pos["entry_atr"]
                if direction == 1:
                    pos["stop"] = max(pos["stop"], pos["trail_ref"] - trail_dist)
                else:
                    pos["stop"] = min(pos["stop"], pos["trail_ref"] + trail_dist)

            exit_price = None
            reason = ""

            # Stop / target are levels on the underlying. If a bar's range spans
            # both, assume the stop triggered first (conservative).
            if direction == 1:
                if low[i] <= pos["stop"]:
                    exit_price, reason = pos["stop"], "stop"
                elif high[i] >= pos["target"]:
                    exit_price, reason = pos["target"], "target"
            else:
                if high[i] >= pos["stop"]:
                    exit_price, reason = pos["stop"], "stop"
                elif low[i] <= pos["target"]:
                    exit_price, reason = pos["target"], "target"

            if exit_price is None and (i - pos["entry_i"]) >= risk.max_hold_bars:
                exit_price, reason = close[i], "max_hold"
            if exit_price is None and is_eod[i]:
                exit_price, reason = close[i], "eod"

            if exit_price is not None:
                if mode == "underlying":
                    pnl = pos["qty"] * (exit_price - pos["entry_underlying"]) * direction
                else:
                    elapsed = (index[i] - pos["entry_time"]).total_seconds()
                    t_remaining = max(pos["t0"] - elapsed / (365 * 24 * 3600), 0.0)
                    mid = bs_price(pos["instrument"], exit_price, pos["strike"],
                                   t_remaining, option_params.risk_free_rate,
                                   option_params.implied_vol)
                    exit_opt = _fill_price(mid, "sell", option_params)
                    pnl = (exit_opt - pos["entry_opt"]) * 100 * pos["qty"]
                    # Per-contract fees are charged on both legs (buy + sell).
                    pnl -= option_params.commission_per_contract * pos["qty"] * 2

                equity += pnl
                peak_equity = max(peak_equity, equity)
                trades.append(Trade(
                    symbol=symbol, direction=direction, instrument=pos["instrument"],
                    entry_time=pos["entry_time"], exit_time=index[i],
                    entry_underlying=pos["entry_underlying"], exit_underlying=exit_price,
                    qty=pos["qty"], entry_value=pos["entry_value"],
                    pnl=pnl, exit_reason=reason,
                ))
                in_position = False
                continue  # no same-bar re-entry

            # Position survives: roll the trailing reference forward with this
            # bar's extreme, then wait -- never stack a second position on top.
            if direction == 1:
                pos["trail_ref"] = max(pos["trail_ref"], high[i])
            else:
                pos["trail_ref"] = min(pos["trail_ref"], low[i])
            continue

        # --- open a new position (only reached while flat) -----------------
        signal = int(entry_signal[i])
        if signal == 0 or is_eod[i]:
            continue
        # Market-wide regime gate: skip entries against the broad-market trend.
        if market_arr is not None and signal * market_arr[i] < 0:
            continue
        # Circuit-breaker: halt new entries while in a deep drawdown.
        drawdown = (peak_equity - equity) / peak_equity if peak_equity else 0.0
        dd_halted = sizing.drawdown_halt(drawdown, dd_halted,
                                         risk.max_drawdown_halt, dd_resume)
        if dd_halted:
            continue
        atr = atr_arr[i]
        if not atr > 0:  # NaN during ATR warm-up, or a degenerate flat bar
            continue

        entry_u = close[i]
        stop = entry_u - signal * risk.stop_atr_mult * atr
        target = entry_u + signal * risk.target_atr_mult * atr

        if mode == "underlying":
            stop_distance = abs(entry_u - stop)
            qty = math.floor((equity * risk.risk_per_trade) / stop_distance)
            if qty <= 0:
                continue
            instrument = "shares"
            entry_value = qty * entry_u
            pos = {"entry_opt": 0.0, "strike": 0.0, "t0": 0.0}
        else:
            instrument = "CALL" if signal == 1 else "PUT"
            strike = float(round(entry_u))
            t0 = option_params.target_dte / 365.0
            mid = bs_price(instrument, entry_u, strike, t0,
                           option_params.risk_free_rate, option_params.implied_vol)
            entry_opt = _fill_price(mid, "buy", option_params)
            # Size the position (risk-based or fixed), then trim to what the live
            # affordability caps allow -- skip only if even one contract is too much.
            base = (sizing.risk_based_contracts(
                        equity, entry_opt, risk.risk_per_trade,
                        option_params.stop_loss_frac)
                    if option_params.size_by_risk else option_params.contracts)
            # Regime-conditional sizing: smaller toward the regime band edges.
            scale = sizing.regime_scale(atr / entry_u,
                                        strategy_params.regime_atr_pct_min,
                                        strategy_params.regime_atr_pct_max,
                                        option_params.regime_size_floor)
            base = max(1, round(base * scale))
            qty = sizing.affordable_contracts(entry_opt, base, equity, equity,
                                              0.0, config.AFFORDABILITY)
            if qty <= 0:
                continue
            entry_value = entry_opt * 100 * qty
            pos = {"entry_opt": entry_opt, "strike": strike, "t0": t0}

        pos.update({
            "direction": signal, "instrument": instrument,
            "entry_i": i, "entry_time": index[i], "entry_underlying": entry_u,
            "stop": stop, "target": target, "entry_atr": atr,
            "trail_ref": high[i] if signal == 1 else low[i],
            "qty": qty, "entry_value": entry_value,
        })
        in_position = True

    stats = summarize(symbol, mode, trades, risk.starting_equity)
    return trades, stats


def summarize(symbol: str, mode: str, trades: list[Trade], starting_equity: float) -> dict:
    """Aggregate performance statistics for a list of trades."""
    n = len(trades)
    wins = [t for t in trades if t.pnl > 0]
    losses = [t for t in trades if t.pnl <= 0]
    total_pnl = sum(t.pnl for t in trades)
    gross_profit = sum(t.pnl for t in wins)
    gross_loss = -sum(t.pnl for t in losses)

    equity, peak, max_dd = starting_equity, starting_equity, 0.0
    for t in trades:
        equity += t.pnl
        peak = max(peak, equity)
        max_dd = max(max_dd, (peak - equity) / peak if peak else 0.0)

    return {
        "symbol": symbol,
        "mode": mode,
        "trades": n,
        "win_rate": len(wins) / n if n else 0.0,
        "total_pnl": total_pnl,
        "return_pct": total_pnl / starting_equity if starting_equity else 0.0,
        "avg_win": gross_profit / len(wins) if wins else 0.0,
        "avg_loss": -gross_loss / len(losses) if losses else 0.0,
        "profit_factor": (gross_profit / gross_loss) if gross_loss else float("inf"),
        "expectancy": total_pnl / n if n else 0.0,
        "max_drawdown": max_dd,
        "ending_equity": starting_equity + total_pnl,
    }