path: root/worker/backtest.py

"""Backtrader runner. User `Strategy` code executes with named feeds and
next-bar fills. Recording happens in two places:
  - a sibling Recorder strategy (never trades) records per-bar cash/equity/closes;
  - a generated subclass of the user's Strategy intercepts notify_order /
    notify_trade (documented hooks) to record orders and fills, then calls the
    user's own hook methods.
"""
from __future__ import annotations

import ast
import math
import resource
import statistics
import time
from pathlib import Path

import backtrader as bt
import pandas as pd

ENGINE_VERSION = bt.__version__
EXECUTION_ASSUMPTIONS = {
    "fill_timing": "signal bar -> order; order fills at NEXT bar open (no same-bar lookahead)",
    "slippage": "percentage of fill price (broker set_slippage_perc)",
    "commission": "percentage of order value, charged at execution",
    "indicator_warmup": "indicators are computed in strategy code; bars before warmup carry no signal",
    "cn_stock_etf_rules": "broker enforces at submission: buys rounded down to lot=100 "
                          "(reject if <100), sells rejected if they would exceed the "
                          "current position (no naked short) or exceed position minus "
                          "same-symbol pending buys (no T+0 sell of same-bar buys). "
                          "With next-bar-open daily fills, a sell submitted on the fill "
                          "day executes next trading day, which is exactly the A-share "
                          "T+1 rule at daily granularity.",
    "unmodeled": "suspension days, price-limit halts, intraday sequencing, "
                 "liquidity/market impact are NOT simulated; no lot-size "
                 "difference between odd lots of pre-held position on sell",
    "etf_rules": "A-share listed ETFs follow the same lot=100 buy rule; ETF sell "
                 "lots in the real market can differ slightly (approximation noted)",
}


class CnDailyRulesBroker(bt.brokers.BackBroker):
    """Broker rules for cn stock/ETF daily simulation, enforced at submission.

    - index feeds (asset_type=index) are nontradable research proxies: any
      order on them is rejected at the BROKER level (they only serve as
      benchmark feeds).
    - buys: size rounded down to a whole lot of 100; rejected if < 100 or not
      an integer number of lots after rounding.
    - sells: rejected when they would exceed the tradable position
      (no naked short; pending same-symbol buys are excluded to prevent T+0
      sell of shares that are only settled by a same-cycle buy).
    """

    def __init__(self, index_feeds=(), rule_assets=(), lot: int = 100):
        super().__init__()
        self.index_feeds = set(index_feeds)
        self.rule_assets = set(rule_assets)
        self.lot = int(lot)
        self.rejections = []  # {symbol, reason}

    def _reject(self, order, reason: str):
        self.rejections.append({"symbol": order.data._name, "reason": reason})
        order.reject()
        self.notify(order)
        return order

    def submit(self, order, check=True):
        feed = order.data._name
        if feed in self.index_feeds:
            return self._reject(order, "index_order_rejected")
        if feed in self.rule_assets:
            size = order.size
            if not float(size).is_integer():
                order.size = int(size)  # floor for buys handled below
            if order.isbuy():
                pending_buy = 0
                for o in (*self.submitted, *self.pending):
                    if o.isbuy() and o.data._name == feed and o.status not in \
                            (o.Rejected, o.Cancelled, o.Completed):
                        pending_buy += abs(int(o.size or 0)) if not o.executed.size else 0
                base = int(size)
                lots = (base + pending_buy) // self.lot * self.lot
                adjusted = max(0, lots - pending_buy)
                if adjusted < self.lot:
                    return self._reject(order, "below_min_lot_100")
                if adjusted != size:
                    order.prereject_size = size  # audit trail of requested size
                    order.size = adjusted
                    order.executed.remsize = adjusted  # _execute reads remsize
            elif order.issell():
                pos = self.getposition(order.data).size
                pending_buy = 0
                for o in (*self.submitted, *self.pending):
                    if o.isbuy() and o.data._name == feed and o.status not in \
                            (o.Rejected, o.Cancelled, o.Completed):
                        pending_buy += abs(int(o.size or 0))
                allowed = max(0, pos - pending_buy)
                if allowed <= 0:
                    return self._reject(order, "no_position_to_sell")
                if abs(order.size) > allowed:
                    order.prereject_size = abs(order.size)
                    order.size = -allowed  # trim to owned shares (no short)
                    order.executed.remsize = order.size
        return super().submit(order, check=check)


def fail(code: str, message: str, details=None, t0: float = 0.0) -> dict:
    return {
        "status": "failed",
        "error": {"code": code, "message": message, "details": details},
        "engine": {"name": "backtrader", "version": ENGINE_VERSION},
        "elapsed_ms": int((time.monotonic() - t0) * 1000),
        "warnings": [],
    }


def _safe_object_path(data_root: str, relpath: str) -> Path:
    root = Path(data_root).resolve()
    p = (root / relpath).resolve()
    if not str(p).startswith(str(root) + "/"):
        raise ValueError(f"path traversal rejected: {relpath}")
    if p.is_symlink():
        raise ValueError(f"symlink rejected in data root: {relpath}")
    if not p.is_file():
        raise FileNotFoundError(relpath)
    return p


def _load_frame(path: Path) -> pd.DataFrame:
    df = pd.read_parquet(path) if path.suffix == ".parquet" else pd.read_csv(path)
    for c in ("date", "open", "high", "low", "close", "volume"):
        if c not in df.columns:
            raise ValueError(f"dataset object {path.name} missing column '{c}'")
    df["date"] = pd.to_datetime(df["date"])
    if df["date"].duplicated().any():
        df = df.drop_duplicates("date", keep="last")
    if not df["date"].is_monotonic_increasing:
        df = df.sort_values("date")
    return df.set_index("date").sort_index()


def _compile_user_strategy(code: str):
    try:
        ast.parse(code)
    except SyntaxError as e:
        raise SyntaxError(f"strategy syntax error at line {e.lineno}: {e.msg}")
    # ONE namespace serves as both exec globals and locals. With separate
    # dicts, module-level imports/definitions land in the locals dict while
    # method bodies resolve names via the globals dict -> NameError in methods
    # (e.g. `import os` at top level, then os.path in __init__).
    # The sandbox is the Docker container itself; do not fake Python globals.
    ns: dict = {"__builtins__": __builtins__}
    exec(compile(code, "<strategy-source>", "exec"), ns)
    cls = ns.get("Strategy")
    if not (isinstance(cls, type) and issubclass(cls, bt.Strategy)):
        raise ValueError("source must define a class named 'Strategy' subclassing backtrader.Strategy")
    return cls


def run_backtest(request: dict) -> dict:
    t0 = time.monotonic()
    cfg = request.get("config") or {}
    manifest = request.get("dataset_manifest") or {}
    warnings = list(manifest.get("warnings") or [])
    if manifest.get("_synthetic"):
        warnings.append("dataset is a synthetic fixture (_synthetic=true); result is for tests only")

    try:
        user_cls = _compile_user_strategy(request["code"])
    except SyntaxError as e:
        return fail("strategy_syntax", str(e), t0=t0)
    except ValueError as e:
        return fail("strategy_invalid", str(e), t0=t0)

    frames = {}
    try:
        for obj in manifest.get("objects", []):
            p = _safe_object_path(request["data_root"], obj["path"])
            frames[p.resolve().name] = _load_frame(p)
        if not frames:
            return fail("data_missing", "dataset manifest has no objects", t0=t0)
    except FileNotFoundError as e:
        return fail("data_missing", f"missing dataset object: {e}", t0=t0)
    except ValueError as e:
        return fail("data_invalid", str(e), t0=t0)

    feed_frames = [(obj, frames[Path(obj["path"]).resolve().name])
                   for obj in manifest["objects"] if Path(obj["path"]).resolve().name in frames]
    symbols = [obj["instrument"]["symbol"] for obj, _ in feed_frames]

    rec = {"orders": [], "fills": [], "closed_trades": [], "rows": []}

    def RecordingStrategy():
        class R(user_cls):
            plotinfo = dict(plot=False, subplot=False)

            def notify_order(self, order):
                rec["orders"].append({
                    "date": self.datas[0].datetime.date(0).strftime("%Y-%m-%d"),
                    "symbol": order.data._name,
                    "side": "buy" if order.isbuy() else "sell",
                    "order_state": order.getstatusname(),
                })
                if order.status == order.Completed:
                    ex = order.executed
                    # turnover: actual executed QUANTITY x executed PRICE.
                    # (Backtrader's `ex.value` for sells reports cost basis, not
                    # executed turnover — honest accounting must use size*price.)
                    rec["fills"].append({
                        "date": self.datas[0].datetime.date(0).strftime("%Y-%m-%d"),
                        "symbol": order.data._name,
                        "side": "buy" if order.isbuy() else "sell",
                        "quantity": abs(float(ex.size)),
                        "price": float(ex.price),
                        "commission": float(ex.comm),
                        "value": abs(float(ex.size) * float(ex.price)),
                    })
                user_cls.notify_order(self, order)

            def notify_trade(self, trade):
                if trade.isclosed:
                    rec["closed_trades"].append({
                        "symbol": trade.data._name,
                        "pnl": float(trade.pnl),
                        "close_date": self.datas[0].datetime.date(0).strftime("%Y-%m-%d"),
                    })
                user_cls.notify_trade(self, trade)
        return R

    class RowsRecorder(bt.Strategy):
        plotinfo = dict(plot=False, subplot=False)

        def next(self):
            closes = {d._name: float(d.close[0]) for d in self.datas if len(d) > 0}
            rec["rows"].append({
                "date": self.datas[0].datetime.date(0).strftime("%Y-%m-%d"),
                "cash": float(self.broker.getcash()),
                "equity": float(self.broker.getvalue()),
                "closes": closes,
                "positions": {d._name: self.getposition(d).size for d in self.datas},
            })

    cerebro = bt.Cerebro(stdstats=False)
    for obj, frame in feed_frames:
        cerebro.adddata(bt.feeds.PandasData(dataname=frame, plot=False,
                                            timeframe=bt.TimeFrame.Days),
                        name=obj["instrument"]["symbol"])
    asset_types = {obj["instrument"]["symbol"]:
                   obj["instrument"].get("asset_type") for obj, _ in feed_frames}
    index_feeds = [s for s, a in asset_types.items() if a == "index"]
    rule_assets = [s for s, a in asset_types.items() if a in ("stock", "etf")]
    broker = CnDailyRulesBroker(index_feeds=index_feeds, rule_assets=rule_assets)
    cerebro.setbroker(broker)
    initial_capital = float(cfg.get("capital", 100000.0))
    cerebro.broker.setcash(initial_capital)
    cerebro.broker.setcommission(commission=float(cfg.get("commission", 0.0003)))
    cerebro.broker.set_slippage_perc(perc=float(cfg.get("slippage", 0.001)))
    cerebro.addstrategy(RowsRecorder)
    cerebro.addstrategy(RecordingStrategy(), **(cfg.get("parameters") or {}))

    try:
        strat_list = cerebro.run(runonce=False, preload=False)
    except Exception as exc:  # user-code runtime errors surface honestly
        return fail("runtime_error", f"{type(exc).__name__}: {exc}", t0=t0)

    reasons = {
        "index_order_rejected": "index feeds are nontradable proxies; broker rejected the order",
        "below_min_lot_100": "buy below minimum lot 100; broker rejected the order",
        "no_position_to_sell": "sell would exceed owned shares (no naked short); broker rejected the order",
    }
    by_reason: dict[tuple[str, str], int] = {}
    for rej in broker.rejections:
        by_reason[(rej["symbol"], rej["reason"])] = by_reason.get((rej["symbol"], rej["reason"]), 0) + 1
    for (sym, reason), n in by_reason.items():
        rec["orders"].append({"date": None, "symbol": sym, "side": None,
                              "order_state": "rejected", "reject_reason": reasons.get(reason, reason),
                              "count": n})
        warnings.append(f"broker rejected {n} order(s) for {sym}: {reasons.get(reason, reason)}")

    assumptions = dict(EXECUTION_ASSUMPTIONS)
    rows = rec["rows"]

    benchmark_symbol = cfg.get("benchmark_symbol")
    bmk_rows = [r["closes"].get(benchmark_symbol) for r in rows if benchmark_symbol]
    bmk_valid = [v for v in bmk_rows if v is not None]
    bmk_base = bmk_valid[0] if bmk_valid else None
    eq, bench_missing = [], False
    for r in rows:
        e = {"date": r["date"], "equity": r["equity"], "cash": r["cash"],
             "closes": r["closes"], "positions": r["positions"]}
        if benchmark_symbol:
            close = r["closes"].get(benchmark_symbol)
            if close is None:
                bench_missing = True
                e["benchmark"] = None
            elif bmk_base:
                # normalized to initial capital so equity and benchmark share a
                # common start (raw price is also available in e["closes"])
                e["benchmark"] = (close / bmk_base) * initial_capital
            else:
                e["benchmark"] = None
        eq.append(e)
    if bmk_rows and benchmark_symbol and bmk_base:
        warnings.append(f"benchmark '{benchmark_symbol}' series normalized to initial capital "
                        f"{initial_capital:g} (raw closes available in equity[].closes)")
    if bench_missing and benchmark_symbol:
        warnings.append(f"benchmark_symbol '{benchmark_symbol}' not present in dataset; benchmark omitted")

    equity_vals = [r["equity"] for r in rows]
    total_return = annual_return = sharpe = max_dd = final = None
    if len(equity_vals) >= 2 and equity_vals[0] > 0:
        final = equity_vals[-1]
        total_return = final / equity_vals[0] - 1.0
        n = len(equity_vals)
        annual_return = (1.0 + total_return) ** (252.0 / n) - 1.0
        peak, max_dd = -float("inf"), 0.0
        for v in equity_vals:
            peak = max(peak, v)
            max_dd = max(max_dd, (peak - v) / peak)
        rets = [b / a - 1.0 for a, b in zip(equity_vals, equity_vals[1:])]
        std = statistics.pstdev(rets)
        sharpe = (statistics.fmean(rets) / std) * (252 ** 0.5) if std > 1e-12 else None
    elif equity_vals:
        final = equity_vals[-1]
    clean = lambda v: v if v is None or math.isfinite(v) else None
    metrics = {
        "total_return": clean(total_return),
        "annual_return": clean(annual_return),
        "max_drawdown": (None if max_dd is None else -max_dd),  # negative loss fraction; 0.0 is valid
        "sharpe": clean(sharpe),
        # trade_count = number of closed round-trip trades (开了又平完成一次), not fills
        "trade_count": len(rec["closed_trades"]),
        "final_equity": clean(final),
    }

    return {
        "status": "succeeded",
        "engine": {"name": "backtrader", "version": ENGINE_VERSION},
        "initial_cash": initial_capital,
        "equity": eq,
        "orders": rec["orders"],
        "trades": rec["fills"],  # executions/fills: {date,symbol,side,quantity,price,commission,value}
        "closed_trades": rec["closed_trades"],
        "metrics": metrics,
        "logs": [
            f"named feeds: {', '.join(symbols)}",
            f"bars in equity series: {len(rows)}",
            f"final equity: {metrics['final_equity']}",
        ],
        "warnings": warnings,
        "assumptions": assumptions,
        "execution_assumptions": EXECUTION_ASSUMPTIONS,
        "elapsed_ms": int((time.monotonic() - t0) * 1000),
        "peak_rss_kb": int(resource.getrusage(resource.RUSAGE_SELF).ru_maxrss),
        "data_manifest_hash": manifest.get("hash"),
        "_synthetic": bool(request.get("_synthetic")) or bool(manifest.get("_synthetic")),
    }