from pathlib import Path

import pandas as pd
workspace_dir = Path(__file__).parent.parent.parent

# 读取你合并后的总表 (253,430条)
df = pd.read_csv(workspace_dir / ".data/ug/validated.tsv", sep='\t')

# 1. 获取所有唯一的说话人
all_speakers = df['client_id'].unique()

# 2. 随机打乱说话人顺序
import random
random.seed(42) # 固定随机种子，保证实验可重复
random.shuffle(all_speakers)

# 3. 挑选验证集说话人，直到录音总数达到 ~8000 条
val_indices = []
val_count = 0
target_val_size = 8000

for speaker in all_speakers:
    speaker_data = df[df['client_id'] == speaker]
    val_indices.extend(speaker_data.index.tolist())
    val_count += len(speaker_data)
    if val_count >= target_val_size:
        break

# 4. 划分文件
df_val = df.loc[val_indices]
df_train = df.drop(val_indices)

print(f"训练集条数: {len(df_train)}")
print(f"验证集条数: {len(df_val)}")

df_train.to_csv(workspace_dir / ".data/ug/train_new.tsv", sep='\t', index=False)
df_val.to_csv(workspace_dir / ".data/ug/val_new.tsv", sep='\t', index=False)
