"""策略基类与回测主循环实现。

BaseStrategy 定义：
- 账户与持仓管理；
- 买卖接口；
- 回测主循环 run_backtest；
- 集成止盈止损和仓位管理钩子；
子类只需实现 on_bar，在每个交易日根据行情数据决策买卖。
"""
from __future__ import annotations

from abc import ABC, abstractmethod
from dataclasses import dataclass
from typing import Dict, List, Optional

import numpy as np
import pandas as pd

from utils.logger import setup_logger

logger = setup_logger(__name__)


@dataclass
class Position:
    """单只股票的持仓信息。"""

    ts_code: str
    quantity: int
    cost: float  # 成本价
    days_held: int = 0


class BaseStrategy(ABC):
    """回测策略抽象基类。

    - 管理资金与持仓；
    - 提供买卖接口；
    - 实现回测主循环；
    - 支持止盈止损和仓位管理；
    - 子类只需实现 on_bar 方法。
    """

    def __init__(
        self,
        initial_cash: float,
        max_positions: int = 2,
        stop_loss: Optional[object] = None,
        take_profit: Optional[object] = None,
        position_sizer: Optional[object] = None,
        date_index_dict: Optional[Dict[str, Dict[str, int]]] = None,
        buy_signal_index: Optional[Dict[str, List[str]]] = None,
    ):
        """初始化策略。
        
        参数：
            initial_cash: 初始资金
            max_positions: 最大持仓数
            stop_loss: 止损管理器（StopLoss实例）
            take_profit: 止盈管理器（StopLoss实例）
            position_sizer: 仓位管理器（PositionSizing实例）
            date_index_dict: 日期索引字典 {ts_code: {date: idx}}
            buy_signal_index: 买入信号索引 {date: [ts_code1, ts_code2, ...]}
        """
        self.initial_cash = float(initial_cash)
        self.cash: float = float(initial_cash)
        self.max_positions: int = int(max_positions)
        self.positions: Dict[str, Position] = {}
        self.equity_curve: List[Dict] = []
        self.trade_count: int = 0  # 交易次数统计（买入+卖出）
        
        # 交易历史记录（用于计算胜率和盈亏比）
        self.trade_history: List[Dict] = []  # 记录每笔完整交易（买入->卖出）
        
        # 风险管理模块（可选）
        self.stop_loss = stop_loss
        self.take_profit = take_profit
        self.position_sizer = position_sizer
        
        # 性能优化：日期索引字典（避免deepcopypandas.DataFrame.attrs）
        self.date_index_dict = date_index_dict if date_index_dict is not None else {}
        self.buy_signal_index = buy_signal_index if buy_signal_index is not None else {}

    # ------------ 需要子类实现的接口 ------------
    @abstractmethod
    def on_bar(self, current_date: str, data_dict: Dict[str, pd.DataFrame]) -> None:
        """每个交易日调用一次，由子类实现交易逻辑。"""

    # ------------ 买卖接口 ------------
    def buy(self, ts_code: str, price: float, quantity: int) -> None:
        """以指定价格和数量买入股票。

        A股交易规则：
        - 最小交易单位为 1 手（100 股）；
        - 买入数量必须为 100 的整数倍。
        """
        # 向下取整到 100 的整数倍
        quantity = (quantity // 100) * 100
        if quantity <= 0:
            return
        cost_amount = float(price) * int(quantity)
        if cost_amount > self.cash:
            return

        self.cash -= cost_amount
        if ts_code in self.positions:
            pos = self.positions[ts_code]
            total_cost = pos.cost * pos.quantity + cost_amount
            total_qty = pos.quantity + quantity
            pos.cost = total_cost / total_qty
            pos.quantity = total_qty
            # 继续累积 days_held
        else:
            self.positions[ts_code] = Position(ts_code=ts_code, quantity=quantity, cost=float(price))

        self.trade_count += 1  # 记录买入交易

    def sell(self, ts_code: str, price: float, quantity: int) -> None:
        """以指定价格和数量卖出股票。"""
        if ts_code not in self.positions:
            return
        pos = self.positions[ts_code]
        sell_qty = min(quantity, pos.quantity)
        if sell_qty <= 0:
            return

        # 计算盈亏（用于更新仓位管理器统计）
        profit_pct = (price - pos.cost) / pos.cost
        profit_amount = (price - pos.cost) * sell_qty
        
        # 记录交易历史
        self.trade_history.append({
            "ts_code": ts_code,
            "buy_price": pos.cost,
            "sell_price": price,
            "quantity": sell_qty,
            "profit_pct": profit_pct,
            "profit_amount": profit_amount,
            "is_win": profit_pct > 0,
        })
        
        self.cash += float(price) * sell_qty
        pos.quantity -= sell_qty
        if pos.quantity == 0:
            del self.positions[ts_code]
            # 清理跟踪止盈状态
            if self.stop_loss is not None:
                self.stop_loss.reset_tracking(ts_code)
            if self.take_profit is not None:
                self.take_profit.reset_tracking(ts_code)
        
        # 更新仓位管理器的交易统计（用于Kelly公式）
        if self.position_sizer is not None and hasattr(self.position_sizer, 'update_trade_stats'):
            self.position_sizer.update_trade_stats(ts_code, profit_pct)

        self.trade_count += 1  # 记录卖出交易

    # ------------ 回测主循环 ------------
    def _calc_market_value(self, current_date: str, data_dict: Dict[str, pd.DataFrame]) -> float:
        """计算当前持仓市值。"""
        total = 0.0
        for ts_code, pos in self.positions.items():
            df = data_dict.get(ts_code)
            if df is None or df.empty:
                continue
            
            # 性能优化：使用预计算的日期索引
            date_index = self.date_index_dict.get(ts_code)
            if date_index is not None and current_date in date_index:
                idx = date_index[current_date]
                price = float(df.iloc[idx]["close"])
            else:
                # 备用方案：若当日无行情，则以最近一条收盘价估值
                price = float(df["close"].iloc[-1])
            
            total += price * pos.quantity
        return total

    def _update_days_held(self, current_date: str, data_dict: Dict[str, pd.DataFrame]) -> None:
        """默认每天将所有持仓的持有天数 +1。

        若子类需要更复杂逻辑，可以覆盖或在 on_bar 中自行管理。
        """
        for pos in self.positions.values():
            pos.days_held += 1
    
    def _check_stop_loss_take_profit(self, current_date: str, data_dict: Dict[str, pd.DataFrame]) -> None:
        """检查所有持仓的止盈止损条件。
        
        在每个交易日开盘前或收盘后调用，自动触发止损/止盈卖出。
        子类可以覆盖此方法以实现自定义逻辑。
        """
        for ts_code in list(self.positions.keys()):
            df = data_dict.get(ts_code)
            if df is None or df.empty:
                continue
            
            # 性能优化：使用预计算的日期索引
            date_index = self.date_index_dict.get(ts_code)
            if date_index is None or current_date not in date_index:
                continue
            
            idx = date_index[current_date]
            row = df.iloc[idx]
            
            pos = self.positions[ts_code]
            close = float(row["close"])
            high = float(row["high"]) if "high" in df.columns else None
            low = float(row["low"]) if "low" in df.columns else None
            
            # 计算ATR（如果需要）
            atr = None
            if self.stop_loss is not None and self.stop_loss.method == "atr":
                from risk.stop_loss import calculate_atr
                atr_series = calculate_atr(df, self.stop_loss.atr_period)
                if not atr_series.empty:
                    atr = float(atr_series.iloc[-1])
            
            # 检查止损
            if self.stop_loss is not None:
                should_exit, reason = self.stop_loss.should_exit(
                    ts_code=ts_code,
                    current_price=close,
                    cost_price=pos.cost,
                    high=high,
                    low=low,
                    atr=atr,
                )
                if should_exit:
                    quantity = pos.quantity
                    self.sell(ts_code, close, quantity)
                    # logger.info(f"{current_date} 止损卖出 {ts_code} 数量 {quantity} 价格 {close:.2f} 原因: {reason}")  # 已移除：止损日志改用进度条
                    continue  # 已卖出，不再检查止盈
            
            # 检查止盈
            if self.take_profit is not None:
                should_exit, reason = self.take_profit.should_exit(
                    ts_code=ts_code,
                    current_price=close,
                    cost_price=pos.cost,
                    high=high,
                    low=low,
                    atr=atr,
                )
                if should_exit:
                    quantity = pos.quantity
                    self.sell(ts_code, close, quantity)
                    # logger.info(f"{current_date} 止盈卖出 {ts_code} 数量 {quantity} 价格 {close:.2f} 原因: {reason}")  # 已移除：止盈日志改用进度条

    def run_backtest(self, data_dict: Dict[str, pd.DataFrame], calendar: List[str]) -> pd.DataFrame:
        """主回测入口。

        参数：
        - data_dict: {ts_code: DataFrame}，每个 df 至少含 ['trade_date', 'open', 'high', 'low', 'close']；
        - calendar: 统一交易日列表（升序，字符串 YYYYMMDD）。

        返回：
        - 资金曲线 DataFrame，列为 ['trade_date', 'total_asset', 'cash', 'market_value']。
        """
        from tqdm import tqdm
        
        logger.info(f"回测开始，共 {len(calendar)} 个交易日")

        # 使用进度条显示回测进度
        for current_date in tqdm(calendar, desc="回测进度", unit="日"):
            # 更新持有天数
            self._update_days_held(current_date, data_dict)
            
            # 检查止盈止损（在策略决策前）
            self._check_stop_loss_take_profit(current_date, data_dict)

            # 策略决策
            try:
                self.on_bar(current_date, data_dict)
            except Exception as e:  # noqa: BLE001
                logger.error(f"on_bar 异常: {e}")

            # 计算当日资产
            market_value = self._calc_market_value(current_date, data_dict)
            total_asset = self.cash + market_value

            self.equity_curve.append(
                {
                    "trade_date": current_date,
                    "total_asset": total_asset,
                    "cash": self.cash,
                    "market_value": market_value,
                }
            )

        logger.info("回测结束")
        return pd.DataFrame(self.equity_curve)
