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

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

def preprocess_data(excel_path: str, output_dir: str, test_size: float, random_seed: int):
    """
    读取Excel数据，进行清洗和格式转换，并划分为训练集和测试集。
    
    Args:
        excel_path: 原始Excel文件的路径。
        output_dir: 存放处理后CSV文件的目录。
        test_size: 测试集占总数据的比例。
        random_seed: 随机种子。
    """
    print(f"--- 1. 读取数据: {excel_path} ---")
    try:
        # 尝试读取Excel文件的第一个工作表
        df = pd.read_excel(excel_path)
    except FileNotFoundError:
        print(f"错误：文件未找到在路径: {excel_path}")
        return
    except Exception as e:
        print(f"读取Excel文件时发生错误: {e}")
        return
    
    # --- 2. 数据清洗与格式转换 ---
    
    # 检查所需的列是否存在
    required_cols = ['sentence', 'Envir']
    if not all(col in df.columns for col in required_cols):
        print(f"错误：Excel中缺少必需的列。找到的列有: {df.columns.tolist()}")
        print(f"必需的列是: {required_cols}")
        return
        
    # 重命名列以符合通用NLP任务格式 (text 和 label)
    df = df.rename(columns={'sentence': 'text', 'Envir': 'label'})
    
    # 确保'label'列是整数类型 (0或1)
    df['label'] = df['label'].astype(int)
    
    # 仅保留 'text' 和 'label' 两列
    df = df[['text', 'label']].dropna()
    print(f"原始数据条数: {len(df)}")
    
    # --- 3. 划分数据集 ---
    print(f"--- 划分数据集 (训练集:{1-test_size}, 测试集:{test_size}) ---")
    
    # 将数据划分为训练集和测试集，使用分层抽样 (stratify) 确保标签比例一致
    train_df, test_df = train_test_split(
        df, 
        test_size=test_size, 
        random_state=random_seed,
        stratify=df['label']
    )
    
    val_size_from_train = 0.1 / (1 - test_size)
    train_df, val_df = train_test_split(
        train_df,
        test_size=val_size_from_train,
        random_state=random_seed,
        stratify=train_df['label']
    )
    
    # --- 4. 保存为CSV文件 ---
    os.makedirs(output_dir, exist_ok=True)
    
    train_output_path = os.path.join(output_dir, 'train.csv')
    val_output_path = os.path.join(output_dir, 'validation.csv')
    test_output_path = os.path.join(output_dir, 'test.csv')

    train_df.to_csv(train_output_path, index=False)
    val_df.to_csv(val_output_path, index=False)
    test_df.to_csv(test_output_path, index=False)
    
    print("--- 结果保存成功 ---")
    print(f"训练集条数: {len(train_df)}. 保存至: {train_output_path}")
    print(f"验证集条数: {len(val_df)}. 保存至: {val_output_path}")
    print(f"测试集条数: {len(test_df)}. 保存至: {test_output_path}")

if __name__ == "__main__":
    preprocess_data(EXCEL_FILE_PATH, OUTPUT_DIR, TEST_SIZE, RANDOM_SEED)