import pandas as pd
import os
from sklearn.model_selection import train_test_split
BASE_DIR =os.path.dirname(os.path.dirname(os.path.abspath(__file__)))

def train_val_split(train_x_seg_path,train_y_seg_path,val_x_seg_path,val_y_seg_path):
    train_df_x=pd.read_csv(train_x_seg_path,encoding='utf-8')
    train_df_y=pd.read_csv(train_y_seg_path,encoding='utf-8')
    print(len(train_df_x))
    print(len(train_df_y))
    #训练集 验证集 划分
    #train_test_split的可选参数test_size , train_size, random_state(相当于随机数种子）, shuffle（洗牌）, stratify（按照谁的比例）
    X_train,X_val,y_train,y_val = train_test_split(train_df_x,train_df_y,
                                                   test_size=0.002,
                                                   )
    X_train.to_csv(train_x_seg_path,sep='\t',index=None,header=False)
    y_train.to_csv(train_y_seg_path,sep='\t',index=None,header=False)
    X_val.to_csv(val_x_seg_path,sep='\t',index=None,header=False)
    y_val.to_csv(val_y_seg_path,sep='\t',index=None,header=False)

if __name__ == "__main__":
    train_val_split('{}/datasets/train_set.seg_x.txt'.format(BASE_DIR),
                    '{}/datasets/train_set.seg_y.txt'.format(BASE_DIR),
                    '{}/datasets/val_set.seg_x.txt'.format(BASE_DIR),
                    '{}/datasets/val_set.seg_y.txt'.format(BASE_DIR))
