import pandas as pd
import numpy as np

def holder_trade_factors(all_data_df: pd.DataFrame,
                         stk_holdertrade_df: pd.DataFrame) -> pd.DataFrame:
    """
    计算基于股东增减持数据的因子。

    Args:
        all_data_df (pd.DataFrame): 包含每日股票数据的 DataFrame，
                                    必须包含 'ts_code' 和 'trade_date' 列。
        stk_holdertrade_df (pd.DataFrame): 包含股东增减持信息的 DataFrame。
                                           必须包含 'ts_code', 'ann_date',
                                           'in_de' (例如, '增持', '减持'),
                                           和 'change_ratio'。

    Returns:
        pd.DataFrame: 添加了新的股东增减持因子的 all_data_df DataFrame。
    """
    print("开始计算股东增减持因子...")

    # --- 1. 预处理 stk_holdertrade_df ---
    # 创建副本以避免修改原始传入的DataFrame
    stk_trade_processed_df = stk_holdertrade_df.copy()

    # 确保 'ann_date' 是 datetime 类型
    stk_trade_processed_df['ann_date'] = pd.to_datetime(stk_trade_processed_df['ann_date'])

    # 将 'in_de' 映射为数值: 1 代表 '增持', -1 代表 '减持'
    # 请根据你数据中实际的 'in_de' 字符串调整
    in_de_map = {'增持': 1, '减持': -1} # 假设是这两个值
    # 如果你的值是 '1' 和 '2' (1代表增持, 2代表减持)，则映射应相应调整
    # 或者如果 'in_de' 已经是 1 和 -1 (或类似数值)，则可以跳过映射，但要确保类型正确
    stk_trade_processed_df['_direction'] = stk_trade_processed_df['in_de'].map(in_de_map)
    # 如果 _direction 列在映射后可能产生NaN (因为in_de中有未覆盖的值)，需要处理
    if stk_trade_processed_df['_direction'].isnull().any():
        print("警告: 'in_de' 列中存在未映射的值，可能导致 _direction 列出现NaN。")
        # 可以选择填充NaN，例如用0填充，或者移除这些行
        # stk_trade_processed_df['_direction'].fillna(0, inplace=True)

    # 计算有效变动比例 (方向 * 变动比例)
    # 确保 change_ratio 是数值类型。假设 change_ratio 是一个正确的比例值 (例如 0.005 代表 0.5%)。
    # 如果 change_ratio 是百分点 (例如 0.5 代表 0.5%)，你可能需要除以 100。
    stk_trade_processed_df['change_ratio'] = pd.to_numeric(stk_trade_processed_df['change_ratio'], errors='coerce')
    stk_trade_processed_df['_effective_change'] = stk_trade_processed_df['_direction'] * stk_trade_processed_df['change_ratio']

    # 按股票代码和公告日期聚合当日的多次增减持操作
    daily_trade_agg = stk_trade_processed_df.groupby(['ts_code', 'ann_date']).agg(
        net_change_ratio_daily=('_effective_change', 'sum'),
        any_increase_daily=('_direction', lambda x: (x == 1).any().astype(int)),
        any_decrease_daily=('_direction', lambda x: (x == -1).any().astype(int))
    ).reset_index()

    # 将 'ann_date' 重命名为 'trade_date' 以便与 all_data_df 合并
    daily_trade_agg = daily_trade_agg.rename(columns={'ann_date': 'trade_date'})

    # --- 2. 与 all_data_df 合并 ---
    # 创建 all_data_df 的副本进行操作
    df_merged = all_data_df.copy()
    df_merged['trade_date'] = pd.to_datetime(df_merged['trade_date']) # 确保日期类型一致

    # 使用左合并保留 all_data_df 的所有行，并在有增减持公告的日期添加信息
    df_merged = pd.merge(df_merged, daily_trade_agg, on=['ts_code', 'trade_date'], how='left')

    # 对于没有增减持公告的日期，填充NaN值为0（表示当天无变动或无公告）
    df_merged['net_change_ratio_daily'] = df_merged['net_change_ratio_daily'].fillna(0)
    df_merged['any_increase_daily'] = df_merged['any_increase_daily'].fillna(0)
    df_merged['any_decrease_daily'] = df_merged['any_decrease_daily'].fillna(0)

    # --- 3. 计算滚动因子 ---
    # !!! 关键步骤：确保在 groupby().rolling() 之前按分组键和时间键排序 !!!
    # 这一步至关重要，以保证滚动计算后的 Series 在 reset_index 后能正确对齐
    df_merged = df_merged.sort_values(['ts_code', 'trade_date']).reset_index(drop=True)

    grouped = df_merged.groupby('ts_code', group_keys=False) # group_keys=False 避免在结果中保留分组键作为索引层级

    # 因子: 过去N日净变动比例之和
    for N in [10]:
        rolling_series = grouped['net_change_ratio_daily'].rolling(window=N, min_periods=1).sum()
        df_merged[f'holder_net_change_sum_{N}d'] = rolling_series.reset_index(level=0, drop=True)

    # 因子: 过去N日发生增持的天数
    for N in [10]:
        rolling_series = grouped['any_increase_daily'].rolling(window=N, min_periods=1).sum()
        df_merged[f'holder_increase_days_{N}d'] = rolling_series.reset_index(level=0, drop=True)

    # 因子: 过去N日发生减持的天数
    for N in [10]:
        rolling_series = grouped['any_decrease_daily'].rolling(window=N, min_periods=1).sum()
        df_merged[f'holder_decrease_days_{N}d'] = rolling_series.reset_index(level=0, drop=True)

    # 因子: 过去N日是否发生过增持 (布尔标志)
    for N in [10]:
        rolling_series_sum = grouped['any_increase_daily'].rolling(window=N, min_periods=1).sum()
        df_merged[f'holder_any_increase_flag_{N}d'] = (rolling_series_sum > 0).astype(int).reset_index(level=0, drop=True)

    # 因子: 过去N日是否发生过减持 (布尔标志)
    for N in [10]:
        rolling_series_sum = grouped['any_decrease_daily'].rolling(window=N, min_periods=1).sum()
        df_merged[f'holder_any_decrease_flag_{N}d'] = (rolling_series_sum > 0).astype(int).reset_index(level=0, drop=True)

    # 因子: 过去N日净“活动”得分 (增持天数 - 减持天数)
    for N in [10]:
        # 直接使用已经计算好且索引对齐的列
        df_merged[f'holder_direction_score_{N}d'] = df_merged[f'holder_increase_days_{N}d'] - df_merged[f'holder_decrease_days_{N}d']

    # 清理在合并过程中产生的每日临时列（如果不再需要它们）
    df_merged.drop(columns=['net_change_ratio_daily', 'any_increase_daily', 'any_decrease_daily'], inplace=True, errors='ignore')

    print("股东增减持因子计算完成。")
    return df_merged