quant-trading/precompute_backtest.py

"""
Pre-compute backtest results for all strategies, for multiple capital scenarios.

Each strategy is backtested ONCE per market with that market's live starting
capital (MONITOR_CAPITAL: US $10k, CN ¥70k); the
resulting equity series is then replayed with different cash-injection
schedules so we get multiple scenarios without re-running each strategy.

Usage:
    uv run python precompute_backtest.py [--output PATH] [--years N]

Output JSON shape:
    {
        "computedAt": "...",
        "scenarios": {
            "lump_sum_10k": {label, capital, injectionSchedule, markets: {us, cn}},
            "dca_feb_aug":  {label, capital, injectionSchedule, markets: {us, cn}}
        }
    }
"""

import argparse
import json
import os
import sys
from datetime import datetime

import pandas as pd

from backtest_runner import run_market
from bridge import (
    STRATEGY_META,
    _build_strategy_result,
    _sanitize_for_json,
)
from trader import MONITOR_CAPITAL, STRATEGY_REGISTRY, strategy_supports_market


# Initial capital per market, in local currency — the same amounts the live
# monitor starts with (US $10k, CN ¥70k ≈ $10k). Scenario amounts below are
# multiples of it, so both markets model the same real money.
MARKET_CAPITAL = dict(MONITOR_CAPITAL)


SCENARIOS = [
    {
        "id": "lump_sum_10k",
        "label": "起始资金一次性投入",
        "description": "US $10,000 / CN ¥70,000 一次性投入，无后续追加",
        "injection": None,
    },
    {
        "id": "dca_feb_aug",
        "label": "起始资金 + 每年 2/8 月追加一半",
        "description": "US $10,000 + 每年 2/8 月各 $5,000；CN ¥70,000 + 每年 2/8 月各 ¥35,000",
        "injection": {
            "amount_ratio": 0.5,
            "frequency": "yearly_months",
            "months": [2, 8],
            "startYear": 1,
        },
    },
]


def _injection_for_market(injection: dict | None, capital: float) -> dict | None:
    """Materialize a scenario's injection schedule in a market's currency."""
    if not injection:
        return None
    out = {k: v for k, v in injection.items() if k != "amount_ratio"}
    out["amount"] = round(capital * injection["amount_ratio"], 2)
    return out


def log(msg: str) -> None:
    print(f"[precompute] {msg}", file=sys.stderr)


def _equity_series_from_curve(curve: list[dict]) -> pd.Series:
    """Reconstruct a pandas Series from a precomputed equity curve list."""
    if not curve:
        return pd.Series(dtype=float)
    dates = pd.to_datetime([p["date"] for p in curve])
    values = [float(p["value"]) for p in curve]
    return pd.Series(values, index=dates)


def build_scenario_for_market(
    raw: dict,
    scenario: dict,
    years: int,
    capital: float,
) -> dict:
    """Materialize a scenario (with optional injections) for one market."""
    injection_schedule = _injection_for_market(scenario.get("injection"), capital)
    benchmark_curve = raw["benchmark_curve"]
    benchmark_label = raw["benchmark_label"]

    strategy_results = []
    for run in raw["strategy_runs"]:
        equity = run["equity"]
        try:
            result = _build_strategy_result(
                run["name"], equity, run["trades"], injection_schedule, capital,
            )
            strategy_results.append(result)
        except Exception as e:
            log(f"    [scenario={scenario['id']}] FAILED {run['name']}: {e}")

    # Re-scale benchmark for DCA scenarios so the comparison stays apples-to-apples.
    final_bench_curve = benchmark_curve
    if injection_schedule and benchmark_curve:
        bench_series = _equity_series_from_curve(benchmark_curve)
        bench_returns = bench_series.pct_change().fillna(0.0)
        from bridge import build_injection_schedule
        from research.permanent_yearly import equity_with_cashflows
        contribs, _ = build_injection_schedule(bench_series, injection_schedule)
        bench_with_inj = equity_with_cashflows(bench_returns, contribs, capital)
        final_bench_curve = [
            {"date": d.strftime("%Y-%m-%d"), "value": round(float(v), 2)}
            for d, v in bench_with_inj.items()
        ]

    return {
        "meta": {
            "market": raw["market"],
            "years": years,
            "initialCapital": capital,
            "startDate": raw["start_date"],
            "endDate": raw["end_date"],
            "txCost": raw["tx_cost"],
            "fixedFee": raw["fixed_fee"],
            "injectionSchedule": injection_schedule,
        },
        "strategies": strategy_results,
        "benchmark": {
            "name": benchmark_label,
            "equityCurve": final_bench_curve,
        },
    }


def main():
    parser = argparse.ArgumentParser(description="Pre-compute backtest results")
    parser.add_argument(
        "--output", default=None,
        help="Output directory (default: data/precomputed, read by the web UI)",
    )
    parser.add_argument(
        "--years", type=int, default=10,
        help="Backtest period in years (default: 10)",
    )
    args = parser.parse_args()

    output_dir = args.output or os.path.join(
        os.path.dirname(os.path.abspath(__file__)), "data", "precomputed"
    )
    os.makedirs(output_dir, exist_ok=True)

    us_strategies = []
    cn_strategies = []
    for name in STRATEGY_REGISTRY:
        meta = STRATEGY_META.get(name, {"markets": ["us", "cn"]})
        markets = [m for m in meta.get("markets", ["us", "cn"]) if strategy_supports_market(name, m)]
        if "us" in markets:
            us_strategies.append(name)
        if "cn" in markets:
            cn_strategies.append(name)

    us_capital = MARKET_CAPITAL["us"]
    cn_capital = MARKET_CAPITAL["cn"]

    # Run each market once (cached price data → fast subsequent runs across scenarios).
    log("=" * 60)
    log(f"US MARKET: {len(us_strategies)} strategies @ ${us_capital:,.0f}")
    log("=" * 60)
    # US fee schedule (IBKR-style tiered):
    #   fee = max(2.0, 1.88 + 0.009 * shares)
    us_raw = run_market(
        "us", us_strategies,
        years=args.years, capital=us_capital,
        tx_cost=0.001, fixed_fee=2.0,
        fee_base=1.88, fee_per_share=0.009,
    )

    log("=" * 60)
    log(f"CN MARKET: {len(cn_strategies)} strategies @ ¥{cn_capital:,.0f}")
    log("=" * 60)
    cn_raw = run_market(
        "cn", cn_strategies,
        years=args.years, capital=cn_capital,
        tx_cost=0.001, fixed_fee=5.0,
    )

    # Materialize every scenario × market by applying its injection schedule.
    scenarios_out: dict[str, dict] = {}
    for scenario in SCENARIOS:
        log(f"--- scenario: {scenario['id']} ---")
        scenarios_out[scenario["id"]] = {
            "id": scenario["id"],
            "label": scenario["label"],
            "description": scenario["description"],
            "capital": us_capital,
            "capitalByMarket": {"us": us_capital, "cn": cn_capital},
            "injectionSchedule": _injection_for_market(scenario["injection"], us_capital),
            "markets": {},
        }
        if us_raw:
            scenarios_out[scenario["id"]]["markets"]["us"] = build_scenario_for_market(
                us_raw, scenario, args.years, us_capital
            )
        if cn_raw:
            scenarios_out[scenario["id"]]["markets"]["cn"] = build_scenario_for_market(
                cn_raw, scenario, args.years, cn_capital
            )

    output = {
        "computedAt": datetime.now().isoformat(),
        "config": {"years": args.years, "initialCapital": MARKET_CAPITAL},
        "scenarios": scenarios_out,
    }

    output_path = os.path.join(output_dir, "backtest_results.json")
    with open(output_path, "w") as f:
        json.dump(_sanitize_for_json(output), f, default=str)

    file_size = os.path.getsize(output_path) / (1024 * 1024)
    log("=" * 60)
    log(f"DONE: Written to {output_path} ({file_size:.1f} MB)")
    for sid, sc in scenarios_out.items():
        for mk, mv in sc["markets"].items():
            log(f"  [{sid}] {mk.upper()}: {len(mv['strategies'])} strategies")
    log("=" * 60)


if __name__ == "__main__":
    main()