"""
ROE 因子。
"""

import pandas as pd

from factors.base import BaseFactor
from factors.fundamental._mapping import map_fundamental_to_daily, effective_available_dates


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:
        """将财务数据（季度）映射到日线索引，按披露可用日消除前视。"""
        return map_fundamental_to_daily(daily_df, fina_df, column)


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)

        if "roe" not in self._financial_df.columns or "end_date" not in self._financial_df.columns:
            return pd.Series(float("nan"), index=df.index)

        # 基于财务帧计算披露可用日（含 ann_date 优先/法定滞后兜底）
        fina = effective_available_dates(self._financial_df)
        fina = fina[["_avail", "end_date", "roe"]].dropna(subset=["_avail", "roe"]).copy()
        if fina.empty:
            return pd.Series(float("nan"), index=df.index)

        # 报告期月份（A 股季末：03/06/09/12）
        end_dt = pd.to_datetime(fina["end_date"].astype(str), format="%Y%m%d", errors="coerce")
        fina["report_month"] = end_dt.dt.month
        fina["report_year"] = end_dt.dt.year
        # 本年取数映射：{ (year, month): roe }
        cur_map = dict(zip(zip(fina["report_year"], fina["report_month"]), fina["roe"]))
        # 按披露可用日排序，逐期覆盖区间
        fina = fina.sort_values("_avail")

        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():
            avail_dt = row["_avail"]
            if pd.isna(avail_dt):
                continue
            prev_roe = cur_map.get((row["report_year"] - 1, row["report_month"]))
            if prev_roe is None:
                continue
            delta = row["roe"] - prev_roe
            mask = dates >= avail_dt
            roe_delta[mask] = delta

        return roe_delta.astype("float64")