import os
import argparse
import random
import warnings

import math
from shutil import copyfile, move
from tqdm import tqdm
import concurrent.futures
import numpy as np


def copy_file_fnc(s_d):
    s, d = s_d
    file_name = os.path.split(s)[-1]
    copyfile(s, os.path.join(d, file_name))
    return int(os.path.exists(d))


def move_file_fnc(s_d):
    s, d = s_d
    file_name = os.path.split(s)[-1]
    move(s, os.path.join(d, file_name))
    return int(os.path.exists(d))


def split(train_perc: float=0.7,
          validation_perc: float=0.15,
          test_perc: float=0.15,
          data_dir: str='',
          file_format: str='h5',
          random_split: bool=True,
          move_files: bool=False):
    assert train_perc + validation_perc + test_perc == 1.0, 'Train+Validation+Test != 1 (100%)'
    assert train_perc > 0 and test_perc > 0, 'Train and test percentages must be greater than zero'

    file_set = [os.path.join(data_dir, f) for f in os.listdir(data_dir) if f.endswith(file_format)]
    random.shuffle(file_set) if random_split else file_set.sort()

    num_files = len(file_set)
    num_validation = math.floor(num_files * validation_perc)
    num_test = math.floor(num_files * test_perc)
    num_train = num_files - num_test - num_validation

    dataset_root, dataset_name = os.path.split(data_dir)
    dst_train = os.path.join(dataset_root, 'SPLIT_' + dataset_name, 'train_set')
    dst_validation = os.path.join(dataset_root, 'SPLIT_' + dataset_name, 'validation_set')
    dst_test = os.path.join(dataset_root, 'SPLIT_' + dataset_name, 'test_set')

    print('OUTPUT INFORMATION\n=============')
    print('Train:\t\t{}'.format(num_train))
    print('Validation:\t{}'.format(num_validation))
    print('Test:\t\t{}'.format(num_test))
    print('Num. samples\t{}'.format(num_files))
    print('Path:\t\t', os.path.join(dataset_root, 'SPLIT_' + dataset_name))

    dest = [dst_train] * num_train + [dst_validation] * num_validation + [dst_test] * num_test

    os.makedirs(dst_train, exist_ok=True)
    os.makedirs(dst_validation, exist_ok=True)
    os.makedirs(dst_test, exist_ok=True)

    progress_bar = tqdm(zip(file_set, dest), desc='Copying files', total=num_files)
    operation = move_file_fnc if move_files else copy_file_fnc
    desc = 'Moving files' if move_files else 'Copying files'
    with concurrent.futures.ProcessPoolExecutor(max_workers=10) as ex:
        results = list(tqdm(ex.map(operation, zip(file_set, dest)), desc=desc, total=num_files))

    num_copies = np.sum(results)
    if num_copies == num_files:
        print('Done successfully')
    else:
        warnings.warn('Missing files: {}'.format(num_files - num_copies))

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--train', '-t', type=float, default=.70, help='Train percentage. Default: 0.70')
    parser.add_argument('--validation', '-v', type=float, default=0.15, help='Validation percentage. Default: 0.15')
    parser.add_argument('--test', '-s', type=float, default=0.15, help='Test percentage. Default: 0.15')
    parser.add_argument('-d', '--dir', type=str, help='Directory where the data is')
    parser.add_argument('-f', '--format', type=str, help='Format of the data files. Default: h5', default='h5')
    parser.add_argument('-r', '--random', help='Randomly split the dataset or not. Default: True', action='store_true', default=True)
    parser.add_argument('-m', '--movefiles', help='Move files. Otherwise copy. Default: False', action='store_true', default=False)

    args = parser.parse_args()

    split(args.train, args.validation, args.test, args.dir, args.format, args.random, args.movefiles)