"""
策略模块 — 信号生成、权重分配、组合构建
"""
from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from typing import Dict, List, Optional, Set, Tuple

import numpy as np
import pandas as pd


# ---------------------------------------------------------------------------
# 信号生成器
# ---------------------------------------------------------------------------

class SignalGenerator(ABC):
    """信号生成器基类 — 将因子值转换为交易信号"""

    @abstractmethod
    def generate(self, factor_df: pd.DataFrame, **kwargs) -> pd.DataFrame:
        """
        factor_df: DataFrame, index=date, columns=stock_codes, values=factor值
        返回信号 DataFrame，同结构，1=做多, 0=平仓, -1=做空
        """
        ...


class QuantileSignal(SignalGenerator):
    """基于因子分位数的信号生成器"""

    def __init__(
        self,
        n_quantiles: int = 5,
        long_quantile: int = 5,
        short_quantile: int = 1,
        min_stocks: int = 10,
    ):
        self.n_quantiles = n_quantiles
        self.long_quantile = long_quantile
        self.short_quantile = short_quantile
        self.min_stocks = min_stocks

    def generate(self, factor_df: pd.DataFrame, **kwargs) -> pd.DataFrame:
        """
        逐截面按因子值分组：
        long_quantile 档 → +1
        short_quantile 档 → -1
        其余 → 0
        """
        signal = pd.DataFrame(0, index=factor_df.index, columns=factor_df.columns)

        for date in factor_df.index:
            row = factor_df.loc[date].dropna()
            if len(row) < self.min_stocks:
                continue
            try:
                bins = pd.qcut(
                    row, self.n_quantiles, labels=False, duplicates="drop"
                )
            except ValueError:
                continue
            signal.loc[date, bins == self.long_quantile - 1] = 1
            signal.loc[date, bins == self.short_quantile - 1] = -1

        return signal


class ZScoreSignal(SignalGenerator):
    """基于 Z-Score 阈值的信号生成器"""

    def __init__(self, entry_z: float = 1.0, exit_z: float = 0.5):
        self.entry_z = entry_z
        self.exit_z = exit_z

    def generate(self, factor_df: pd.DataFrame, **kwargs) -> pd.DataFrame:
        """
        因子截面标准化（Z-Score）:
        > entry_z  → +1
        < -entry_z → -1
        其余       → 0
        """
        signal = pd.DataFrame(0, index=factor_df.index, columns=factor_df.columns)
        for date in factor_df.index:
            row = factor_df.loc[date]
            valid = row.notna()
            if valid.sum() < 3:
                continue
            z = (row - row[valid].mean()) / (row[valid].std() + 1e-12)
            signal.loc[date, z > self.entry_z] = 1
            signal.loc[date, z < -self.entry_z] = -1
        return signal


# ---------------------------------------------------------------------------
# 权重分配器
# ---------------------------------------------------------------------------

class WeightAllocator(ABC):
    """权重分配器基类"""

    @abstractmethod
    def allocate(
        self,
        signals: pd.DataFrame,
        prices: pd.DataFrame,
        cash: float,
        positions: Dict[str, int],
        **kwargs,
    ) -> Dict[str, float]:
        """
        返回目标权重字典 {stock_code: weight}
        weight 为目标持仓市值占比 (0~1)
        """
        ...


class EqualWeightAllocator(WeightAllocator):
    """等权分配器"""

    def __init__(self, max_positions: int = 30):
        self.max_positions = max_positions

    def allocate(
        self,
        signals: pd.DataFrame,
        prices: pd.DataFrame,
        cash: float,
        positions: Dict[str, int],
        **kwargs,
    ) -> Dict[str, float]:
        latest_signal = signals.iloc[-1] if len(signals) > 0 else pd.Series(dtype=float)
        long_stocks = latest_signal[latest_signal > 0].index.tolist()

        if not long_stocks:
            return {}

        # 限制持仓数
        long_stocks = long_stocks[: self.max_positions]
        weight = 1.0 / len(long_stocks)
        return {s: weight for s in long_stocks}


class FactorWeightAllocator(WeightAllocator):
    """因子值加权分配器 — 因子越大权重越大"""

    def __init__(self, max_positions: int = 30, min_weight: float = 0.005):
        self.max_positions = max_positions
        self.min_weight = min_weight

    def allocate(
        self,
        signals: pd.DataFrame,
        prices: pd.DataFrame,
        cash: float,
        positions: Dict[str, int],
        factor_df: Optional[pd.DataFrame] = None,
        **kwargs,
    ) -> Dict[str, float]:
        if factor_df is None:
            return EqualWeightAllocator(self.max_positions).allocate(
                signals, prices, cash, positions
            )

        latest_signal = signals.iloc[-1] if len(signals) > 0 else pd.Series(dtype=float)
        latest_factor = factor_df.iloc[-1] if len(factor_df) > 0 else pd.Series(dtype=float)

        long_stocks = latest_signal[latest_signal > 0].index
        # 按因子值排序
        valid = latest_factor[long_stocks].dropna().sort_values(ascending=False)
        selected = valid.head(self.max_positions)
        if selected.empty:
            return {}

        raw_w = selected.values / (selected.values.sum() + 1e-12)
        raw_w = np.clip(raw_w, self.min_weight, 1.0)
        raw_w = raw_w / raw_w.sum()
        return dict(zip(selected.index.tolist(), raw_w.tolist()))


# ---------------------------------------------------------------------------
# 策略上下文
# ---------------------------------------------------------------------------

@dataclass
class Strategy:
    """
    量化策略 — 组合因子、信号生成器、权重分配器，定义一个完整策略。
    """

    name: str
    factors: List  # 因子对象列表
    signal_generator: SignalGenerator
    weight_allocator: WeightAllocator
    factor_weights: Optional[Dict[str, float]] = None  # 因子合成权重
    filter_universe: Optional[Set[str]] = None          # 可选股票池
    description: str = ""

    def compute_composite_factor(self, factor_data: Dict[str, pd.DataFrame]) -> pd.DataFrame:
        """
        合成多因子值。
        factor_data: {factor_name: DataFrame(index=date, columns=stocks)}
        """
        if not factor_data:
            raise ValueError("factor_data 为空")

        # 统一列（股票池）合并
        all_stocks = sorted(
            set().union(*[set(df.columns) for df in factor_data.values()])
        )
        all_dates = sorted(
            set().union(*[set(df.index) for df in factor_data.values()])
        )

        composite = pd.DataFrame(0.0, index=all_dates, columns=all_stocks)
        n_factors = len(factor_data)

        for name, df in factor_data.items():
            w = (
                self.factor_weights.get(name, 1.0 / n_factors)
                if self.factor_weights
                else 1.0 / n_factors
            )
            aligned = df.reindex(index=all_dates, columns=all_stocks)
            # 截面标准化
            z = aligned.sub(aligned.mean(axis=1), axis=0).div(
                aligned.std(axis=1) + 1e-12, axis=0
            )
            composite += w * z

        return composite

    def run_step(
        self,
        date: pd.Timestamp,
        factor_data: Dict[str, pd.DataFrame],
        prices: pd.DataFrame,
        cash: float,
        positions: Dict[str, int],
    ) -> Dict[str, float]:
        """
        单步执行：生成信号 → 分配权重 → 返回目标持仓权重。
        """
        composite = self.compute_composite_factor(factor_data)
        # 仅取当前日期截面
        if date in composite.index:
            current_slice = composite.loc[[date]]
        else:
            current_slice = composite.iloc[-1:]

        signals = self.signal_generator.generate(current_slice)
        weights = self.weight_allocator.allocate(
            signals, prices, cash, positions, factor_df=composite
        )
        return weights