import pandas as pd
from sklearn.model_selection import train_test_split
import os

# --- 配置参数 ---
EXCEL_FILE_PATH = "/home/hsichen/part_time/BERT_finetune/标注数据_更正后.xlsx"
OUTPUT_DIR = "./processed_data_task2_fixed"
# 划分数据集的比例 (训练集:测试集)
TEST_SIZE = 0.2
# 随机种子，用于确保每次划分结果一致
RANDOM_SEED = 42

def preprocess_data(excel_path: str, output_dir: str, test_size: float, random_seed: int):
    """
    读取Excel数据，进行清洗和格式转换，并划分为训练集、验证集和测试集。
    使用标签数量 (Label Count) 进行分层抽样。
    """
    print(f"--- 1. 读取数据: {excel_path} ---")
    try:
        df = pd.read_excel(excel_path)
    except FileNotFoundError:
        print(f"错误：文件未找到在路径: {excel_path}")
        return
    except Exception as e:
        print(f"读取Excel文件时发生错误: {e}")
        return
    
    # --- 2. 数据清洗与格式转换 ---
    
    # 1. 筛选数据：只保留 Envir=1 的行
    df = df[df['Envir'] == 1].copy()
    print(f"筛选 Envir=1 后数据条数: {len(df)}")
    
    # 2. 整合标签
    TAG_COLS = ['Data', 'Action', 'Gain', 'Regu', 'Vague']
    
    # 将标签列转换为列表
    df['labels'] = df[TAG_COLS].values.tolist()
    
    df = df.rename(columns={'sentence': 'text'})

# 3. 统计标签组合及其个数 (用于分析，保留逻辑)
    print("--- 3. 标签组合类型统计 ---")
    
    # 将标签列表转换为元组
    df['label_tuple'] = df['labels'].apply(tuple)
    
    # 将元组转换为字符串，作为 train_test_split 的 stratify 参数
    # 【新增/修改】：创建用于分层的字符串列
    df['stratify_col'] = df['label_tuple'].astype(str) 
    
    print("-" * 30)
    
    # 仅保留 'text', 'labels', 'stratify_col' 列用于划分 (注意不再需要 'label_count')
    df = df[['text', 'labels', 'stratify_col']].copy()
    
    # --- 4. 划分数据集 (使用 stratify_col 进行分层) ---
    print(f"--- 划分数据集 (训练集:{1-test_size}, 测试集:{test_size}) ---")
    
    # 第一次划分：训练集+验证集 vs 测试集
    train_val_df, test_df = train_test_split(
        df, 
        test_size=test_size, 
        random_state=random_seed,
        # 【关键修改】：使用 'stratify_col' 进行分层
        stratify=df['stratify_col'] 
    )
    
    # 第二次划分：训练集 vs 验证集
    val_size_from_train = 0.1 / (1 - test_size)
    
    train_df, val_df = train_test_split(
        train_val_df,
        test_size=val_size_from_train,
        random_state=random_seed,
        # 【关键修改】：使用 'stratify_col' 进行分层
        stratify=train_val_df['stratify_col']
    )
    
    # --- 5. 保存为CSV文件 ---
    os.makedirs(output_dir, exist_ok=True)
    
    # 保存时仅保留 BERT 需要的 'text' 和 'labels' 列
    train_df[['text', 'labels']].to_csv(os.path.join(output_dir, 'train.csv'), index=False)
    val_df[['text', 'labels']].to_csv(os.path.join(output_dir, 'validation.csv'), index=False)
    test_df[['text', 'labels']].to_csv(os.path.join(output_dir, 'test.csv'), index=False)
    
    print("--- 结果保存成功 ---")
    print(f"训练集条数: {len(train_df)}. 保存至: {os.path.join(output_dir, 'train.csv')}")
    print(f"验证集条数: {len(val_df)}. 保存至: {os.path.join(output_dir, 'validation.csv')}")
    print(f"测试集条数: {len(test_df)}. 保存至: {os.path.join(output_dir, 'test.csv')}")

if __name__ == "__main__":
    # 请确保您已安装必要的库: pip install pandas openpyxl scikit-learn
    preprocess_data(EXCEL_FILE_PATH, OUTPUT_DIR, TEST_SIZE, RANDOM_SEED)