"""
情绪因子。

将 Qwen 输出的情绪分数转换为量化因子值。
"""

import numpy as np
import pandas as pd

from factors.base import BaseFactor


class NewsSentimentFactor(BaseFactor):
    """
    新闻情绪因子。

    将多条新闻的情绪分数按时间加权聚合到每个交易日。

    参数:
        window: 滚动窗口（交易日）
        decay: 指数衰减系数（越大衰减越快），0 表示等权
        sentiment_df: 情绪分析结果 DataFrame
                      (date, sentiment_score, confidence, title, content)
    """

    category = "sentiment"

    def __init__(self, window: int = 5, decay: float = 0.3, sentiment_df: pd.DataFrame | None = None):
        self.window = window
        self.decay = decay
        self.sentiment_df = sentiment_df
        self.name = f"news_sent_{window}"

    def calculate(self, df: pd.DataFrame) -> pd.Series:
        if self.sentiment_df is None or self.sentiment_df.empty:
            return pd.Series(float("nan"), index=df.index, name=self.name)

        return _aggregate_sentiment(
            df, self.sentiment_df, self.window, self.decay
        )

    def get_required_columns(self) -> list[str]:
        return []


class SentimentMomentumFactor(BaseFactor):
    """
    情绪动量因子。

    当前情绪 - N 日前情绪，衡量情绪变化方向。
    """

    category = "sentiment"

    def __init__(self, period: int = 5, sentiment_df: pd.DataFrame | None = None):
        self.period = period
        self.sentiment_df = sentiment_df
        self.name = f"sent_delta_{period}"

    def calculate(self, df: pd.DataFrame) -> pd.Series:
        base = NewsSentimentFactor(window=1, sentiment_df=self.sentiment_df).calculate(df)
        return base.diff(self.period)

    def get_required_columns(self) -> list[str]:
        return []


class SentimentConfidenceFactor(BaseFactor):
    """
    情绪置信度因子。

    新闻情绪分析的置信度越高，因子绝对值越大（方向同 sentiment_score）。
    sentiment_score × confidence → 高置信利好=正值大，高置信利空=负值大。
    """

    category = "sentiment"

    def __init__(self, window: int = 5, sentiment_df: pd.DataFrame | None = None):
        self.window = window
        self.sentiment_df = sentiment_df
        self.name = f"news_conf_{window}"

    def calculate(self, df: pd.DataFrame) -> pd.Series:
        if self.sentiment_df is None or self.sentiment_df.empty:
            return pd.Series(float("nan"), index=df.index, name=self.name)

        sdf = self.sentiment_df.copy()
        # 设置加权分数
        if "confidence" in sdf.columns and "sentiment_score" in sdf.columns:
            sdf["weighted_score"] = sdf["sentiment_score"] * sdf["confidence"]
        else:
            return pd.Series(float("nan"), index=df.index, name=self.name)

        return _aggregate_sentiment(df, sdf, self.window, decay=0.3, score_col="weighted_score")

    def get_required_columns(self) -> list[str]:
        return []


# ── 情绪聚合工具函数 ──────────────────────────────────────

def _aggregate_sentiment(
    daily_df: pd.DataFrame,
    sentiment_df: pd.DataFrame,
    window: int,
    decay: float,
    score_col: str = "sentiment_score",
) -> pd.Series:
    """
    将情绪分数按时间加权聚合到交易日。

    逻辑：
    1. 对每个交易日 t，找到 [t - window + 1, t] 范围内的所有新闻
    2. 按 time_decay = exp(-decay * days_from_t) 加权
    3. 按 confidence（如有）加权
    4. 返回加权平均情绪分数

    参数:
        daily_df: 日线 DataFrame（提供 index 和日期对齐）
        sentiment_df: 情绪 DataFrame（date 列 + score_col）
        window: 窗口大小
        decay: 衰减系数
        score_col: 情绪分数列名
    """
    if sentiment_df.empty:
        return pd.Series(float("nan"), index=daily_df.index, name=score_col)

    # 统一日期格式
    sdf = sentiment_df.copy()
    sdf["date"] = pd.to_datetime(sdf["date"], format="%Y%m%d", errors="coerce")
    sdf = sdf.dropna(subset=["date"])
    sdf = sdf.sort_values("date")

    daily_idx = pd.to_datetime(daily_df.index, format="%Y%m%d", errors="coerce")
    if daily_idx.isna().all():
        daily_idx = pd.to_datetime(daily_df.index)

    result = pd.Series(float("nan"), index=daily_df.index)

    # 对 news 日期建立搜索索引
    news_dates = sdf["date"].values

    for i, dt in enumerate(daily_idx):
        if pd.isna(dt):
            continue
        # 窗口起始
        window_start = dt - pd.Timedelta(days=window * 2)  # 宽窗覆盖非交易日

        mask = (news_dates >= window_start) & (news_dates <= dt)
        candidates = sdf[mask]

        if candidates.empty:
            continue

        # 时间衰减权重
        days_diff = (dt - candidates["date"]).dt.days
        time_weights = np.exp(-decay * days_diff)

        # 置信度权重（如有）
        conf_weights = candidates.get("confidence", pd.Series(1.0, index=candidates.index)).fillna(0.5)
        scores = candidates[score_col].fillna(0.0)

        total_weight = (time_weights * conf_weights).sum()
        if total_weight > 0:
            result.iloc[i] = (scores * time_weights * conf_weights).sum() / total_weight

    result.name = score_col
    return result.astype("float64")