sanguo_vnpy_v2/sanguo_backtest/cta_engine.py

"""CTA strategy backtesting engine wrapper using vnpy_ctastrategy.backtesting."""
import sys
import os
import traceback
import uuid
from datetime import datetime
from pathlib import Path

import pandas as pd

# Add vnpy source to path for local development
_VNPY_SRC = os.path.join(os.path.dirname(__file__), "..", "vnpy_v4.4.0")
_VNPY_SRC = os.path.abspath(_VNPY_SRC)
if _VNPY_SRC not in sys.path:
    sys.path.insert(0, _VNPY_SRC)

from sanguo_backtest.result_store import BacktestResult, save_result


# Mock Exchange enum for local use (replaces vnpy.trader.constant.Exchange)
class MockExchange:
    SSE = "SSE"  # Shanghai Stock Exchange
    SZSE = "SZSE"  # Shenzhen Stock Exchange

    class Exchange:
        SSE = "SSE"
        SZSE = "SZSE"

        def __init__(self, value):
            self.value = value

        def __repr__(self):
            return f"Exchange.{self.value}"


Exchange = MockExchange.Exchange


def guess_exchange(symbol: str) -> Exchange:
    """按代码前缀判断交易所：6/68/5x→SSE，0/3/15x→SZSE"""
    if symbol.startswith(("60", "68", "51", "56", "58")):
        return Exchange("SSE")
    if symbol.startswith(("00", "30", "15")):
        return Exchange("SZSE")
    return Exchange("SSE")


def run_cta_backtest(strategy_class, symbol: str, params: dict, start: str, end: str, cfg, db_path: str) -> BacktestResult:
    """
    Run CTA strategy backtest using vnpy_ctastrategy BacktestingEngine.

    Args:
        strategy_class: CTA strategy class to backtest
        symbol: Stock symbol (e.g., "600000")
        params: Strategy parameters dict
        start: Backtest start date (YYYY-MM-DD format)
        end: Backtest end date (YYYY-MM-DD format)
        cfg: Configuration object (may contain data paths)
        db_path: SQLite database path for saving results

    Returns:
        BacktestResult: Result object with backtest statistics and status
    """
    # Generate unique task ID
    task_id = f"cta_{uuid.uuid4().hex[:8]}"

    try:
        # Lazy import of BacktestingEngine (local env may not have vnpy_ctastrategy)
        from vnpy_ctastrategy.backtesting import BacktestingEngine

        # Build vt_symbol for A-shares
        vt_symbol = f"{symbol}.{guess_exchange(symbol).value}"

        # Convert date strings to datetime objects
        start_dt = datetime.strptime(start, "%Y-%m-%d")
        end_dt = datetime.strptime(end, "%Y-%m-%d") if end else None

        # Create and configure backtesting engine
        engine = BacktestingEngine()

        # Set parameters with A-share specific values
        engine.set_parameters(
            vt_symbol=vt_symbol,
            interval="d",  # Interval.DAILY.value — vnpy enum uses "d" not "1d"
            start=start_dt,
            end=end_dt,
            rate=0.001,  # Commission rate (0.1% for A-shares)
            slippage=0,  # No slippage for simplicity
            size=1,  # Contract size (1 for stocks)
            pricetick=0.01,  # Minimum price tick (0.01 yuan for A-shares)
            capital=1_000_000  # Starting capital — 0 causes instant liquidation on first trade
        )

        # Add strategy
        engine.add_strategy(strategy_class, params)

        # Load historical data
        engine.load_data()

        # Run backtesting
        engine.run_backtesting()

        # Calculate statistics — calculate_result() returns a daily DataFrame,
        # calculate_statistics(df) returns the stats dict (sharpe/drawdown/etc.)
        daily_df = engine.calculate_result()
        raw_stats = engine.calculate_statistics(daily_df, output=False) or {}
        # Ensure JSON-serializable (vnpy may include Timestamp / non-numeric values)
        statistics = {
            k: (v if isinstance(v, (int, float, str, bool)) or v is None else str(v))
            for k, v in raw_stats.items()
        }

        # Build equity curve DataFrame (S1.2): engine.get_all_daily_results()
        # returns a list of dicts; keep date + balance for the chart + parquet.
        daily_results = engine.get_all_daily_results()
        if isinstance(daily_results, list) and daily_results:
            equity_df = pd.DataFrame(daily_results)
            cols = [c for c in ("date", "balance") if c in equity_df.columns]
            equity_df = equity_df[cols] if cols else pd.DataFrame()
        else:
            equity_df = pd.DataFrame()

        # Build trades DataFrame (S1.2): engine.trades is dict[vt_tradeid, TradeData].
        trades_dict = engine.trades if isinstance(engine.trades, dict) else {}
        trades_df = pd.DataFrame([
            {
                "datetime": str(t.datetime),
                "direction": str(t.direction),
                "offset": str(t.offset),
                "price": t.price,
                "volume": t.volume,
                "vt_symbol": getattr(t, "vt_symbol", ""),
            }
            for t in trades_dict.values()
        ])

        # Build result object
        result = BacktestResult(
            task_id=task_id,
            type="cta",
            status="done",
            strategy=strategy_class.__name__,
            symbol=symbol,
            params=params,
            start=start,
            end=end,
            statistics=statistics,
            equity_curve=equity_df,
            trades=trades_df,
        )

    except Exception as e:
        # Handle any exceptions and return failed result
        error_msg = f"{type(e).__name__}: {e}\n{traceback.format_exc()}"

        result = BacktestResult(
            task_id=task_id,
            type="cta",
            status="failed",
            strategy=strategy_class.__name__,
            symbol=symbol,
            params=params,
            start=start,
            end=end,
            statistics={},
            equity_curve=None,
            trades=None,
            error_msg=error_msg
        )

    # Save result to database. file_dir = db dir so equity_curve/trades persist
    # to parquet (S1.1) and reload via result.id.
    save_result(result, db_path=db_path, file_dir=os.path.dirname(os.path.abspath(db_path)))

    return result


# Module-level reference for mocking in tests
BacktestingEngine = None  # Will be set when imported inside run_cta_backtest