NewQuant/src/execution_simulator.py

# src/execution_simulator.py (修改部分)
from datetime import datetime
from typing import Dict, List, Optional
import pandas as pd
from .core_data import Order, Trade, Bar, PortfolioSnapshot


class ExecutionSimulator:
    """
    模拟交易执行和管理账户资金、持仓。
    """

    def __init__(self, initial_capital: float,
                 slippage_rate: float = 0.0001,
                 commission_rate: float = 0.0002,
                 initial_positions: Optional[Dict[str, int]] = None):
        """
        Args:
            initial_capital (float): 初始资金。
            slippage_rate (float): 滑点率（相对于成交价格的百分比）。
            commission_rate (float): 佣金率（相对于成交金额的百分比）。
            initial_positions (Optional[Dict[str, int]]): 初始持仓，格式为 {symbol: quantity}。
        """
        self.initial_capital = initial_capital
        self.cash = initial_capital
        self.positions: Dict[str, int] = initial_positions if initial_positions is not None else {}
        # 新增：跟踪持仓的平均成本 {symbol: average_cost}
        self.average_costs: Dict[str, float] = {}
        # 如果有初始持仓，需要设置初始成本（简化为0，或在外部配置）
        if initial_positions:
            for symbol, qty in initial_positions.items():
                # 初始持仓成本，如果需要精确，应该从外部传入
                self.average_costs[symbol] = 0.0  # 简化处理，初始持仓成本为0

        self.slippage_rate = slippage_rate
        self.commission_rate = commission_rate
        self.trade_log: List[Trade] = []  # 存储所有成交记录
        self.pending_orders: Dict[str, Order] = {}  # {order_id: Order_object}
        self._current_time = None

        print(
            f"模拟器初始化：初始资金={self.initial_capital:.2f}, 滑点率={self.slippage_rate}, 佣金率={self.commission_rate}")
        if self.positions:
            print(f"初始持仓：{self.positions}")

    def update_time(self, current_time: datetime):
        """
        更新模拟器的当前时间。
        这个方法由 BacktestEngine 在遍历 K 线时调用。
        """
        self._current_time = current_time

    # --- 新增的公共方法 ---
    def get_current_time(self) -> datetime:
        """
        获取模拟器的当前时间。
        """
        if self._current_time is None:
            # 可以在这里抛出错误或者返回一个默认值，取决于你对未初始化时间的处理
            # 抛出错误可以帮助你发现问题，例如在模拟器时间未设置时就尝试获取
            # raise RuntimeError("Simulator time has not been set. Ensure update_time is called.")
            return None
        return self._current_time

    def _calculate_fill_price(self, order: Order, current_bar: Bar) -> float:
        """
        内部方法：根据订单类型和滑点计算实际成交价格。
        - 市价单通常以当前K线的开盘价成交（考虑滑点）。
        - 限价单判断是否触及限价，如果触及，以限价成交（考虑滑点）。
        """
        fill_price = -1.0  # 默认未成交

        if order.price_type == "MARKET":
            # 市价单通常以开盘价成交，或者根据你的策略需求选择收盘价
            # 这里我们仍然使用开盘价作为市价单的基准成交价
            base_price = current_bar.open

            if order.direction == "BUY" or order.direction == "CLOSE_SHORT":  # 买入或平空，价格向上偏离
                fill_price = base_price * (1 + self.slippage_rate)
            elif order.direction == "SELL" or order.direction == "CLOSE_LONG":  # 卖出或平多，价格向下偏离
                fill_price = base_price * (1 - self.slippage_rate)
            else:  # 默认情况，理论上不应该到这里，因为方向应该明确
                fill_price = base_price  # 不考虑滑点

            # 市价单只要有价格就会成交，无需额外价格区间判断

        elif order.price_type == "LIMIT" and order.limit_price is not None:
            limit_price = order.limit_price
            high = current_bar.high
            low = current_bar.low
            open_price = current_bar.open  # 也可以在限价单成交时用开盘价作为实际成交价，或限价本身

            if order.direction == "BUY" or order.direction == "CLOSE_SHORT":  # 限价买入或限价平空
                # 买入：如果K线最低价 <= 限价，则限价单可能成交
                if low <= limit_price:
                    # 成交价通常是限价，但考虑滑点后可能略高（对买方不利）
                    # 或者可以以 open/low 之间的一个价格成交，这里简化为限价+滑点
                    fill_price_candidate = limit_price * (1 + self.slippage_rate)

                    # 确保成交价不会比当前K线的最低价还低（如果不是在最低点成交）
                    # 也可以简单就返回 limit_price * (1 + self.slippage_rate)

                    # 如果开盘价低于或等于限价，那么通常会以开盘价成交（或者略差）
                    # 否则，如果价格是从上方跌落到限价区，那么会在限价附近成交
                    # 这里简化：如果K线触及限价，则以限价成交（考虑滑点）。
                    # 更精细的模拟会考虑价格穿越顺序 (例如，是否开盘就跳过了限价)

                    # 如果开盘价已经低于或等于限价，那么就以开盘价成交（考虑滑点）
                    if open_price <= limit_price:
                        fill_price = open_price * (1 + self.slippage_rate)
                    else:  # 价格从高处跌落到限价
                        fill_price = limit_price * (1 + self.slippage_rate)

                    # 确保成交价不会超过限价 (虽然加滑点可能会略超，但这是交易成本的一部分)
                    # 这个检查是为了避免逻辑错误，理论上加滑点后应该可以接受比限价略高的价格
                    # 如果你严格要求成交价不能高于限价，则需要移除滑点或者更复杂的逻辑
                    # 这里我们接受加滑点后的价格
                    if fill_price > high and fill_price > limit_price:  # 如果计算出来的成交价高于K线最高价，则可能不合理
                        fill_price = -1.0  # 理论上不该发生，除非滑点过大
                else:
                    return -1.0  # 未触及限价

            elif order.direction == "SELL" or order.direction == "CLOSE_LONG":  # 限价卖出或限价平多
                # 卖出：如果K线最高价 >= 限价，则限价单可能成交
                if high >= limit_price:
                    # 成交价通常是限价，但考虑滑点后可能略低（对卖方不利）
                    fill_price_candidate = limit_price * (1 - self.slippage_rate)

                    # 如果开盘价已经高于或等于限价，那么就以开盘价成交（考虑滑点）
                    if open_price >= limit_price:
                        fill_price = open_price * (1 - self.slippage_rate)
                    else:  # 价格从低处上涨到限价
                        fill_price = limit_price * (1 - self.slippage_rate)

                    # 确保成交价不会低于限价
                    if fill_price < low and fill_price < limit_price:  # 如果计算出来的成交价低于K线最低价，则可能不合理
                        fill_price = -1.0  # 理论上不该发生
                else:
                    return -1.0  # 未触及限价

        # 最后检查成交价是否有效
        if fill_price <= 0:
            return -1.0  # 如果计算出来价格无效，返回未成交

        return fill_price

    def send_order(self, order: Order, current_bar: Bar) -> Optional[Trade]:
        """
        接收策略发出的订单，并模拟执行。
        如果订单未立即成交，则加入待处理订单列表。
        特殊处理：如果 order.direction 是 "CANCEL"，则调用 cancel_order。

        Args:
            order (Order): 待执行的订单对象。
            current_bar (Bar): 当前的Bar数据，用于确定成交价格。

        Returns:
            Optional[Trade]: 如果订单成功执行则返回 Trade 对象，否则返回 None。
        """
        # --- 处理撤单指令 ---
        if order.direction == "CANCEL":
            success = self.cancel_order(order.id)
            if success:
                # print(f"[{current_bar.datetime}] 模拟器: 收到并成功处理撤单指令 for Order ID: {order.id}")
                pass
            else:
                # print(f"[{current_bar.datetime}] 模拟器: 收到撤单指令 for Order ID: {order.id}, 但订单已成交或不存在。")
                pass
            return None  # 撤单操作不返回Trade

        # --- 正常买卖订单处理 ---
        symbol = order.symbol
        volume = order.volume

        # 尝试计算成交价格
        fill_price = self._calculate_fill_price(order, current_bar)

        executed_trade: Optional[Trade] = None
        realized_pnl = 0.0  # 初始化实现盈亏
        is_open_trade = False
        is_close_trade = False

        if fill_price <= 0:  # 表示未成交或不满足限价条件
            if order.price_type == "LIMIT":
                self.pending_orders[order.id] = order
                # print(f'撮合失败,order id:{order.id},fill_price:{fill_price}')
            return None  # 未成交，返回None

        # --- 以下是订单成功成交的逻辑 ---
        trade_value = volume * fill_price
        commission = trade_value * self.commission_rate

        current_position = self.positions.get(symbol, 0)
        current_average_cost = self.average_costs.get(symbol, 0.0)

        actual_direction = order.direction
        if order.direction == "CLOSE_SHORT":
            actual_direction = "BUY"
        elif order.direction == "CLOSE_LONG":
            actual_direction = "SELL"

        is_close_order_intent = (order.direction == "CLOSE_LONG" or
                                 order.direction == "CLOSE_SHORT")

        if actual_direction == "BUY":  # 处理买入 (开多 / 平空)
            # 开多仓或平空仓
            if current_position >= 0:  # 当前持有多仓或无仓位 (开多)
                is_open_trade = not is_close_order_intent  # 如果是平仓意图，则不是开仓交易
                # 更新平均成本 (加权平均)
                new_total_cost = (current_average_cost * current_position) + (fill_price * volume)
                new_total_volume = current_position + volume
                self.average_costs[symbol] = new_total_cost / new_total_volume if new_total_volume > 0 else 0.0
                self.positions[symbol] = new_total_volume
            else:  # 当前持有空仓 (平空)
                is_close_trade = is_close_order_intent  # 这是平仓交易
                # 计算平空盈亏
                pnl_per_share = current_average_cost - fill_price
                realized_pnl = pnl_per_share * volume

                # 更新持仓和成本
                self.positions[symbol] += volume
                if self.positions[symbol] == 0:  # 如果全部平仓
                    del self.positions[symbol]
                    if symbol in self.average_costs: del self.average_costs[symbol]
                elif self.positions[symbol] > 0 and current_position < 0:  # 部分平空，且空头仓位被买平为多头仓位
                    # 这是从空头转为多头的复杂情况。需要重新计算平均成本
                    # 简单处理：将剩余的多头仓位成本设为当前价格
                    self.average_costs[symbol] = fill_price

            if self.cash >= trade_value + commission:
                self.cash -= (trade_value + commission)
            else:
                print(f"[{current_bar.datetime}] 资金不足，无法执行买入 {volume} {symbol}")
                return None


        elif actual_direction == "SELL":  # 处理卖出 (开空 / 平多)
            # 开空仓或平多仓
            if current_position <= 0:  # 当前持有空仓或无仓位 (开空)
                is_open_trade = not is_close_order_intent  # 如果是平仓意图，则不是开仓交易
                # 更新平均成本 (空头成本为负值)
                # 对于空头，平均成本通常是指你卖出开仓的平均价格
                # 这里需要根据你的空头成本计算方式来调整
                # 常见的做法是：总卖出价值 / 总卖出数量
                new_total_value = (current_average_cost * abs(current_position)) + (fill_price * volume)
                new_total_volume = abs(current_position) + volume
                self.average_costs[symbol] = new_total_value / new_total_volume if new_total_volume > 0 else 0.0
                self.positions[symbol] -= volume  # 空头数量增加，持仓量变为负更多

            else:  # 当前持有多仓 (平多)
                is_close_trade = is_close_order_intent  # 这是平仓交易
                # 计算平多盈亏
                pnl_per_share = fill_price - current_average_cost
                realized_pnl = pnl_per_share * volume

                # 更新持仓和成本
                self.positions[symbol] -= volume
                if self.positions[symbol] == 0:  # 如果全部平仓
                    del self.positions[symbol]
                    if symbol in self.average_costs: del self.average_costs[symbol]
                elif self.positions[symbol] < 0 and current_position > 0:  # 部分平多，且多头仓位被卖平为空头仓位
                    # 从多头转为空头的复杂情况
                    self.average_costs[symbol] = fill_price  # 简单将剩余空头仓位成本设为当前价格

            if self.cash >= commission:
                self.cash -= commission
                self.cash += trade_value
            else:
                print(f"[{current_bar.datetime}] 资金不足，无法执行卖出 {volume} {symbol}")
                return None
        else:  # 既不是 BUY 也不是 SELL，且也不是 CANCEL。这可能是未知的 direction
            print(
                f"[{current_bar.datetime}] 模拟器: 收到未知订单方向 {order.direction} for Order ID: {order.id}. 订单未处理。")
            return None

            # 创建 Trade 对象
        executed_trade = Trade(
            order_id=order.id, fill_time=current_bar.datetime, symbol=symbol,
            direction=order.direction,  # 记录原始订单方向 (BUY/SELL/CLOSE_X)
            volume=volume, price=fill_price, commission=commission,
            cash_after_trade=self.cash, positions_after_trade=self.positions.copy(),
            realized_pnl=realized_pnl,  # 填充实现盈亏
            is_open_trade=is_open_trade,
            is_close_trade=is_close_trade
        )
        self.trade_log.append(executed_trade)

        # 如果订单成交，无论它是市价单还是限价单，都从待处理订单中移除
        if order.id in self.pending_orders:
            del self.pending_orders[order.id]

        return executed_trade

    def cancel_order(self, order_id: str) -> bool:
        """
        尝试取消一个待处理订单。

        Args:
            order_id (str): 要取消的订单ID。

        Returns:
            bool: 如果成功取消则返回 True，否则返回 False（例如，订单不存在或已成交）。
        """
        if order_id in self.pending_orders:
            # print(f"订单 {order_id} 已成功取消。")
            del self.pending_orders[order_id]
            return True
        # print(f"订单 {order_id} 不存在或已成交，无法取消。")
        return False

    def get_pending_orders(self) -> Dict[str, Order]:
        """
        获取当前所有待处理订单的副本。
        """
        return self.pending_orders.copy()

    def get_portfolio_value(self, current_bar: Bar) -> float:
        """
        计算当前的投资组合总价值（包括现金和持仓市值）。
        Args:
            current_bar (Bar): 当前的Bar数据，用于计算持仓市值。
        Returns:
            float: 当前的投资组合总价值。
        """
        total_value = self.cash

        # 在单品种场景下，我们假设 self.positions 最多只包含一个品种
        # 并且这个品种就是 current_bar.symbol 所代表的品种
        symbol_in_position = list(self.positions.keys())[0] if self.positions else None

        if symbol_in_position and symbol_in_position == current_bar.symbol:
            quantity = self.positions[symbol_in_position]
            # 持仓市值 = 数量 * 当前市场价格 (current_bar.close)
            # 无论多头(quantity > 0)还是空头(quantity < 0)，这个计算都是正确的
            total_value += quantity * current_bar.open

            # 您也可以选择在这里打印调试信息
            # print(f"  DEBUG Portfolio Value Calculation: Cash={self.cash:.2f}, "
            #       f"Position for {symbol_in_position}: {quantity} @ {current_bar.close:.2f}, "
            #       f"Position Value={quantity * current_bar.close:.2f}, Total Value={total_value:.2f}")

        # 如果没有持仓，或者持仓品种与当前Bar品种不符 (理论上单品种不会发生)
        # 那么 total_value 依然是 self.cash

        return total_value

    def get_current_positions(self) -> Dict[str, int]:
        """
        返回当前持仓字典的副本。
        """
        return self.positions.copy()

    def get_trade_history(self) -> List[Trade]:
        """
        返回所有成交记录的副本。
        """
        return self.trade_log.copy()

    def reset(self, new_initial_capital: float = None, new_initial_positions: Dict[str, int] = None) -> None:
        """
        重置模拟器状态到新的初始条件。
        可以在总回测开始时调用，或在合约切换时调整资金和持仓。
        """
        print("ExecutionSimulator: 重置状态。")
        self.cash = new_initial_capital if new_initial_capital is not None else self.initial_capital
        self.positions = new_initial_positions.copy() if new_initial_positions is not None else {}
        self.trade_history = []
        self.current_orders = {}

    def clear_trade_history(self) -> None:
        """
        清空当前模拟器的交易历史。
        在每个合约片段结束时调用，以便我们只收集当前片段的交易记录。
        """
        print("ExecutionSimulator: 清空交易历史。")
        self.trade_history = []

    def get_average_position_price(self, symbol: str) -> Optional[float]:
        """
        获取指定合约的平均持仓成本。
        如果无持仓或无该合约记录，返回 None。
        """
        # 返回 average_costs 字典中对应 symbol 的值
        # 如果没有持仓或者没有记录，返回 None
        if symbol in self.positions and self.positions[symbol] != 0:
            return self.average_costs.get(symbol)
        return None