import pandas as pd
import numpy as np

def calculate_sharpe_sortino_labels_efficient(
    df: pd.DataFrame,
    N_days: int = 5,  # 未来考察的天数，用于计算收益率序列
    risk_free_rate_annual: float = 0.02, # 年化无风险利率 (例如 0.02 表示 2%)
    annualization_factor: float = np.sqrt(252), # 年化因子，日数据通常用sqrt(252)
    min_periods_for_std: int = 2 # 计算标准差所需的最小有效数据点数
) -> pd.DataFrame:
    """
    高效计算每只股票在每个交易日未来 N 日的年化夏普比率和索提诺比率，作为训练的Label。
    函数会获取未来 N 天的每日收益率序列，并基于此序列计算夏普/索提诺。

    Args:
        df (pd.DataFrame): 输入 DataFrame，需包含 'ts_code', 'trade_date', 'close' 列。
                           'close' 列必须是股票收盘价。
        N_days (int): 未来考察的天数。例如，N_days=5 表示考察未来5个交易日。
        risk_free_rate_annual (float): 年化无风险利率。
        annualization_factor (float): 用于将日收益率转化为年化收益率的因子 (例如 np.sqrt(252))。
        min_periods_for_std (int): 计算标准差所需的最小有效数据点数。
                                   如果未来N天有效数据少于此值，则 Label 为 NaN。

    Returns:
        pd.DataFrame: 原始DataFrame，新增 'sharpe_ratio_label' 和 'sortino_ratio_label' 列。
                      注意：数据末尾 N_days 行的 Label 将为 NaN。
    """
    df_copy = df
    
    # 确保数据已按股票和日期排序，这对于高效计算至关重要
    df_copy = df_copy.sort_values(by=['ts_code', 'trade_date'])

    # # 计算每日收益率
    # df_copy['pct_chg'] = df_copy.groupby('ts_code')['close'].pct_change()

    # 将年化无风险利率转换为每日无风险利率
    daily_risk_free_rate = risk_free_rate_annual / 252 # 假设每年252个交易日

    def _calculate_metrics_for_window(returns_series):
        """
        辅助函数：计算单个滚动窗口内的夏普和索提诺比率。
        这个函数将被 apply 调用于每个 (N_days) 收益率序列。
        """
        if len(returns_series.dropna()) < min_periods_for_std:
            return np.nan, np.nan # 数据不足，返回NaN

        excess_returns = returns_series - daily_risk_free_rate

        # --- 夏普比率 ---
        mean_excess_return = excess_returns.mean()
        std_dev_returns = returns_series.std() # 夏普比率使用总收益率的标准差

        sharpe = np.nan
        if std_dev_returns != 0:
            sharpe = (mean_excess_return / std_dev_returns) * annualization_factor

        # --- 索提诺比率 ---
        downside_returns = returns_series[returns_series < daily_risk_free_rate]
        sortino = np.nan
        if not downside_returns.empty:
            downside_deviation = np.sqrt(np.mean((downside_returns - daily_risk_free_rate)**2))
            if downside_deviation != 0:
                sortino = (mean_excess_return / downside_deviation) * annualization_factor

        return sharpe, sortino

    # 使用 groupby().rolling() 结合 apply 来高效计算
    # 注意: 这里使用 shift(-N_days) 来获取未来 N_days 的窗口
    # 同时 rolling 窗口需要是 N_days 大小，并对 pct_chg 序列进行操作
    # 结果会被放置在窗口的末尾，但我们实际上需要它对应到窗口的起始位置
    # 因此，我们先计算，然后将其向上 shift N_days + 1 （或者根据 rolling 的行为调整）
    
    # 获取未来 N_days 的每日收益率序列
    # 为了避免在每一行中循环，我们构造一个辅助的DataFrame，其中包含未来N天的收益率
    # 这里使用一个更高效的思路：对每只股票，生成一个N天的收益率列表
    
    # 使用 groupby().apply() 和 list comprehension 来构建未来N天的收益率列表
    # 为了获得未来N天的收益率，需要对pct_chg进行shift(-N_days)
    # 然后再对每个N_days的窗口进行处理
    
    # 方法一：先获取未来N天的收益率列，然后进行滑动窗口计算 （更推荐，效率高）
    # 创建一个辅助列，表示未来 N 天的每日收益率列表
    # 注意：这里需要 N_days 长度的窗口，且数据点至少为 min_periods_for_std
    
    # 对每只股票进行分组，然后计算未来 N 天的夏普/索提诺
    # 这仍然涉及到对每个时间点获取未来N天的收益序列，Pandas 的 rolling 默认是向后看的
    # 为了实现向后看 N_days (即从当前日期 t 预测 t+1 到 t+N_days)，我们需要一些技巧
    
    # 最直接且通常高效的方法是：
    # 1. 计算日收益率。
    # 2. 对每个股票，对日收益率进行反向滚动，然后计算指标。
    #    或者更简单地，使用 shift(-N_days) 创建“未来的”日收益率列，然后进行滚动。
    
    # 为了简化且保持效率，我们可以这样做：
    # 1. 计算每个股票的每日收益率。
    # 2. 对于每个 (ts_code, trade_date) 行，我们将查看其未来的 N_days 每日收益率。
    #    这在 Pandas 中不直接通过 .rolling() 实现，因为 rolling 是“当前及之前”的窗口。
    #    我们可以用 groupby().apply() 配合自定义函数来模拟这个“未来窗口”的行为。

    def _apply_metrics_per_stock(stock_df_group):
        """
        对单只股票的数据帧进行处理，计算其每个日期的未来 N 日夏普/索提诺比率。
        """
        sharpe_labels = [np.nan] * len(stock_df_group)
        sortino_labels = [np.nan] * len(stock_df_group)

        # 遍历当前股票的每个日期，获取未来的 N_days 收益率序列
        for i in range(len(stock_df_group) - N_days):
            # 获取未来 N_days 的每日收益率序列
            # 从当前日期 t 的下一天 (t+1) 开始，到 t+N_days
            future_pct_chgs = stock_df_group['pct_chg'].iloc[i+1 : i+1+N_days].dropna()
            
            # 计算夏普和索提诺
            sharpe, sortino = _calculate_metrics_for_window(future_pct_chgs)
            
            # 赋值到对应日期
            sharpe_labels[i] = sharpe
            sortino_labels[i] = sortino
        
        stock_df_group['sharpe_ratio_label'] = sharpe_labels
        stock_df_group['sortino_ratio_label'] = sortino_labels
        return stock_df_group

    # 对每个股票分组应用这个函数
    df_result = df_copy.groupby('ts_code', group_keys=False).apply(_apply_metrics_per_stock)

    # # 清理不再需要的中间列
    # df_result = df_result.drop(columns=['pct_chg'])

    return df_result