# 数据读取和划分
# 读取./Task2.xlsx文件, 其中"Yes or No"sheet中为判断数据集，"Factoid"sheet中为填空数据集
# 两类数据均有200条，列名为"序号"、"问题"、"答案"、文本
# 每大类的200条数据分为5小类问题，每类40条，按序号顺序依次排列（如地质灾害发育特征、地质灾害经济损失评估等）
# 考虑到构建数据集的时候，连续相邻问题的文本可能来源于同一段，因此切分测试集和训练集的时候不能直接按原始顺序划分，要随机抽取构建初始数据集。
# 随机抽取时，每小类20条测试+20训练条训练（每大类共计100测试100训练）。


def split_data(data_path, output_path):
    import pandas as pd
    import random
    # 读取数据
    data_yes_no = pd.read_excel(data_path, sheet_name='Yes or No')
    data_factoid = pd.read_excel(data_path, sheet_name='Factoid')

    # 划分数据 (Yes or No)
    train_data_yes_no = []
    test_data_yes_no = []
    for category in range(5):
        category_data = data_yes_no[category * 40:(category + 1) * 40]
        random.seed(42)
        train = category_data.sample(n=20, random_state=42)
        test = category_data.drop(train.index)
        train_data_yes_no.append(train)
        test_data_yes_no.append(test)

    # 划分数据 (Factoid)
    train_data_factoid = []
    test_data_factoid = []
    for category in range(5):
        category_data = data_factoid[category * 40:(category + 1) * 40]
        random.seed(42)
        train = category_data.sample(n=20, random_state=42)
        test = category_data.drop(train.index)
        train_data_factoid.append(train)
        test_data_factoid.append(test)

    # 合并两个sheet的数据
    train_data_yes_no = pd.concat(train_data_yes_no)
    test_data_yes_no = pd.concat(test_data_yes_no)
    train_data_factoid = pd.concat(train_data_factoid)
    test_data_factoid = pd.concat(test_data_factoid)

    # 创建一个 Excel writer 对象
    writer = pd.ExcelWriter(output_path, engine='xlsxwriter')

    # 将数据写入不同的 sheet
    train_data_yes_no.to_excel(writer, sheet_name='Yes or No Train', index=False)
    test_data_yes_no.to_excel(writer, sheet_name='Yes or No Test', index=False)
    train_data_factoid.to_excel(writer, sheet_name='Factoid Train', index=False)
    test_data_factoid.to_excel(writer, sheet_name='Factoid Test', index=False)

    # 保存 Excel 文件
    writer.close()
    print(f"数据已保存到 {output_path}")

if __name__ == "__main__":
    import pandas as pd
    import random
    data_path = './data/Task2.xlsx'
    output_path = './data/train_test_data.xlsx'
    split_data(data_path, output_path)