quant-trading/backtest_runner.py

"""Shared, feasible replay for the API, precompute and strategy comparisons."""
from __future__ import annotations

import sys
import numpy as np
import pandas as pd

import data_manager
import universe_history as uh
from execution import ENGINE_VERSION, simulate_book
from trader import (ETF_STRATEGY_UNIVERSES, STRATEGY_REGISTRY, RESIZE_BAND, AGENT_REVIEW_STRATEGIES,
                    convert_quote_currency, price_request, strategy_extra_tickers,
                    strategy_fields, strategy_signal_tickers, make_strategy,
                    session_signals, strategy_supports_market, strategy_resize_band)
from universe import UNIVERSES

MIN_TRADE_VALUE = 50.0


def _log(msg):
    print(msg, file=sys.stderr)


def trade_events(diff: pd.DataFrame, threshold: float = 0.01) -> list[dict]:
    rows = []
    for day, row in diff.iterrows():
        buy, sell = row[row > threshold], row[row < -threshold]
        if len(buy) + len(sell):
            rows.append({"date": str(day.date()), "buys": len(buy), "sells": len(sell),
                         "buyTickers": list(buy.nlargest(3).index),
                         "sellTickers": list(sell.nsmallest(3).index)})
    return rows


def order_events(trades: list[dict]) -> list[dict]:
    days = {}
    for t in trades:
        day = days.setdefault(t["date"], {"date": t["date"], "buys": 0, "sells": 0,
                                         "buyTickers": [], "sellTickers": []})
        side = "buy" if t["action"] == "BUY" else "sell"
        day[side + "s"] += 1
        day[side + "Tickers"].append(t["ticker"])
    for day in days.values():
        for key in ("buyTickers", "sellTickers"):
            day[key] = day[key][:3]
    return list(days.values())


def benchmark_curve(prices: pd.Series, dates: pd.DatetimeIndex, capital: float) -> list[dict]:
    series = prices.reindex(dates).ffill().dropna()
    if series.empty:
        return []
    return [{"date": str(d.date()), "value": round(float(v), 2)}
            for d, v in (series / series.iloc[0] * capital).items()]


def _panels(data_market, tickers, fields, refresh):
    filled = (data_manager.update(data_market, tickers, fields=fields) if refresh
              else data_manager.load_fields(data_market, fields))
    raw = data_manager.load_raw(data_market)
    if raw is None:
        raise FileNotFoundError(f"No local prices for {data_market}")
    # Only valuation forward-fills. Indicators and tradability see actual bars.
    filled["close"] = raw.reindex_like(filled["close"])
    return filled


def load_stock_panels(market: str, strategy_names: list[str], years: int | None,
                      *, refresh: bool = True) -> dict:
    u = UNIVERSES[market]
    # Refresh membership explicitly only in the data-loading stage.
    if refresh:
        if market == "us":
            uh.load_sp500_history()
        else:
            uh.load_csi300_history()
    sample = pd.DatetimeIndex([pd.Timestamp.today().normalize()])
    ever = list(uh.pit_membership(market, sample).columns)
    extras = {t for n in strategy_names for t in strategy_extra_tickers(n) + strategy_signal_tickers(n, market)}
    fields = sorted({"close", "open", "volume"} | {f for n in strategy_names for f in strategy_fields(n)})
    panels = _panels(market, sorted(set(ever) | extras | {u['benchmark']}), fields, refresh)
    close = panels["close"]
    member = uh.pit_membership(market, close.index)
    stocks = sorted(set(member.columns) & set(close.columns) - extras - {u['benchmark']})
    observed = close.reindex(columns=member.columns).notna() & member
    coverage = observed.sum(axis=1).div(member.sum(axis=1).replace(0, np.nan))
    return {"panels": panels, "member": member, "stocks": stocks,
            "benchmark": u["benchmark"], "benchmark_label": u["benchmark_label"],
            "coverage": coverage}


def load_etf_panels(market: str, strategy_names: list[str], *, refresh: bool = True) -> dict:
    """Close/open panels in USD for the pure-ETF strategies among ``strategy_names``."""
    request = sorted({t for n in strategy_names for t in ETF_STRATEGY_UNIVERSES[n]} | {"SPY"})
    etfs = _panels("etfs", price_request(market, request), ["close", "open"], refresh)
    return {f: convert_quote_currency(p) for f, p in etfs.items()}


def prepare_strategy(name: str, market: str, panels: dict, cols: list[str], *,
                     member: pd.DataFrame | None = None, years: int | None = None,
                     start=None, end=None, strategy=None) -> dict:
    close = panels["close"]
    if end is not None:
        close = close.loc[:end]
    benchmark = "SPY" if name in ETF_STRATEGY_UNIVERSES else UNIVERSES[market]["benchmark"]
    if benchmark in close:
        close = close.loc[close[benchmark].notna()]
    extra = strategy_extra_tickers(name) + strategy_signal_tickers(name, market)
    if name in ETF_STRATEGY_UNIVERSES:
        # A pure-ETF strategy handed the stock panel finds none of its tickers
        # and would replay as an all-cash book instead of failing.
        missing = sorted(set(ETF_STRATEGY_UNIVERSES[name]) - (set(cols) & set(close.columns)))
        if missing:
            raise ValueError(f"{name} needs its ETF universe; missing {missing}")
    cols = sorted(set(cols + extra) & set(close.columns))
    bars = {f: panels[f].reindex(index=close.index, columns=cols)
            for f in strategy_fields(name)}
    targets = session_signals(market, name, cols, close.index, close,
                              panels["open"].reindex(index=close.index),
                              bars=bars, member=member, strategy=strategy)
    cutoff = pd.Timestamp(start) if start is not None else (
        close.index[-1] - pd.DateOffset(years=years) if years else close.index[0])
    days = close.index[close.index >= cutoff]
    if len(days) < 2:
        raise ValueError("Need at least two replay sessions")
    return {"name": name, "market": market, "targets": targets.loc[days],
            "prices": close.reindex(columns=targets.columns).loc[days],
            "warmup_rows": int((close.index < days[0]).sum()),
            "resize_band": strategy_resize_band(name),
            "previous_prices": close.reindex(columns=targets.columns).ffill().shift(1).loc[days],
            "volume": panels["volume"].reindex(index=days, columns=targets.columns) if "volume" in panels else None,
            # Tickers outside the membership history (ETFs) count as members: never settled.
            "member": (member.reindex(index=days, columns=targets.columns, fill_value=True)
                       if member is not None else None)}


def replay_prepared(plan: dict, capital: float, tx_cost: float, fixed_fee: float,
                    fee_base: float = 0.0, fee_per_share: float = 0.0,
                    *, contributions=None, integer_shares: bool = True,
                    resize_band: float | None = None, cash_returns=None) -> dict:
    result = simulate_book(plan["targets"], plan["prices"], capital, market=plan["market"],
                           previous_prices=plan.get("previous_prices"), volume=plan.get("volume"),
                           member=plan.get("member"),
                           transaction_cost=tx_cost, fixed_fee=fixed_fee,
                           fee_base=fee_base, fee_per_share=fee_per_share,
                           integer_shares=integer_shares,
                           lot_size=100 if plan["market"] == "cn" else 1,
                           min_trade_value=MIN_TRADE_VALUE,
                           resize_band=plan.get("resize_band", RESIZE_BAND) if resize_band is None else resize_band,
                           contributions=contributions, cash_returns=cash_returns)
    # Include the initial capital so day-one commissions count in all metrics.
    anchor = plan["prices"].index[0] - pd.Timedelta(days=1)
    for key in ("equity", "unit_equity"):
        result[key] = pd.concat([pd.Series([capital], index=[anchor]), result[key]])
    result.update(name=plan["name"], trades=order_events(result["trade_log"]),
                  warmup_rows=plan.get("warmup_rows", 0))
    return result


def _run(strategy, panels, cols, capital, tx_cost, fixed_fee, fee_base, fee_per_share,
         member=None, min_trade_value=MIN_TRADE_VALUE):
    # Compatibility for callers that already constructed a strategy.
    from signals import session_targets
    from strategies.base import BarStrategy
    close = panels["close"][cols]
    bars = {f: panels[f][cols] for f in strategy.fields} if isinstance(strategy, BarStrategy) else None
    targets = session_targets(strategy, close, panels["open"][cols], bars, member)
    result = simulate_book(targets, close, capital, transaction_cost=tx_cost, fixed_fee=fixed_fee,
                           fee_base=fee_base, fee_per_share=fee_per_share,
                           min_trade_value=min_trade_value, resize_band=RESIZE_BAND)
    return result["equity"], order_events(result["trade_log"])


def run_market(market: str, strategy_names: list[str], years: int | None, capital: float,
               tx_cost: float, fixed_fee: float, fee_base: float = 0.0,
               fee_per_share: float = 0.0, *, refresh: bool = True,
               start=None, end=None) -> dict | None:
    if any(n in AGENT_REVIEW_STRATEGIES for n in strategy_names):
        raise ValueError("Agent-reviewed strategies are forward-only; use their dated live record")
    invalid = [n for n in strategy_names if not strategy_supports_market(n, market)]
    if invalid:
        raise ValueError(f"Unsupported strategies for {market}: {invalid}")
    fees = dict(tx_cost=tx_cost, fixed_fee=fixed_fee, fee_base=fee_base, fee_per_share=fee_per_share)
    stocks = [n for n in strategy_names if n not in ETF_STRATEGY_UNIVERSES]
    stock = load_stock_panels(market, stocks, years, refresh=refresh) if stocks else None
    etf_names = [n for n in strategy_names if n in ETF_STRATEGY_UNIVERSES]
    etfs = load_etf_panels(market, etf_names, refresh=refresh) if etf_names else None
    histories = []
    for panel, benchmark in ((stock["panels"], stock["benchmark"]) if stock else (None, None),
                             (etfs, "SPY")):
        if panel is not None:
            sessions = panel["close"].index[panel["close"][benchmark].notna()]
            if end is not None:
                sessions = sessions[sessions <= pd.Timestamp(end)]
            if len(sessions) <= 315:
                raise ValueError("Registered replay needs 315 pre-window sessions")
            histories.append(sessions)
    common_end = min(h[-1] for h in histories)
    requested_start = (pd.Timestamp(start) if start is not None else
                       common_end - pd.DateOffset(years=years) if years else min(h[0] for h in histories))
    common_start = max(requested_start, max(h[315] for h in histories))
    runs, errors = [], []
    for name in strategy_names:
        try:
            panels = etfs if name in ETF_STRATEGY_UNIVERSES else stock["panels"]
            cols = ETF_STRATEGY_UNIVERSES[name] if name in ETF_STRATEGY_UNIVERSES else stock["stocks"]
            member = None if name in ETF_STRATEGY_UNIVERSES else stock["member"]
            plan = prepare_strategy(name, market, panels, cols, member=member, start=common_start, end=common_end)
            result = replay_prepared(plan, capital, **fees)
            result["_plan"] = plan
            result["_fees"] = fees
            runs.append(result)
            _log(f"  {name}: {result['equity'].iloc[-1]:,.2f}")
        except (ValueError, KeyError, FileNotFoundError) as exc:
            errors.append({"strategy": name, "error": str(exc)})
            _log(f"  FAILED {name}: {exc}")
    if not runs:
        if errors:
            raise ValueError(str(errors))
        return None
    first, last = runs[0]["equity"].index[[0, -1]]
    bench = stock["benchmark"] if stock else "SPY"
    panels = stock["panels"] if stock else etfs
    dates = runs[0]["equity"].index
    curve = benchmark_curve(panels["close"][bench], dates[1:], capital)
    curve.insert(0, {"date": str(first.date()), "value": capital})
    coverage = stock["coverage"].loc[first:last].dropna() if stock else pd.Series(dtype=float)
    warnings = ["使用复权价格：历史整股、整手及公司行动的模拟仍有近似。",
                ("美股持仓停止报价 5 个交易日且已移出指数后按最后成交价结算；" if market == "us"
                 else "停牌或退市的 A 股持仓保持冻结，不做退市结算；")
                + "缺少市场冲击数据，收益为研究估计。"]
    if market == "cn":
        warnings.append("沪深300历史成分股按月记录；缺少完整历史ST状态及集合竞价数据，涨跌停成交采用保守限制。")
    if common_start > requested_start:
        warnings.append(f"为保留315个交易日的预热历史，统一统计窗口从 {common_start.date()} 开始；更早区间需补齐行情。")
    if len(coverage) and coverage.min() < .99:
        warnings.append(f"历史成分股报价覆盖率最低为 {coverage.min():.1%}，仍存在样本缺失偏差。")
    return {"market": market, "benchmark_label": stock["benchmark_label"] if stock else "SPY",
            "benchmark_curve": curve, "start_date": str(first.date()), "end_date": str(last.date()),
            **fees, "strategy_runs": runs, "engine_version": ENGINE_VERSION,
            "actual_years": round((last - first).days / 365.25, 1),
            "warnings": warnings, "errors": errors,
            "coverage": {"minimum": float(coverage.min()), "median": float(coverage.median())}
                        if len(coverage) else None}