import pandas as pd
import numpy as np
import statsmodels.api as sm # 用于中性化回归
from tqdm import tqdm

epsilon = 1e-10 # 防止除零

def zscore_standardize(train_df: pd.DataFrame, test_df: pd.DataFrame, features: list, epsilon: float = 1e-10):
    """
    对指定特征列进行截面 Z-Score 标准化 (原地修改)。
    方法: Z = (value - cross_sectional_mean) / (cross_sectional_std + epsilon)

    Args:
        df (pd.DataFrame): 输入 DataFrame，需包含 'trade_date' 和 features 列。
        features (list): 需要处理的特征列名列表。
        epsilon (float): 防止除以零的小常数。

    WARNING: 此函数会原地修改输入的 DataFrame 'df'。
    """
    print("开始截面 Z-Score 标准化...")
    if not all(col in train_df.columns for col in features):
        missing = [col for col in features if col not in train_df.columns]
        print(f"错误: DataFrame 中缺少以下特征列: {missing}。跳过标准化处理。")
        return


    for col in tqdm(features, desc="Standardizing"):
        try:
            # 使用 transform 计算截面均值和标准差
            mean = train_df[col].transform('mean')
            std = train_df[col].transform('std')

            # 计算 Z-Score 并原地赋值
            train_df[col] = (train_df[col] - mean) / (std + epsilon)
            test_df[col] = (test_df[col] - mean) / (std + epsilon)

        except KeyError:
            print(f"警告: 列 '{col}' 可能不存在或在分组中出错，跳过此列的标准化处理。")
        except Exception as e:
            print(f"警告: 处理列 '{col}' 时发生错误: {e}，跳过此列的标准化处理。")

    print("截面 Z-Score 标准化完成。")


# --- 1. 中位数去极值 (MAD) ---

def cs_mad_filter(df: pd.DataFrame,
                  features: list,
                  k: float = 3.0,
                  scale_factor: float = 1.4826):
    """
    对指定特征列进行截面 MAD 去极值处理 (原地修改)。

    方法: 对每日截面数据，计算 median 和 MAD，
          将超出 [median - k * scale * MAD, median + k * scale * MAD] 范围的值
          替换为边界值 (Winsorization)。
          scale_factor=1.4826 使得 MAD 约等于正态分布的标准差。

    Args:
        df (pd.DataFrame): 输入 DataFrame，需包含 'trade_date' 和 features 列。
        features (list): 需要处理的特征列名列表。
        k (float): MAD 的倍数，用于确定边界。默认为 3.0。
        scale_factor (float): MAD 的缩放因子。默认为 1.4826。

    WARNING: 此函数会原地修改输入的 DataFrame 'df'。
    """
    print(f"开始截面 MAD 去极值处理 (k={k})...")
    if not all(col in df.columns for col in features):
        missing = [col for col in features if col not in df.columns]
        print(f"错误: DataFrame 中缺少以下特征列: {missing}。跳过去极值处理。")
        return

    grouped = df.groupby('trade_date')

    for col in tqdm(features, desc="MAD Filtering"):
        try:
            # 计算截面中位数
            median = grouped[col].transform('median')
            # 计算截面 MAD (Median Absolute Deviation from Median)
            mad = (df[col] - median).abs().groupby(df['trade_date']).transform('median')

            # 计算上下边界
            lower_bound = median - k * scale_factor * mad
            upper_bound = median + k * scale_factor * mad

            # 原地应用 clip
            df[col] = np.clip(df[col], lower_bound, upper_bound)

        except KeyError:
            print(f"警告: 列 '{col}' 可能不存在或在分组中出错，跳过此列的 MAD 处理。")
        except Exception as e:
            print(f"警告: 处理列 '{col}' 时发生错误: {e}，跳过此列的 MAD 处理。")

    print("截面 MAD 去极值处理完成。")


# --- 2. 行业市值中性化 ---

def cs_neutralize_industry_cap(df: pd.DataFrame,
                               features: list,
                               industry_col: str = 'cat_l2_code',
                               market_cap_col: str = 'circ_mv'):
    """
    对指定特征列进行截面行业和对数市值中性化 (原地修改)。
    使用 OLS 回归: feature ~ 1 + log(market_cap) + C(industry)
    将回归残差写回原特征列。

    Args:
        df (pd.DataFrame): 输入 DataFrame，需包含 'trade_date', features 列,
                           industry_col, market_cap_col。
        features (list): 需要处理的特征列名列表。
        industry_col (str): 行业分类列名。
        market_cap_col (str): 流通市值列名。

    WARNING: 此函数会原地修改输入的 DataFrame 'df' 的 features 列。
             计算量较大，可能耗时较长。
             需要安装 statsmodels 库 (pip install statsmodels)。
    """
    print("开始截面行业市值中性化...")
    required_cols = features + ['trade_date', industry_col, market_cap_col]
    if not all(col in df.columns for col in required_cols):
        missing = [col for col in required_cols if col not in df.columns]
        print(f"错误: DataFrame 中缺少必需列: {missing}。无法进行中性化。")
        return

    # 预处理：计算 log 市值，处理 industry code 可能的 NaN
    log_cap_col = '_log_market_cap'
    df[log_cap_col] = np.log1p(df[market_cap_col]) # log1p 处理 0 值
    # df[industry_col] = df[industry_col].cat.add_categories('UnknownIndustry')
    # df[industry_col] = df[industry_col].fillna('UnknownIndustry') # 填充行业 NaN
    # df[industry_col] = df[industry_col].astype('category') # 转为类别，ols 会自动处理

    dates = df['trade_date'].unique()
    all_residuals = [] # 用于收集所有日期的残差

    for date in tqdm(dates, desc="Neutralizing"):
        daily_data = df.loc[df['trade_date'] == date, features + [log_cap_col, industry_col]].copy() # 使用 .loc 获取副本

        # 准备自变量 X (常数项 + log市值 + 行业哑变量)
        X = daily_data[[log_cap_col]]
        X = sm.add_constant(X, prepend=True) # 添加常数项
        # 创建行业哑变量 (drop_first=True 避免共线性)
        industry_dummies = pd.get_dummies(daily_data[industry_col], prefix=industry_col, drop_first=True)
        industry_dummies = industry_dummies.astype(int)
        X = pd.concat([X, industry_dummies], axis=1)

        daily_residuals = daily_data[[col for col in features]].copy() # 创建用于存储残差的df

        for col in features:
            Y = daily_data[col]

            # 处理 NaN 值，确保 X 和 Y 在相同位置有有效值
            valid_mask = Y.notna() & X.notna().all(axis=1)
            if valid_mask.sum() < (X.shape[1] + 1): # 数据点不足以估计模型
                print(f"警告: 日期 {date}, 特征 {col} 有效数据不足 ({valid_mask.sum()}个)，无法中性化，填充 NaN。")
                daily_residuals[col] = np.nan
                continue

            Y_valid = Y[valid_mask]
            X_valid = X[valid_mask]

            # 执行 OLS 回归
            try:
                model = sm.OLS(Y_valid.to_numpy(), X_valid.to_numpy())
                results = model.fit()
                # 将残差填回对应位置
                daily_residuals.loc[valid_mask, col] = results.resid
                daily_residuals.loc[~valid_mask, col] = np.nan # 原本无效的位置填充 NaN
            except Exception as e:
                print(f"警告: 日期 {date}, 特征 {col} 回归失败: {e}，填充 NaN。")
                daily_residuals[col] = np.nan
                break

        all_residuals.append(daily_residuals)

    # 合并所有日期的残差结果
    if all_residuals:
        residuals_df = pd.concat(all_residuals)
        # 将残差结果更新回原始 df (原地修改)
        # 使用 update 比 merge 更适合基于索引的原地更新
        # 确保 residuals_df 的索引与 df 中对应部分一致
        df.update(residuals_df)
    else:
        print("没有有效的残差结果可以合并。")


    # 清理临时列
    df.drop(columns=[log_cap_col], inplace=True)
    print("截面行业市值中性化完成。")


# --- 3. Z-Score 标准化 ---

def cs_zscore_standardize(df: pd.DataFrame, features: list, epsilon: float = 1e-10):
    """
    对指定特征列进行截面 Z-Score 标准化 (原地修改)。
    方法: Z = (value - cross_sectional_mean) / (cross_sectional_std + epsilon)

    Args:
        df (pd.DataFrame): 输入 DataFrame，需包含 'trade_date' 和 features 列。
        features (list): 需要处理的特征列名列表。
        epsilon (float): 防止除以零的小常数。

    WARNING: 此函数会原地修改输入的 DataFrame 'df'。
    """
    print("开始截面 Z-Score 标准化...")
    if not all(col in df.columns for col in features):
        missing = [col for col in features if col not in df.columns]
        print(f"错误: DataFrame 中缺少以下特征列: {missing}。跳过标准化处理。")
        return

    grouped = df.groupby('trade_date')

    for col in tqdm(features, desc="Standardizing"):
        try:
            # 使用 transform 计算截面均值和标准差
            mean = grouped[col].transform('mean')
            std = grouped[col].transform('std')

            # 计算 Z-Score 并原地赋值
            df[col] = (df[col] - mean) / (std + epsilon)

        except KeyError:
            print(f"警告: 列 '{col}' 可能不存在或在分组中出错，跳过此列的标准化处理。")
        except Exception as e:
            print(f"警告: 处理列 '{col}' 时发生错误: {e}，跳过此列的标准化处理。")

    print("截面 Z-Score 标准化完成。")

def fill_nan_with_daily_median(df: pd.DataFrame, feature_columns: list[str]) -> pd.DataFrame:
    """
    对指定特征列进行每日截面中位数填充缺失值 (NaN)。

    参数:
    df (pd.DataFrame): 包含多日数据的DataFrame，需要包含 'trade_date' 和 feature_columns 中的列。
    feature_columns (list[str]): 需要进行缺失值填充的特征列名称列表。

    返回:
    pd.DataFrame: 包含缺失值填充后特征列的DataFrame。在输入DataFrame的副本上操作。
    """
    processed_df = df.copy() # 在副本上操作，保留原始数据

    # 确保 trade_date 是 datetime 类型以便正确分组
    processed_df['trade_date'] = pd.to_datetime(processed_df['trade_date'])

    def _fill_daily_nan(group):
        # group 是某一个交易日的 DataFrame

        # 遍历指定的特征列
        for feature_col in feature_columns:
            # 检查列是否存在于当前分组中
            if feature_col in group.columns:
                # 计算当日该特征的中位数
                median_val = group[feature_col].median()

                # 使用当日中位数填充该特征列的 NaN 值
                # inplace=True 会直接修改 group DataFrame
                group[feature_col].fillna(median_val, inplace=True)
            # else:
            #     print(f"Warning: Feature column '{feature_col}' not found in daily group for {group['trade_date'].iloc[0]}. Skipping.")

        return group

    # 按交易日期分组，并应用每日填充函数
    # group_keys=False 避免将分组键添加到结果索引中
    filled_df = processed_df.groupby('trade_date', group_keys=False).apply(_fill_daily_nan)

    return filled_df