import numpy as np
from pathlib import Path
import os.path
import ntpath
import os
import csv 
from shutil import copy

# get all subdirectories of a folder
def get_folders(a_dir):
    return [name for name in os.listdir(a_dir)
            if os.path.isdir(os.path.join(a_dir, name))]

# get all files in a folder
def get_files(a_dir):
    a_dir = Path(a_dir)
    return [f for f in a_dir.glob('**/*') if f.is_file()]

subj_num = int(input('Which subject would you like to process? '))

# change to folder where frames of the video are
frame_folder = 'framedata/S{}C2/frames'.format(subj_num)
file_list = get_files(frame_folder)
print(len(file_list))

# change to frame rate of video
framerate = 30000./1001.

# change to folders for training, validation, test sets
train_path = 'framedata/Data/Training/'
valid_path = 'framedata/Data/Validation/'
test_path = 'framedata/Data/Test/'

training = get_folders(train_path)
validation = get_folders(valid_path)
test = get_folders(test_path)

# iterate through each frame
for cur in file_list:
    with open(cur, 'r') as img:
        # name of the frame should be the number frame that it is in the video
        filename = ntpath.basename(cur)
        frame_num = int(filename[:-4])
        frame_num = (frame_num - 1) * framerate
        # For each subset of data, check the label data generated by splitdata.py and insert the 
        # frame file name in a txt file that will contain the filepath of all frames of that subset
        for folder in training:
            train_csvfile = get_files(train_path + folder)
            with open(train_path + 'train_list.txt', 'a', newline='') as output:
                if len(train_csvfile) > 0:
                    with open(train_csvfile[0], 'r', encoding='latin-1', newline='') as source:
                        train_reader = csv.reader(source, delimiter = ',')
                        for trdata in train_reader:
                            if frame_num > float(trdata[0]) and frame_num < float(trdata[1]):
                                output.write(ntpath.basename(train_csvfile[0])[:-4] + '/')
                                output.write(filename + '\n')                       
        for folder in validation:
            valid_csvfile = get_files(valid_path + folder)
            if len(valid_csvfile) > 0:
                with open(valid_csvfile[0], 'r', encoding='latin-1', newline='') as source:
                    valid_reader = csv.reader(source, delimiter = ',')
                    for vdata in valid_reader:
                        if frame_num > float(vdata[0]) and frame_num < float(vdata[1]):
                            with open(valid_path + 'valid_list.txt', 'a', newline='') as output:
                                output.write(ntpath.basename(valid_csvfile[0])[:-4] + '/')
                                output.write(filename + '\n')
        for folder in test:
            test_csvfile = get_files(test_path + folder)
            if len(test_csvfile) > 0:
                with open(test_csvfile[0], 'r', encoding='latin-1', newline='') as source:
                    test_reader = csv.reader(source, delimiter = ',')
                    for tsdata in test_reader:
                        if frame_num > float(tsdata[0]) and frame_num < float(tsdata[1]):
                            with open(test_path + 'test_list.txt', 'a', newline='') as output:
                                output.write(ntpath.basename(test_csvfile[0])[:-4] + '/')
                                output.write(filename + '\n')

# use the text file given by the earlier operations to copy the frames into the correct folder
with open(train_path + 'train_list.txt', 'r') as train_list:
    for line in train_list:
         class_folder = line.rsplit('/',1)[0]
         copy(frame_folder + '/' + line[-11:-1], train_path + class_folder + '/')
         print('Copied {} to training folder {}.'.format(line[-11:-1], train_path + class_folder))
with open(valid_path + 'valid_list.txt', 'r') as valid_list:
    for line in valid_list:
         class_folder = line.rsplit('/',1)[0]
         copy(frame_folder + '/' + line[-11:-1], valid_path + class_folder + '/')
         print('Copied {} to validation folder {}.'.format(line[-11:-1], valid_path + class_folder))
with open(test_path + 'test_list.txt', 'r') as test_list:
    for line in test_list:
         class_folder = line.rsplit('/',1)[0]
         copy(frame_folder + '/' + line[-11:-1], test_path + class_folder + '/')
         print('Copied {} to test folder {}.'.format(line[-11:-1], test_path + class_folder))
