"""
VectorBT 回测引擎封装。

统一接口：engine.run(strategy, price_df, factor_df) → BacktestReport
"""

import numpy as np
import pandas as pd
import vectorbt as vbt

from backtest.base import BaseStrategy
from backtest.report import BacktestReport


class VectorBTEngine:
    """
    VectorBT 回测引擎。

    只做多，不做空。
    """

    def __init__(
        self,
        initial_capital: float = 100_000,
        commission: float = 0.0003,  # 万三
        freq: str = "D",
    ):
        self.initial_capital = initial_capital
        self.commission = commission
        self.freq = freq

    # ── 单股票回测 ────────────────────────────────────────

    def run(
        self,
        strategy: BaseStrategy,
        price_df: pd.DataFrame,
        factor_df: pd.DataFrame | None = None,
    ) -> BacktestReport:
        """
        单股票回测。

        参数:
            strategy: 策略实例
            price_df: 价格数据，index=trade_date，必须有 'close' 列
            factor_df: 因子数据，index=trade_date。
                       None 时使用 price_df 作为因子数据源。

        返回:
            BacktestReport
        """
        if factor_df is None:
            factor_df = price_df

        # 1. 对齐日期
        common_idx = price_df.index.intersection(factor_df.index)
        if len(common_idx) < 2:
            return BacktestReport()

        price_df = price_df.loc[common_idx].sort_index()
        factor_df = factor_df.loc[common_idx].sort_index()

        # 2. 合并 close 到 factor_df（策略可能需要）
        if "close" not in factor_df.columns:
            factor_df = factor_df.copy()
            factor_df["close"] = price_df["close"]

        # 3. 生成信号
        raw_signals = strategy.generate_signals(factor_df)

        # 4. 信号 → VectorBT entries/exits
        entries, exits = self._signals_to_entries(raw_signals, price_df.index)

        # 5. 运行回测
        close = price_df["close"]
        pf = vbt.Portfolio.from_signals(
            close,
            entries=entries,
            exits=exits,
            init_cash=self.initial_capital,
            fees=self.commission,
            freq=self.freq,
            direction="longonly",
        )

        return BacktestReport.from_vbt_result(pf, close)

    # ── 截面回测（多股票） ──────────────────────────────────

    def run_cross_section(
        self,
        strategy: BaseStrategy,
        price_universe: dict[str, pd.DataFrame],
        factor_universe: dict[str, pd.DataFrame] | None = None,
        rebalance_freq: str = "M",
    ) -> BacktestReport:
        """
        截面策略回测（多股票 + 定期调仓）。

        对每只股票独立回测，合并权益曲线。

        参数:
            strategy: 策略实例
            price_universe: {ts_code: price_df}
            factor_universe: {ts_code: factor_df}
            rebalance_freq: 调仓频率 'D'/'W'/'M'，用于合并时对齐

        返回:
            BacktestReport
        """
        if factor_universe is None:
            factor_universe = price_universe

        stock_equities = {}
        stock_reports = {}

        # 逐股票回测
        for ts_code in price_universe:
            price_df = price_universe[ts_code]
            if "close" not in price_df.columns or price_df.empty:
                continue

            factor_df = factor_universe.get(ts_code, price_df)

            report = self.run(strategy, price_df, factor_df)
            if report is not None and len(report.equity_curve) > 0:
                stock_equities[ts_code] = report.equity_curve
                stock_reports[ts_code] = report

        if not stock_equities:
            return BacktestReport()

        # 合并：等权分配资金到各股票
        return self._merge_equities(stock_equities)

    # ── 信号转换 ──────────────────────────────────────────

    @staticmethod
    def _signals_to_entries(
        raw_signals: pd.Series,
        target_index: pd.Index,
    ) -> tuple[pd.Series, pd.Series]:
        """
        将策略信号转为 VectorBT entries/exits。

        信号格式:
            1 → 买入
            0 → 平仓
           -1 → 继续持有/不操作

        entries: True 时开仓
        exits: True 时平仓
        """
        # 对齐到目标 index
        aligned = pd.Series(-1, index=target_index)
        common = target_index.intersection(raw_signals.index)
        aligned.loc[common] = raw_signals.loc[common].values

        entries = pd.Series(False, index=target_index)
        exits = pd.Series(False, index=target_index)

        in_position = False
        for i in range(len(aligned)):
            sig = aligned.iloc[i]
            if not in_position and sig == 1:
                entries.iloc[i] = True
                in_position = True
            elif in_position and sig == 0:
                exits.iloc[i] = True
                in_position = False

        return entries, exits

    # ── 合并多股票权益 ─────────────────────────────────────

    def _merge_equities(
        self, stock_equities: dict[str, pd.Series]
    ) -> BacktestReport:
        """等权合并多股票权益曲线，构建组合级报告。"""
        equity_df = pd.DataFrame(stock_equities)
        equity_df = equity_df.ffill().fillna(0)
        # 转为 DatetimeIndex
        if not isinstance(equity_df.index, pd.DatetimeIndex):
            equity_df.index = pd.to_datetime(equity_df.index, format="%Y%m%d")

        n_stocks = len(stock_equities)
        weight = 1.0 / n_stocks if n_stocks > 0 else 1.0

        # 加权组合收益
        returns_df = equity_df.pct_change().fillna(0)
        portfolio_ret = returns_df.mean(axis=1)  # 等权 = 逐行平均

        # 组合净值
        portfolio_equity = self.initial_capital * (1 + portfolio_ret).cumprod()

        dd = portfolio_equity / portfolio_equity.cummax() - 1
        years = max(len(portfolio_ret) / 252, 0.02)

        total_return = (portfolio_equity.iloc[-1] / portfolio_equity.iloc[0] - 1) * 100
        cagr = ((total_return / 100 + 1) ** (1 / years) - 1) * 100
        mdd = dd.min() * 100
        mean_ret = portfolio_ret.mean() * 252
        std_ret = portfolio_ret.std() * np.sqrt(252)
        sharpe = mean_ret / std_ret if std_ret > 0 else 0
        calmar = cagr / abs(mdd) if abs(mdd) > 0 else 0

        try:
            monthly = portfolio_equity.resample("ME").last().pct_change()
        except Exception:
            monthly = pd.Series(dtype=float)

        return BacktestReport(
            total_return=round(total_return, 2),
            cagr=round(cagr, 2),
            max_drawdown=round(mdd, 2),
            sharpe_ratio=round(sharpe, 2),
            calmar_ratio=round(calmar, 2),
            annual_volatility=round(std_ret * 100 if std_ret != 0 else 0, 2),
            total_trades=0,
            equity_curve=portfolio_equity,
            drawdown_curve=dd,
            monthly_returns=monthly,
        )
