import pandas as pd
import numpy as np
from tqdm import tqdm

def _prepare_concept_df(concept_dict: dict) -> pd.DataFrame:
    """将 concept_dict 转换为长格式的 DataFrame。"""
    records = []
    for date_str, inner_dict in concept_dict.items():
        trade_date = pd.to_datetime(date_str, format='%Y%m%d')
        for concept_name, stock_list in inner_dict.items():
            for ts_code in stock_list:
                records.append((trade_date, concept_name, ts_code))
    
    if not records:
        return pd.DataFrame(columns=['trade_date', 'concept_name', 'ts_code'])
        
    concept_df = pd.DataFrame(records, columns=['trade_date', 'concept_name', 'ts_code'])
    concept_df = concept_df.drop_duplicates(subset=["trade_date", "ts_code"], keep="first")

    return concept_df

def generate_concept_factors(df: pd.DataFrame, concept_dict: dict) -> pd.DataFrame:
    """
    基于热门概念数据生成因子。

    Args:
        df (pd.DataFrame): 所有股票所有日期的数据，需包含 'ts_code', 'trade_date'
                           以及用于聚合的列 (如 'pct_chg', 'turnover_rate')。
        concept_dict (dict): 每日热门概念及其股票列表的字典。

    Returns:
        pd.DataFrame: 添加了概念相关因子的原始 DataFrame。
    """
    print("开始生成概念相关因子...")
    
    # 0. 准备工作，创建副本以避免修改原始df
    df = df.copy()
    df['trade_date'] = pd.to_datetime(df['trade_date'])

    # 1. 将 concept_dict 转换为适合合并的 DataFrame
    concept_df = _prepare_concept_df(concept_dict).sort_values(by=['trade_date'])
    if concept_df.empty:
        print("警告: concept_dict 为空或格式不正确，无法生成概念因子。")
        return df

    # 2. 将概念信息合并到主数据 df
    df = pd.merge(df, concept_df, on=['trade_date', 'ts_code'], how='left')
    # --- 因子计算 ---

    # 因子 1: 是否属于当日热门概念
    df['cat_hot_concept_stock'] = df['concept_name'].notna().astype(np.int8)

    # 因子 2 & 3: 概念内的截面排序因子
    # 创建一个掩码，只对热门概念股进行后续计算，以提高效率
    hot_mask = df['concept_name'].notna()
    
    # 定义需要在概念内部进行截面排序的特征列表
    # 确保这些列存在于你的 df 中
    features_to_rank = ['pct_chg', 'turnover_rate', 'volume_ratio'] 
    
    # 筛选出 df 中实际存在的特征列
    existing_features_to_rank = [f for f in features_to_rank if f in df.columns]
    if not existing_features_to_rank:
        print("警告: df 中缺少用于概念内排序的特征列，跳过相关因子计算。")
    else:
        print(f"开始计算概念内截面排序因子，基于: {existing_features_to_rank}")
        
        # 使用 groupby().rank() 高效计算截面排名
        grouped = df[hot_mask].groupby(['trade_date', 'concept_name'])
        
        for feature in tqdm(existing_features_to_rank, desc="Ranking Features in Concepts"):
            # 计算百分比排名 (0到1之间)，值越大表示排名越靠前
            rank_col_name = f'concept_rank_{feature}'
            df[rank_col_name] = grouped[feature].rank(pct=True)

    # --- 清理 & 返回 ---
    # `concept_name` 列包含了有用的信息，可以选择保留或删除
    # 这里我们选择保留，以便后续分析。如果不需要，可以取消下面这行注释。
    df.drop(columns=['concept_name'], inplace=True)
    
    print("概念相关因子生成完毕。")
    return df