#Takes the available numpy features and randomly assigns them to either the training
#or test set. This is currently not what we want to do but can be readjusted.

import os
import random
import shutil
import sys

def splitTrainTest(splitRate):
    mainDataDir = '/home-2/fyu12@jhu.edu/work/fyu12/CATARACT/raw_data/cataract_phase_separated/'
    phaseLabelsTrain = open(mainDataDir + '/train_phase_labels.txt', 'w')
    phaseLabelsTest = open(mainDataDir + '/test_phase_labels.txt', 'w')
    for ph in range(1, 11):
        dataDir = mainDataDir + str(ph)
        clips = os.listdir(dataDir)
        for clip in clips:
            if clip[-3:] == 'avi':
                split_name = clip.split('_')
                p = split_name[0][1:]
                n = split_name[1][1:]
                num = split_name[3][:3]
                video_tag = p + '\t' + n + '\t' + num
            
                randNum = random.random()
                if randNum < splitRate:
                    phaseLabelsTrain.write(video_tag + '\n')
                else:
                    phaseLabelsTest.write(video_tag+ '\n')
    
    phaseLabelsTrain.close()
    phaseLabelsTest.close()

if __name__ == "__main__":
    splitRate = eval(sys.argv[1])
    if splitRate > 1:
        splitRate = 0.8
    splitTrainTest(splitRate)
