src/engine.py

"""Backtest engine.

Two properties are enforced structurally rather than by convention:

**Fills happen at the next bar's open.** A strategy's decision at bar `t` is
shifted forward one bar by `_shift_decisions`, which `run_backtest` always
applies and a strategy cannot bypass, and the fill price handed to vectorbt is
`open`. There is no code path that enters on the signal bar.

**Strategies cannot see the future.** `assert_causal` perturbs the tail of a
price series, re-runs the indicator, and requires that every output before the
perturbation is bit-identical. A strategy that reads bar `t+1` changes its
earlier outputs when the future changes, and is caught. This is a property
test on the function, not a promise in a docstring.
"""

from __future__ import annotations

import hashlib
import json
from dataclasses import dataclass, field, asdict, replace

import numpy as np
import pandas as pd
import vectorbt as vbt

from . import config
from .metrics import Metrics, compute_metrics

# --------------------------------------------------------------------------
# Errors
# --------------------------------------------------------------------------


class EngineError(RuntimeError):
    pass


class LookaheadError(AssertionError):
    """A strategy's output at bar t depended on data after bar t."""


# --------------------------------------------------------------------------
# Configuration
# --------------------------------------------------------------------------


@dataclass(frozen=True)
class Costs:
    """Trading frictions. On by default -- turning them off is how a backtest lies."""

    enabled: bool = True
    commission_bps: float = config.DEFAULT_COMMISSION_BPS  # per side
    slippage_bps: float = config.DEFAULT_SLIPPAGE_BPS
    slippage_model: str = "fixed"  # "fixed" | "volume_scaled"
    volume_scale_k: float = 0.5

    @property
    def commission_rate(self) -> float:
        return (self.commission_bps / 10_000.0) if self.enabled else 0.0

    @property
    def slippage_rate(self) -> float:
        return (self.slippage_bps / 10_000.0) if self.enabled else 0.0


@dataclass(frozen=True)
class Sizing:
    mode: str = "fixed_pct"  # "fixed_pct" | "vol_target" | "fixed_units"
    pct: float = 1.0  # fraction of equity when fixed_pct
    units: float = 1.0  # units when fixed_units
    vol_target_ann: float = 0.15
    vol_lookback: int = 20
    leverage: float = 1.0
    max_position: float = 1.0  # cap as fraction of equity


@dataclass(frozen=True)
class Stops:
    sl_pct: float | None = None
    tp_pct: float | None = None
    trail_pct: float | None = None

    @property
    def any_set(self) -> bool:
        return any(v is not None for v in (self.sl_pct, self.tp_pct, self.trail_pct))


@dataclass(frozen=True)
class Validation:
    """How the period is carved into in-sample and out-of-sample."""

    mode: str = "holdout"  # "none" | "split" | "walk_forward" | "holdout"
    split_frac: float = 0.7
    train_months: int = 12
    test_months: int = 3
    roll_months: int = 3
    holdout_months: int = config.DEFAULT_HOLDOUT_MONTHS


@dataclass(frozen=True)
class BacktestConfig:
    asset: str = "BTC-USD"
    timeframe: str = "1d"
    strategy: str = "Buy & Hold"
    params: dict = field(default_factory=dict)
    costs: Costs = field(default_factory=Costs)
    sizing: Sizing = field(default_factory=Sizing)
    stops: Stops = field(default_factory=Stops)
    validation: Validation = field(default_factory=Validation)
    init_cash: float = config.DEFAULT_INIT_CASH
    fill: str = "next_open"  # the only supported value; see module docstring
    direction: str = "longonly"  # "longonly" | "both"
    # Used for R-multiple when no stop is configured. Documented, not hidden.
    risk_per_trade_pct: float = 0.02

    def fingerprint(self) -> str:
        """Stable hash of everything that can change results -- determinism key."""
        payload = json.dumps(asdict(self), sort_keys=True, default=str)
        return hashlib.sha256(payload.encode()).hexdigest()[:16]


# --------------------------------------------------------------------------
# Strategy output
# --------------------------------------------------------------------------


@dataclass
class StrategyOutput:
    """Decisions aligned to bar *close*. The engine shifts them, not you."""

    entries: pd.Series
    exits: pd.Series
    triggers: pd.Series | None = None
    short_entries: pd.Series | None = None
    short_exits: pd.Series | None = None
    shifted: bool = False

    def validate(self, index: pd.Index) -> None:
        for name in ("entries", "exits", "short_entries", "short_exits"):
            s = getattr(self, name)
            if s is None:
                continue
            if not s.index.equals(index):
                raise EngineError(f"{name} index does not match the price index")


def _shift_decisions(out: StrategyOutput) -> StrategyOutput:
    """Move every decision one bar forward, so it can only act at the next open."""
    if out.shifted:
        raise EngineError("decisions were already shifted; shift exactly once")

    def sh(s):
        if s is None:
            return None
        # Cast before filling: filling an object-dtype series with False and
        # then downcasting is deprecated in pandas and would change behaviour.
        return s.astype("boolean").shift(1).fillna(False).astype(bool)

    def sh_obj(s):
        if s is None:
            return None
        return s.shift(1)

    return StrategyOutput(
        entries=sh(out.entries),
        exits=sh(out.exits),
        triggers=sh_obj(out.triggers),
        short_entries=sh(out.short_entries),
        short_exits=sh(out.short_exits),
        shifted=True,
    )


# --------------------------------------------------------------------------
# Structural causality check
# --------------------------------------------------------------------------


def assert_causal(
    indicator_fn,
    prices: pd.DataFrame,
    *,
    probe_fracs: tuple[float, ...] = (0.4, 0.55, 0.7, 0.85),
    bumps: tuple[float, ...] = (1.35, 0.65),
) -> None:
    """Fail if `indicator_fn`'s output before bar k depends on data after bar k.

    `indicator_fn(prices) -> StrategyOutput | Series | DataFrame`. For each
    probe point the tail of the price frame is scaled, and every output value
    strictly before the probe must be unchanged.

    The tail is scaled both up and down. One direction alone is not enough: a
    boolean comparison that is already `True` can survive an upward bump
    unchanged, so a peeking strategy would slip through. Perturbing both ways
    forces any future-dependent comparison to flip in at least one of them.
    """
    baseline = _as_frame(indicator_fn(prices))
    n = len(prices)

    for frac in probe_fracs:
        k = int(n * frac)
        if k < 2 or k >= n:
            continue
        for bump in bumps:
            perturbed = prices.copy()
            tail = perturbed.index[k:]
            for col in ("open", "high", "low", "close"):
                if col in perturbed.columns:
                    perturbed.loc[tail, col] = perturbed.loc[tail, col] * bump
            if "volume" in perturbed.columns:
                perturbed.loc[tail, "volume"] = perturbed.loc[tail, "volume"] * bump

            probed = _as_frame(indicator_fn(perturbed))
            a, b = baseline.iloc[:k], probed.iloc[:k]
            if not _frames_equal(a, b):
                where = _first_difference(a, b)
                raise LookaheadError(
                    f"output changed before the perturbation point (bar {k}, "
                    f"{prices.index[k]}, tail scaled by {bump}): first divergence "
                    f"at {where}. The strategy is using information from after "
                    f"the bar it acts on."
                )


def _as_frame(obj) -> pd.DataFrame:
    if isinstance(obj, StrategyOutput):
        cols = {"entries": obj.entries, "exits": obj.exits}
        if obj.short_entries is not None:
            cols["short_entries"] = obj.short_entries
        if obj.short_exits is not None:
            cols["short_exits"] = obj.short_exits
        return pd.DataFrame(cols)
    if isinstance(obj, pd.Series):
        return obj.to_frame("value")
    if isinstance(obj, pd.DataFrame):
        return obj
    raise EngineError(f"cannot interpret indicator output of type {type(obj)}")


def _frames_equal(a: pd.DataFrame, b: pd.DataFrame) -> bool:
    if list(a.columns) != list(b.columns) or len(a) != len(b):
        return False
    for c in a.columns:
        x, y = a[c], b[c]
        if x.dtype == bool or y.dtype == bool:
            xa = x.astype('boolean').fillna(False).astype(bool).to_numpy()
            ya = y.astype('boolean').fillna(False).astype(bool).to_numpy()
            if not (xa == ya).all():
                return False
        elif np.issubdtype(x.dtype, np.number):
            if not np.allclose(x.fillna(0).to_numpy(), y.fillna(0).to_numpy(),
                               rtol=1e-12, atol=1e-12):
                return False
        else:
            if not (x.fillna("").astype(str).to_numpy()
                    == y.fillna("").astype(str).to_numpy()).all():
                return False
    return True


def _first_difference(a: pd.DataFrame, b: pd.DataFrame) -> str:
    for c in a.columns:
        x = a[c].fillna(0) if np.issubdtype(a[c].dtype, np.number) else a[c].fillna("")
        y = b[c].fillna(0) if np.issubdtype(b[c].dtype, np.number) else b[c].fillna("")
        diff = np.where(x.to_numpy() != y.to_numpy())[0]
        if len(diff):
            return f"column {c!r}, bar {int(diff[0])} ({a.index[int(diff[0])]})"
    return "unknown"


# --------------------------------------------------------------------------
# Validation plans
# --------------------------------------------------------------------------


@dataclass(frozen=True)
class Window:
    idx: int
    train_start: pd.Timestamp
    train_end: pd.Timestamp
    test_start: pd.Timestamp
    test_end: pd.Timestamp

    def validate(self) -> None:
        if self.train_start > self.train_end:
            raise EngineError(f"window {self.idx}: inverted train range")
        if self.test_start > self.test_end:
            raise EngineError(f"window {self.idx}: inverted test range")
        if self.test_start <= self.train_end:
            raise EngineError(
                f"window {self.idx}: test starts {self.test_start} which is not "
                f"after train end {self.train_end} -- train/test would overlap"
            )


@dataclass
class ValidationPlan:
    mode: str
    index: pd.DatetimeIndex
    holdout_start: pd.Timestamp | None = None
    is_end: pd.Timestamp | None = None
    windows: list[Window] = field(default_factory=list)
    # Honest notes about what this plan could not do, surfaced in the UI.
    notes: list[str] = field(default_factory=list)

    @property
    def holdout_mask(self) -> pd.Series:
        if self.holdout_start is None:
            return pd.Series(False, index=self.index)
        return pd.Series(self.index >= self.holdout_start, index=self.index)

    def selectable_index(self) -> pd.DatetimeIndex:
        """Bars any parameter-selection path is permitted to touch.

        The holdout is excluded here and nowhere else, so a caller cannot fit
        on it by accident.
        """
        if self.holdout_start is None:
            return self.index
        return self.index[self.index < self.holdout_start]

    def segment_of(self, ts: pd.Timestamp) -> str:
        if self.holdout_start is not None and ts >= self.holdout_start:
            return "holdout"
        if self.is_end is not None:
            return "IS" if ts <= self.is_end else "OOS"
        for w in self.windows:
            if w.test_start <= ts <= w.test_end:
                return "OOS"
        return "IS"


def build_validation_plan(index: pd.DatetimeIndex, v: Validation) -> ValidationPlan:
    if len(index) == 0:
        return ValidationPlan(mode=v.mode, index=index)

    start, end = index[0], index[-1]

    # A locked holdout applies to every mode except "none". Walk-forward in
    # particular must not roll its windows into reserved data -- the design
    # shows rolling windows and an "OOS holdout LAST 6MO" side by side.
    holdout_start = None
    if v.mode != "none" and v.holdout_months > 0:
        holdout_start = end - pd.DateOffset(months=v.holdout_months)
        if holdout_start <= start:
            holdout_start = None

    plan = ValidationPlan(mode=v.mode, index=index, holdout_start=holdout_start)

    if v.mode == "none":
        plan.is_end = end
        return plan

    if v.mode == "split":
        cut = int(len(index) * v.split_frac)
        cut = min(max(cut, 1), len(index) - 1)
        plan.is_end = index[cut - 1]
        return plan

    if v.mode == "holdout":
        plan.is_end = (holdout_start - pd.Timedelta(seconds=1)) if holdout_start else end
        return plan

    if v.mode == "walk_forward":
        usable = plan.selectable_index() if holdout_start is not None else index
        if len(usable) == 0:
            return plan
        w_start = usable[0]
        i = 0
        last = usable[-1]
        while True:
            train_start = w_start
            train_end = train_start + pd.DateOffset(months=v.train_months)
            test_start = train_end
            test_end = test_start + pd.DateOffset(months=v.test_months)
            if test_start >= last:
                break
            # train_end is exclusive of the test side: nudge back one instant so
            # the two never share a boundary bar.
            win = Window(
                idx=i,
                train_start=train_start,
                train_end=train_end - pd.Timedelta(seconds=1),
                test_start=test_start,
                test_end=min(test_end - pd.Timedelta(seconds=1), last),
            )
            win.validate()
            plan.windows.append(win)
            i += 1
            w_start = w_start + pd.DateOffset(months=v.roll_months)
            if i > 200:
                break
        if not plan.windows:
            span_days = (usable[-1] - usable[0]).days
            plan.notes.append(
                f"Walk-forward produced no out-of-sample windows: the selectable "
                f"period is {span_days} days, which is shorter than one "
                f"{v.train_months}-month train plus {v.test_months}-month test "
                f"window. There is no OOS result to report — widen the date range "
                f"or shorten the training window."
            )
        return plan

    raise EngineError(f"unknown validation mode {v.mode!r}")


# --------------------------------------------------------------------------
# Results
# --------------------------------------------------------------------------


@dataclass
class WindowResult:
    window: Window
    metrics: Metrics
    equity: pd.Series


@dataclass
class BacktestResult:
    config: BacktestConfig
    equity: pd.Series
    benchmark_equity: pd.Series
    position: pd.Series
    trades: pd.DataFrame
    metrics_all: Metrics
    metrics_is: Metrics
    metrics_oos: Metrics
    metrics_holdout: Metrics | None
    plan: ValidationPlan
    windows: list[WindowResult] = field(default_factory=list)
    prices: pd.DataFrame | None = None
    fingerprint: str = ""
    elapsed_s: float = 0.0

    @property
    def costs_paid(self) -> float:
        return float(self.trades["costs"].sum()) if len(self.trades) else 0.0

    def summary(self) -> dict:
        return {
            "fingerprint": self.fingerprint,
            "trades": int(len(self.trades)),
            "total_return": self.metrics_all.total_return,
            "oos_sharpe": self.metrics_oos.sharpe,
            "max_drawdown": self.metrics_all.max_drawdown,
            "costs_paid": self.costs_paid,
        }


# --------------------------------------------------------------------------
# Sizing / slippage helpers
# --------------------------------------------------------------------------


def _size_args(prices: pd.DataFrame, sizing: Sizing) -> tuple:
    """Return (size, size_type) for vectorbt."""
    if sizing.mode == "fixed_units":
        return sizing.units, "amount"

    if sizing.mode == "vol_target":
        rets = prices["close"].pct_change()
        realised = rets.rolling(sizing.vol_lookback).std()
        bpy = config.bars_per_year("BTC-USD", "1d")  # scale set by caller's tf
        ann = realised * np.sqrt(bpy)
        target = (sizing.vol_target_ann / ann.replace(0.0, np.nan))
        target = target.clip(upper=sizing.max_position * sizing.leverage)
        target = target.fillna(sizing.pct).clip(lower=0.0)
        return target.to_numpy(), "percent"

    frac = min(sizing.pct * sizing.leverage, sizing.max_position * sizing.leverage)
    return float(frac), "percent"


def _slippage_arg(prices: pd.DataFrame, costs: Costs):
    base = costs.slippage_rate
    if base == 0.0:
        return 0.0
    if costs.slippage_model != "volume_scaled":
        return base
    vol = prices["volume"].astype("float64")
    med = float(vol.replace(0.0, np.nan).median())
    if not np.isfinite(med) or med <= 0:
        return base
    # Thin bars cost more to trade than typical ones.
    scale = 1.0 + costs.volume_scale_k * ((med / vol.replace(0.0, np.nan)) - 1.0)
    return (base * scale.clip(lower=0.5, upper=5.0).fillna(1.0)).to_numpy()


# --------------------------------------------------------------------------
# Trade list construction
# --------------------------------------------------------------------------


def _build_trades(
    pf, prices: pd.DataFrame, triggers: pd.Series | None,
    cfg: BacktestConfig, plan: ValidationPlan,
) -> pd.DataFrame:
    """Turn vectorbt's records into the trade list the UI and metrics use."""
    cols = ["id", "entry_ts", "exit_ts", "side", "entry_px", "exit_px", "size",
            "gross_pnl", "costs", "net_pnl", "r_multiple", "mae", "mfe",
            "duration_bars", "trigger", "segment", "status"]
    rec = pf.trades.records_readable
    if rec is None or len(rec) == 0:
        return pd.DataFrame(columns=cols)

    idx = prices.index
    pos_of = {ts: i for i, ts in enumerate(idx)}
    slip = cfg.costs.slippage_rate
    rows = []

    # to_dict("records") preserves the spaced column names; itertuples would
    # rename "Entry Timestamp" to a positional field.
    for n, d in enumerate(rec.to_dict("records")):
        entry_ts = pd.Timestamp(d["Entry Timestamp"])
        exit_ts = pd.Timestamp(d["Exit Timestamp"])
        size = float(d["Size"])
        entry_px = float(d["Avg Entry Price"])
        exit_px = float(d["Avg Exit Price"])
        fees = float(d["Entry Fees"]) + float(d["Exit Fees"])
        net = float(d["PnL"])
        side = str(d["Direction"]).lower()
        status = str(d.get("Status", "Closed"))

        # Slippage is embedded in the fill price rather than booked as a fee, so
        # it is reconstructed from the unslipped reference price. Slippage
        # always works against the trade: a buy fills high, a sell fills low.
        if slip:
            if side == "long":
                raw_entry, raw_exit = entry_px / (1 + slip), exit_px / (1 - slip)
            else:
                raw_entry, raw_exit = entry_px / (1 - slip), exit_px / (1 + slip)
            slip_cost = slip * size * (raw_entry + raw_exit)
        else:
            slip_cost = 0.0
        costs_total = fees + slip_cost
        gross = net + costs_total

        i0, i1 = pos_of.get(entry_ts), pos_of.get(exit_ts)
        mae = mfe = float("nan")
        if i0 is not None and i1 is not None and i1 >= i0:
            window = prices.iloc[i0:i1 + 1]
            lo, hi = float(window["low"].min()), float(window["high"].max())
            if side == "long":
                mae = lo / entry_px - 1.0
                mfe = hi / entry_px - 1.0
            else:
                mae = 1.0 - hi / entry_px
                mfe = 1.0 - lo / entry_px

        if cfg.stops.sl_pct:
            risk = cfg.stops.sl_pct * entry_px * size
        else:
            risk = cfg.risk_per_trade_pct * entry_px * size
        r_mult = (net / risk) if risk else float("nan")

        trig = ""
        if triggers is not None and entry_ts in triggers.index:
            v = triggers.loc[entry_ts]
            trig = "" if (v is None or (isinstance(v, float) and np.isnan(v))) else str(v)

        rows.append({
            "id": n + 1,
            "entry_ts": entry_ts, "exit_ts": exit_ts, "side": side,
            "entry_px": entry_px, "exit_px": exit_px, "size": size,
            "gross_pnl": gross, "costs": costs_total, "net_pnl": net,
            "r_multiple": r_mult, "mae": mae, "mfe": mfe,
            "duration_bars": (i1 - i0) if (i0 is not None and i1 is not None) else np.nan,
            "trigger": trig,
            "segment": plan.segment_of(entry_ts),
            "status": status,
        })

    return pd.DataFrame(rows, columns=cols)


# --------------------------------------------------------------------------
# The run
# --------------------------------------------------------------------------


def run_backtest(
    prices: pd.DataFrame,
    output: StrategyOutput,
    cfg: BacktestConfig,
    *,
    bars_per_year: float | None = None,
) -> BacktestResult:
    """Simulate `output` on `prices` under `cfg`.

    `output` must carry decisions aligned to bar close; the shift to next-bar
    execution is applied here and only here.
    """
    import time

    t0 = time.perf_counter()

    if prices.empty:
        raise EngineError("no price data for the requested range")
    for c in ("open", "high", "low", "close"):
        if c not in prices.columns:
            raise EngineError(f"price frame missing {c!r}")
    if cfg.fill != "next_open":
        raise EngineError(
            f"fill={cfg.fill!r} is not supported; next-bar-open execution is "
            "non-negotiable in this engine"
        )

    prices = prices.sort_index()
    output.validate(prices.index)
    shifted = _shift_decisions(output)
    if not shifted.shifted:
        raise EngineError("internal: decisions were not shifted")

    bpy = bars_per_year if bars_per_year is not None else config.bars_per_year(
        cfg.asset, cfg.timeframe
    )
    plan = build_validation_plan(prices.index, cfg.validation)

    size, size_type = _size_args(prices, cfg.sizing)
    slippage = _slippage_arg(prices, cfg.costs)

    kwargs = dict(
        close=prices["close"],
        entries=shifted.entries,
        exits=shifted.exits,
        price=prices["open"],  # next-bar-open execution
        open=prices["open"], high=prices["high"], low=prices["low"],
        fees=cfg.costs.commission_rate,
        slippage=slippage,
        init_cash=cfg.init_cash,
        size=size, size_type=size_type,
        freq=pd.infer_freq(prices.index) or f"{config.TIMEFRAMES[cfg.timeframe].minutes}min",
        direction=cfg.direction,
        seed=0,  # determinism
    )
    if shifted.short_entries is not None:
        kwargs["short_entries"] = shifted.short_entries
        kwargs["short_exits"] = shifted.short_exits
        kwargs["direction"] = "both"
    if cfg.stops.sl_pct is not None:
        kwargs["sl_stop"] = cfg.stops.sl_pct
    if cfg.stops.trail_pct is not None:
        kwargs["sl_stop"] = cfg.stops.trail_pct
        kwargs["sl_trail"] = True
    if cfg.stops.tp_pct is not None:
        kwargs["tp_stop"] = cfg.stops.tp_pct

    pf = vbt.Portfolio.from_signals(**kwargs)

    equity = pf.value()
    try:
        position = pf.asset_flow().cumsum()
    except Exception:
        position = pd.Series(0.0, index=prices.index)

    trades = _build_trades(pf, prices, shifted.triggers, cfg, plan)

    # Benchmark: buy and hold the same asset, same costs-free basis, so the
    # comparison isolates the strategy rather than the fee schedule.
    bench = prices["close"] / prices["close"].iloc[0] * cfg.init_cash

    def slice_metrics(mask: pd.Series, seg: str) -> Metrics:
        eq = equity[mask]
        if len(eq) < 2:
            return Metrics(segment=seg)
        tr = trades[trades["segment"] == seg] if len(trades) else trades
        return compute_metrics(eq, tr, bpy, segment=seg,
                               position=position[mask] if position is not None else None)

    holdout_mask = plan.holdout_mask
    in_holdout = holdout_mask.to_numpy()
    if plan.is_end is not None:
        is_mask = pd.Series((prices.index <= plan.is_end) & ~in_holdout, index=prices.index)
        oos_mask = pd.Series((prices.index > plan.is_end) & ~in_holdout, index=prices.index)
    else:
        seg = pd.Series([plan.segment_of(t) for t in prices.index], index=prices.index)
        is_mask = (seg == "IS") & ~in_holdout
        oos_mask = (seg == "OOS") & ~in_holdout

    metrics_all = compute_metrics(equity, trades, bpy, segment="all", position=position)
    metrics_is = slice_metrics(is_mask, "IS")
    metrics_oos = slice_metrics(oos_mask, "OOS")
    metrics_holdout = slice_metrics(holdout_mask, "holdout") if in_holdout.any() else None

    windows: list[WindowResult] = []
    for w in plan.windows:
        m = (prices.index >= w.test_start) & (prices.index <= w.test_end)
        eq = equity[m]
        if len(eq) < 2:
            continue
        tr = trades[(trades["entry_ts"] >= w.test_start)
                    & (trades["entry_ts"] <= w.test_end)] if len(trades) else trades
        windows.append(WindowResult(
            window=w,
            metrics=compute_metrics(eq, tr, bpy, segment=f"W{w.idx + 1}"),
            equity=eq,
        ))

    return BacktestResult(
        config=cfg, equity=equity, benchmark_equity=bench, position=position,
        trades=trades, metrics_all=metrics_all, metrics_is=metrics_is,
        metrics_oos=metrics_oos, metrics_holdout=metrics_holdout, plan=plan,
        windows=windows, prices=prices, fingerprint=cfg.fingerprint(),
        elapsed_s=time.perf_counter() - t0,
    )