"""
因子排序轮动策略。

定期按因子值排序，买入排名最高的股票（截面策略）。
"""

import pandas as pd

from backtest.base import BaseStrategy


class FactorRotationStrategy(BaseStrategy):
    """
    因子排序选股策略。

    适用于多股票截面场景：对每只股票计算因子值，
    选排名最高的 top_n 只做多。
    """

    category = "rotation"

    def __init__(self, factor_name: str, top_n: int = 5, bottom_n: int = 0):
        self.factor_name = factor_name
        self.top_n = top_n
        self.bottom_n = bottom_n
        self.name = f"rotation_{factor_name}_top{top_n}"

    def generate_signals(self, factor_df: pd.DataFrame) -> pd.Series:
        """
        单股票/截面模式：factor_df 支持两种输入方式。
        - 单股票: 对每只股票单次调用
        - 截面: 通过 run_cross_section 逐股票调用
        """
        if self.factor_name not in factor_df.columns:
            raise ValueError(f"factor_df 缺少 '{self.factor_name}' 列")

        factor = factor_df[self.factor_name]
        valid = factor.dropna()
        if len(valid) < self.top_n * 2:
            return pd.Series(-1, index=factor_df.index)

        threshold = valid.quantile(1 - self.top_n / max(len(valid), self.top_n))
        signals = pd.Series(-1, index=factor_df.index)
        signals[factor > threshold] = 1

        return signals

    def rank_stocks(
        self, factor_values: dict[str, float]
    ) -> list[str]:
        """
        对股票按因子值排序，返回 top N 的 ts_code 列表。

        参数:
            factor_values: {ts_code: factor_value}
        """
        sorted_stocks = sorted(factor_values, key=factor_values.get, reverse=True)
        return sorted_stocks[: self.top_n]