def split_data(labeledfile, base_name):
    print('SPLITTING2')
    train_name = ('%s_train_split.csv' % base_name)
    dev_name = ('%s_dev_split.csv' % base_name)
    test_name = ('%s_test_split.csv' % base_name)
    train_file = open(train_name, 'w')
    dev_file = open(dev_name, 'w')
    test_file = open(test_name, 'w')
    train_file.write((','.join(['SUBJECT_ID', 'HADM_ID', 'TEXT', 'LABELS']) + '\n'))
    dev_file.write((','.join(['SUBJECT_ID', 'HADM_ID', 'TEXT', 'LABELS']) + '\n'))
    test_file.write((','.join(['SUBJECT_ID', 'HADM_ID', 'TEXT', 'LABELS']) + '\n'))
    hadm_ids = {}
    for splt in ['train', 'dev', 'test']:
        hadm_ids[splt] = set()
        with open(('%s/%s_full_hadm_ids.csv' % (MIMIC_3_DIR, splt)), 'r') as f:
            for line in f:
                hadm_ids[splt].add(line.rstrip())
    with open(labeledfile, 'r') as lf:
        reader = csv.reader(lf)
        next(reader)
        i = 0
        cur_hadm = 0
        for row in reader:
            if ((i % 10000) == 0):
                print((str(i) + ' read'))
            if (len(row) > 0):
                hadm_id = row[1]
                if (hadm_id in hadm_ids['train']):
                    train_file.write((','.join(row) + '\n'))
                elif (hadm_id in hadm_ids['dev']):
                    dev_file.write((','.join(row) + '\n'))
                elif (hadm_id in hadm_ids['test']):
                    test_file.write((','.join(row) + '\n'))
                i += 1
        train_file.close()
        dev_file.close()
        test_file.close()
    return (train_name, dev_name, test_name)
def split_data(labeledfile, base_name):
    print('SPLITTING2')
    train_name = ('%s_train_split.csv' % base_name)
    dev_name = ('%s_dev_split.csv' % base_name)
    test_name = ('%s_test_split.csv' % base_name)
    train_file = open(train_name, 'w')
    dev_file = open(dev_name, 'w')
    test_file = open(test_name, 'w')
    train_file.write((','.join(['SUBJECT_ID', 'HADM_ID', 'TEXT', 'LABELS']) + '\n'))
    dev_file.write((','.join(['SUBJECT_ID', 'HADM_ID', 'TEXT', 'LABELS']) + '\n'))
    test_file.write((','.join(['SUBJECT_ID', 'HADM_ID', 'TEXT', 'LABELS']) + '\n'))
    hadm_ids = {}
    for splt in ['train', 'dev', 'test']:
        hadm_ids[splt] = set()
        with open(('%s/%s_full_hadm_ids.csv' % (MIMIC_3_DIR, splt)), 'r') as f:
            for line in f:
                hadm_ids[splt].add(line.rstrip())
    with open(labeledfile, 'r') as lf:
        reader = csv.reader(lf)
        next(reader)
        i = 0
        cur_hadm = 0
        for row in reader:
            if ((i % 10000) == 0):
                print((str(i) + ' read'))
            if (len(row) > 0):
                hadm_id = row[1]
                if (hadm_id in hadm_ids['train']):
                    train_file.write((','.join(row) + '\n'))
                elif (hadm_id in hadm_ids['dev']):
                    dev_file.write((','.join(row) + '\n'))
                elif (hadm_id in hadm_ids['test']):
                    test_file.write((','.join(row) + '\n'))
                i += 1
        train_file.close()
        dev_file.close()
        test_file.close()
    return (train_name, dev_name, test_name)