"""
ML 模型回测集成。

MLStrategy: 将 ML 预测值作为交易信号接入回测引擎。
MLBenchmark: 多模型基准对比。
"""

import numpy as np
import pandas as pd

from backtest.base import BaseStrategy
from backtest.vectorbt.engine import VectorBTEngine
from models.base import BaseModel
from models.features import FeatureEngine


class MLStrategy(BaseStrategy):
    """
    ML 预测 → 交易信号。

    用模型预测未来 N 日收益，按预测值分位数生成信号：
    - 预测值 > buy_quantile → 买入
    - 预测值 < sell_quantile → 平仓

    参数:
        model: 已训练的 BaseModel
        feature_engine: 已 fit 的 FeatureEngine
        buy_quantile: 买入分位阈值（0.7 = 预测值最高的30%买入）
        sell_quantile: 卖出分位阈值（0.3 = 预测值最低的30%平仓）
        rebalance_freq: 调仓间隔（交易日）
    """

    category = "ml"

    def __init__(
        self,
        model: BaseModel,
        feature_engine: FeatureEngine,
        buy_quantile: float = 0.7,
        sell_quantile: float = 0.3,
        rebalance_freq: int = 5,
    ):
        self.model = model
        self.feature_engine = feature_engine
        self.buy_quantile = buy_quantile
        self.sell_quantile = sell_quantile
        self.rebalance_freq = rebalance_freq
        self.name = f"ml_{model.name}"

    def generate_signals(self, factor_df: pd.DataFrame) -> pd.Series:
        X, _ = self.feature_engine.build(factor_df, factor_df, fit=False)
        if X.empty:
            return pd.Series(-1, index=factor_df.index)

        preds = self.model.predict(X)
        # 用预测值本身的分布作为阈值（相对排序，避免模型偏差影响）
        buy_threshold = preds.quantile(self.buy_quantile)
        sell_threshold = preds.quantile(self.sell_quantile)

        signals = pd.Series(-1, index=factor_df.index)
        common = signals.index.intersection(preds.index)
        buy_mask = preds.loc[common] > buy_threshold
        sell_mask = preds.loc[common] < sell_threshold
        signals.loc[buy_mask[buy_mask].index] = 1
        signals.loc[sell_mask[sell_mask].index] = 0

        signals = self._filter_rebalance(signals)
        return signals

    def _filter_rebalance(self, signals: pd.Series) -> pd.Series:
        """每隔 rebalance_freq 个交易日保留第一个非持有信号。"""
        result = signals.copy()
        last_active = -self.rebalance_freq - 1
        for i in range(len(result)):
            sig = result.iloc[i]
            if sig in (0, 1):
                if i - last_active >= self.rebalance_freq:
                    last_active = i
                else:
                    result.iloc[i] = -1
        return result


class MLBenchmark:
    """ML 模型基准对比测试。"""

    def __init__(
        self,
        models: list[BaseModel],
        feature_engine: FeatureEngine,
        price_df: pd.DataFrame,
        factor_df: pd.DataFrame,
        bt_engine: VectorBTEngine | None = None,
    ):
        self.models = models
        self.feature_engine = feature_engine
        self.price_df = price_df
        self.factor_df = factor_df
        self.bt_engine = bt_engine or VectorBTEngine()

    def run(self) -> pd.DataFrame:
        """对比各模型的预测质量和回测表现。"""
        rows = []
        for model in self.models:
            strategy = MLStrategy(model, self.feature_engine)
            report = self.bt_engine.run(strategy, self.price_df, self.factor_df)

            # OOS 预测 vs 真实值
            X, y_true = self.feature_engine.build(self.factor_df, self.price_df, fit=True)
            y_pred = model.predict(X)

            ic = y_pred.corr(y_true) if len(y_pred) > 0 else 0

            rows.append({
                "model": model.name,
                "ic": round(ic, 4),
                "total_return": report.total_return,
                "cagr": report.cagr,
                "max_dd": report.max_drawdown,
                "sharpe": report.sharpe_ratio,
                "win_rate": report.win_rate,
                "trades": report.total_trades,
            })
        return pd.DataFrame(rows).set_index("model")