import os
import shutil


def split_list(my_list, ratios):
    total_length = len(my_list)
    lengths = [int(total_length * ratio) for ratio in ratios]

    parts = []
    start = 0
    for length in lengths:
        parts.append(my_list[start:start + length])
        start += length
    return parts


def get_plot_files(filename, plot_path):
    files = []
    for sub_path in ['frame', 'layout']:
        file = os.path.join(plot_path, sub_path, filename.replace('.json', '.png'))
        if os.path.exists(file):
            files.append((sub_path, file))
    for sub_path in ['frame_labels', 'layout_labels']:
        file = os.path.join(plot_path, sub_path, filename.replace('.json', '.txt'))
        if os.path.exists(file):
            files.append((sub_path, file))
    return files


if __name__ == '__main__':
    json_path = '/anonymous/data/fp_json_out'
    plot_path = '/anonymous/data/fp_json_plot'
    out_path = '/anonymous/data/fp_json_split'
    json_files = [f for f in os.listdir(json_path) if f.endswith('.json')]
    all_files = []
    for f in os.listdir(json_path):
        if not f.endswith('.json'):
            continue
        plot_files = get_plot_files(f, plot_path)
        if len(plot_files) != 4:
            continue
        plot_files.insert(0, ('', os.path.join(json_path, f)))
        all_files.append(plot_files)
    train_files, val_files, test_files = split_list(all_files, [0.8, 0.1, 0.1])
    for file_li, name in zip([train_files, val_files, test_files], ['train', 'val', 'test']):
        for files in file_li:
            for sub_path, f in files:
                if not os.path.exists(os.path.join(out_path, sub_path, name)):
                    os.makedirs(os.path.join(out_path, sub_path, name))
                shutil.copy(f, os.path.join(out_path, sub_path, name))
