"""
ROE 因子。
"""

import pandas as pd

from factors.base import BaseFactor


class ROEFactor(BaseFactor):
    """
    ROE 因子。

    从财务数据提取 ROE 并映射到日线。

    需要 df 中包含 'roe' 列（由 FactorEngine 合并财务数据后传入），
    或将 financial_df 直接传入构造函数。
    """

    category = "fundamental"

    def __init__(self, financial_df: pd.DataFrame | None = None):
        """
        参数:
            financial_df: 财务数据 DataFrame，columns 含 'end_date', 'roe'。
                          None 时需在 df 参数中直接提供 roe 列。
        """
        self._financial_df = financial_df
        self.name = "roe"

    def calculate(self, df: pd.DataFrame) -> pd.Series:
        if "roe" in df.columns:
            return df["roe"].copy()

        if self._financial_df is None or self._financial_df.empty:
            return pd.Series(float("nan"), index=df.index)

        return self._map_financial_to_daily(
            df, self._financial_df, "roe"
        )

    @staticmethod
    def _map_financial_to_daily(
        daily_df: pd.DataFrame,
        fina_df: pd.DataFrame,
        column: str,
    ) -> pd.Series:
        """将财务数据（季度）映射到日线索引。"""
        fina = fina_df[["end_date", column]].dropna().copy()
        fina["end_date"] = fina["end_date"].astype(str)
        fina = fina.sort_values("end_date")

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

        if fina.empty:
            return result

        dates = pd.to_datetime(daily_df.index, format="%Y%m%d", errors="coerce")
        fina_dates = pd.to_datetime(fina["end_date"], format="%Y%m%d", errors="coerce")

        for i, fina_date in enumerate(fina_dates):
            mask = dates >= fina_date
            if i + 1 < len(fina_dates):
                mask &= dates < fina_dates.iloc[i + 1]
            else:
                pass  # 最新一期覆盖所有后续日期
            result[mask] = fina[column].iloc[i]

        return result.astype("float64")


class ROETTMDeltaFactor(BaseFactor):
    """ROE 同比变化（当前 ROE - 去年同期 ROE）。"""

    category = "fundamental"

    def __init__(self, financial_df: pd.DataFrame | None = None):
        self._financial_df = financial_df
        self.name = "roe_delta"

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

        fina = self._financial_df[["end_date", "roe"]].dropna().copy()
        fina["end_date"] = fina["end_date"].astype(str)
        fina["year"] = fina["end_date"].str[:4].astype(int)
        fina = fina.sort_values("end_date")

        # 按年分组计算 YoY 差值
        roe_delta = pd.Series(float("nan"), index=df.index)
        dates = pd.to_datetime(df.index, format="%Y%m%d", errors="coerce")

        for _, row in fina.iterrows():
            this_year = row["year"]
            prev_row = fina[fina["year"] == this_year - 1]
            if prev_row.empty:
                continue
            delta = row["roe"] - prev_row["roe"].iloc[-1]
            f_date = pd.to_datetime(row["end_date"], format="%Y%m%d")
            mask = dates >= f_date
            roe_delta[mask] = delta

        return roe_delta.astype("float64")