from glob import glob
import sys
import os
import numpy as np
import random
from mpunet.utils import create_folders
import argparse


def get_parser():
    parser = argparse.ArgumentParser(description="Prepare a data folder for a"
                                                 "CV experiment setup.")
    parser.add_argument("--data_dir", type=str,
                        help="Path to data directory")
    parser.add_argument("--CV", type=int, default=5,
                        help="Number of splits (default=5)")
    parser.add_argument("--out_dir", type=str, default="views",
                        help="Directory to store CV subfolders "
                             "(default=views")
    parser.add_argument("--im_sub_dir", type=str, default="images",
                        help="Subfolder under 'data_dir' in which image are "
                             "stored (default=images)")
    parser.add_argument("--lab_sub_dir", type=str, default="labels",
                        help="Subfolder under 'data_dir' in which labels are "
                             "stored (default=labels)")
    parser.add_argument("--copy", action="store_true",
                        help="Copy files to CV-subfolders instead of "
                             "symlinking (not recommended)")
    parser.add_argument("--file_list", action="store_true",
                        help="Create text files with paths pointing to the "
                             "images at the image and labels subdirs under "
                             "each split instead of symlink/copying. This is"
                             " usefull on systems were symlink is not "
                             "supported, but the dataset size is too large to"
                             " store in copies. NOTE: Only one of --copy and "
                             "--file_list flags must be set.")
    parser.add_argument("--file_regex", type=str, default="*.nii*",
                        help="Regex used to select files from the image "
                             "and labels subdirs. (default='*.nii*')")
    parser.add_argument("--validation_fraction", type=float, default=0.20,
                        help="Fraction of OVERALL data size used for "
                             "validation in each split. In a 5-CV setting with "
                             "N=100 and val_frac=0.20, each split will have "
                             "N_train=60, N_val=20 and N_test=20 images")
    parser.add_argument("--test_fraction", type=float, default=0.20,
                        help="Fraction of data size used for test if CV=1.")
    parser.add_argument("--common_prefix_length", type=int, required=False, default=0)
    return parser


def assert_dir_structure(data_dir, im_dir, lab_dir, out_dir):
    for _dir in (data_dir, im_dir, lab_dir):
        if not os.path.exists(_dir):
            raise OSError("Invalid data directory '%s'. Does not exist." % data_dir)
    if os.path.exists(out_dir):
        raise OSError("Output directory at '%s' already exists." % out_dir)


def create_view_folders(out_dir, n_splits):
    if not os.path.exists(out_dir):
        print("Creating directory at %s" % out_dir)
        os.makedirs(out_dir)

    if n_splits > 1:
        for i in range(n_splits):
            split_dir = os.path.join(out_dir, "split_%i" % i)
            print("Creating directory at %s" % split_dir)
            os.mkdir(split_dir)


def pair_by_names(images, common_prefix_length):
    if common_prefix_length == 0:
        return images
    from collections import defaultdict
    names = [os.path.split(i)[-1][:common_prefix_length] for i in images]
    inds = defaultdict(list)
    for i, item in enumerate(names):
        inds[item].append(i)
    pairs = inds.values()
    return [tuple(np.array(images)[i]) for i in pairs]


def add_images(images, im_folder_path, label_folder_path, im_dir, lab_dir,
               link_func=os.symlink):
    for image in images:
        if not isinstance(image, (list, tuple, np.ndarray)):
            image = (image,)
        for im in image:
            # Get file name
            file_name = os.path.split(im)[-1]

            # Get label path (OBS: filenames must match!)
            lab = im.replace(im_dir, lab_dir)

            if not os.path.exists(lab):
                raise OSError("No label file found at '%s'. OBS: image and "
                              "label files must have exactly the same name. "
                              "Images should be located at '%s' and labels at"
                              " '%s'" % (lab, im_folder_path, label_folder_path))

            # Get relative paths
            rel_image = os.path.relpath(im, im_folder_path)
            rel_label = os.path.relpath(lab, label_folder_path)

            # Symlink or copy
            link_func(rel_image, im_folder_path + "/%s" % file_name)
            link_func(rel_label, label_folder_path + "/%s" % file_name)


def _add_to_file_list_fallback(rel_image_path, image_path,
                               fname="LIST_OF_FILES.txt"):

    """
    On some system synlinks are not supported, if --files_list flag is set,
    uses this function to add each absolute file path to a list at the final
    subfolder that is supposed to store images and label links or actual files

    At run-time, these files must be loaded by reading in the path from these
    files instead.
    """
    # Get folder where list of files should be stored
    folder = os.path.split(image_path)[0]

    # Get absolute path to image
    # We change dir to get the correct abs path from the relative
    os.chdir(folder)
    abs_file_path = os.path.abspath(rel_image_path)

    # Get path to the list of files
    list_file_path = os.path.join(folder, fname)

    with open(list_file_path, "a") as out_f:
        out_f.write(abs_file_path + "\n")


def entry_func(args=None):

    # Get parser
    parser = vars(get_parser().parse_args(args))

    # Get arguments
    data_dir = os.path.abspath(parser["data_dir"])
    n_splits = int(parser["CV"])
    if n_splits > 1:
        out_dir = os.path.join(data_dir, parser["out_dir"], "%i_CV" % n_splits)
    else:
        out_dir = os.path.join(data_dir, parser["out_dir"], "fixed_split")
    im_dir = os.path.join(data_dir, parser["im_sub_dir"])
    lab_dir = os.path.join(data_dir, parser["lab_sub_dir"])

    copy = parser["copy"]
    file_list = parser["file_list"]
    regex = parser["file_regex"]
    val_frac = parser["validation_fraction"]
    test_frac = parser["test_fraction"]
    common_prefix_length = parser["common_prefix_length"]

    if n_splits == 1 and not test_frac:
        raise ValueError("Must specify --test_fraction with --CV=1.")
    if copy and file_list:
        raise ValueError("Only one of --copy and --file_list "
                         "flags must be set.")

    # Assert suitable folders
    assert_dir_structure(data_dir, im_dir, lab_dir, out_dir)

    # Create sub-folders
    create_view_folders(out_dir, n_splits)

    # Get images and pair by subject identifier if common_prefix_length > 0
    images = glob(os.path.join(im_dir, regex))
    images = pair_by_names(images, common_prefix_length)
    print("-----")
    print("Found {} images".format(len(images)))

    # Get validation size
    N_total = len(images)
    if n_splits > 1:
        N_test = N_total // n_splits
    else:
        N_test = int(np.ceil(N_total * test_frac))
    N_val = int(np.ceil(N_total * val_frac))
    if N_val + N_test >= N_total:
        raise ValueError("Too large validation_fraction - "
                         "No training samples left!")
    N_train = N_total - N_test - N_val
    print("Total images:".ljust(40), N_total)
    print("Train images pr. split:".ljust(40), N_train)
    print("Validation images pr. split:".ljust(40), N_val)
    print("Test images pr. split:".ljust(40), N_test)

    # Shuffle and split the images into CV parts
    random.shuffle(images)
    splits = np.array_split(images, n_splits)

    # Symlink / copy files
    for i, split in enumerate(splits):
        print("  Split %i/%i" % (i+1, n_splits), end="\r", flush=True)

        # Set root path to split folder
        if n_splits > 1:
            split_path = os.path.join(out_dir, "split_%i" % i)
        else:
            split_path = out_dir
            # Here we kind of hacky force the following code to work with CV=1
            # Define a test set and overwrite the current split (which stores
            # add the data, as splits was never split with n_splits=1
            split = splits[0][:N_test]

            # Overwrite the splits variable to a length 2 array with the
            # remaining data which will be used as val+train. The loop still
            # refers to the old split and thus will only execute once
            splits = [split, splits[0][N_test:]]

        # Define train, val and test sub-dirs
        train_path = os.path.join(split_path, "train")
        train_im_path = os.path.join(train_path, parser["im_sub_dir"])
        train_label_path = os.path.join(train_path, parser["lab_sub_dir"])
        if N_val:
            val_path = os.path.join(split_path, "val")
            val_im_path = os.path.join(val_path, parser["im_sub_dir"])
            val_label_path = os.path.join(val_path, parser["lab_sub_dir"])
        else:
            val_path, val_im_path, val_label_path = (None,) * 3
        test_path = os.path.join(split_path, "test")
        test_im_path = os.path.join(test_path, parser["im_sub_dir"])
        test_label_path = os.path.join(test_path, parser["lab_sub_dir"])

        # Create folders if not existing
        create_folders([train_path, val_path, train_im_path, train_label_path,
                        val_im_path, val_label_path, test_path, test_im_path,
                        test_label_path])

        # Copy or symlink?
        if copy:
            from shutil import copyfile
            move_func = copyfile
        elif file_list:
            move_func = _add_to_file_list_fallback
        else:
            move_func = os.symlink

        # Add test data to test folder
        add_images(split, test_im_path, test_label_path, im_dir, lab_dir, move_func)

        # Join remaining splits into train+val
        remaining = [x for ind, x in enumerate(splits) if ind != i]
        remaining = [item for sublist in remaining for item in sublist]

        # Extract validation data from the remaining
        random.shuffle(remaining)
        validation = remaining[:N_val]
        training = remaining[N_val:]

        # Add
        if validation:
            add_images(validation, val_im_path, val_label_path, im_dir, lab_dir, move_func)
        add_images(training, train_im_path, train_label_path, im_dir, lab_dir, move_func)


if __name__ == "__main__":
    entry_func()

[((11, 13, 12, 72), 'argparse.ArgumentParser', 'argparse.ArgumentParser', (), '', False, 'import argparse\n'), ((55, 7, 55, 30), 'os.path.exists', 'os.path.exists', ({(55, 22, 55, 29): 'out_dir'}, {}), '(out_dir)', False, 'import os\n'), ((76, 11, 76, 28), 'collections.defaultdict', 'defaultdict', ({(76, 23, 76, 27): 'list'}, {}), '(list)', False, 'from collections import defaultdict\n'), ((126, 4, 126, 20), 'os.chdir', 'os.chdir', ({(126, 13, 126, 19): 'folder'}, {}), '(folder)', False, 'import os\n'), ((127, 20, 127, 51), 'os.path.abspath', 'os.path.abspath', ({(127, 36, 127, 50): 'rel_image_path'}, {}), '(rel_image_path)', False, 'import os\n'), ((130, 21, 130, 48), 'os.path.join', 'os.path.join', ({(130, 34, 130, 40): 'folder', (130, 42, 130, 47): 'fname'}, {}), '(folder, fname)', False, 'import os\n'), ((142, 15, 142, 50), 'os.path.abspath', 'os.path.abspath', ({(142, 31, 142, 49): "parser['data_dir']"}, {}), "(parser['data_dir'])", False, 'import os\n'), ((148, 13, 148, 57), 'os.path.join', 'os.path.join', ({(148, 26, 148, 34): 'data_dir', (148, 36, 148, 56): "parser['im_sub_dir']"}, {}), "(data_dir, parser['im_sub_dir'])", False, 'import os\n'), ((149, 14, 149, 59), 'os.path.join', 'os.path.join', ({(149, 27, 149, 35): 'data_dir', (149, 37, 149, 58): "parser['lab_sub_dir']"}, {}), "(data_dir, parser['lab_sub_dir'])", False, 'import os\n'), ((193, 4, 193, 26), 'random.shuffle', 'random.shuffle', ({(193, 19, 193, 25): 'images'}, {}), '(images)', False, 'import random\n'), ((194, 13, 194, 45), 'numpy.array_split', 'np.array_split', ({(194, 28, 194, 34): 'images', (194, 36, 194, 44): 'n_splits'}, {}), '(images, n_splits)', True, 'import numpy as np\n'), ((60, 11, 60, 34), 'os.path.exists', 'os.path.exists', ({(60, 26, 60, 33): 'out_dir'}, {}), '(out_dir)', False, 'import os\n'), ((62, 8, 62, 28), 'os.makedirs', 'os.makedirs', ({(62, 20, 62, 27): 'out_dir'}, {}), '(out_dir)', False, 'import os\n'), ((122, 13, 122, 38), 'os.path.split', 'os.path.split', ({(122, 27, 122, 37): 'image_path'}, {}), '(image_path)', False, 'import os\n'), ((145, 18, 145, 79), 'os.path.join', 'os.path.join', ({(145, 31, 145, 39): 'data_dir', (145, 41, 145, 58): "parser['out_dir']", (145, 60, 145, 78): "'%i_CV' % n_splits"}, {}), "(data_dir, parser['out_dir'], '%i_CV' % n_splits)", False, 'import os\n'), ((147, 18, 147, 74), 'os.path.join', 'os.path.join', ({(147, 31, 147, 39): 'data_dir', (147, 41, 147, 58): "parser['out_dir']", (147, 60, 147, 73): '"""fixed_split"""'}, {}), "(data_dir, parser['out_dir'], 'fixed_split')", False, 'import os\n'), ((171, 18, 171, 45), 'os.path.join', 'os.path.join', ({(171, 31, 171, 37): 'im_dir', (171, 39, 171, 44): 'regex'}, {}), '(im_dir, regex)', False, 'import os\n'), ((182, 16, 182, 43), 'numpy.ceil', 'np.ceil', ({(182, 24, 182, 42): 'N_total * val_frac'}, {}), '(N_total * val_frac)', True, 'import numpy as np\n'), ((216, 21, 216, 54), 'os.path.join', 'os.path.join', ({(216, 34, 216, 44): 'split_path', (216, 46, 216, 53): '"""train"""'}, {}), "(split_path, 'train')", False, 'import os\n'), ((217, 24, 217, 70), 'os.path.join', 'os.path.join', ({(217, 37, 217, 47): 'train_path', (217, 49, 217, 69): "parser['im_sub_dir']"}, {}), "(train_path, parser['im_sub_dir'])", False, 'import os\n'), ((218, 27, 218, 74), 'os.path.join', 'os.path.join', ({(218, 40, 218, 50): 'train_path', (218, 52, 218, 73): "parser['lab_sub_dir']"}, {}), "(train_path, parser['lab_sub_dir'])", False, 'import os\n'), ((225, 20, 225, 52), 'os.path.join', 'os.path.join', ({(225, 33, 225, 43): 'split_path', (225, 45, 225, 51): '"""test"""'}, {}), "(split_path, 'test')", False, 'import os\n'), ((226, 23, 226, 68), 'os.path.join', 'os.path.join', ({(226, 36, 226, 45): 'test_path', (226, 47, 226, 67): "parser['im_sub_dir']"}, {}), "(test_path, parser['im_sub_dir'])", False, 'import os\n'), ((227, 26, 227, 72), 'os.path.join', 'os.path.join', ({(227, 39, 227, 48): 'test_path', (227, 50, 227, 71): "parser['lab_sub_dir']"}, {}), "(test_path, parser['lab_sub_dir'])", False, 'import os\n'), ((230, 8, 232, 41), 'mpunet.utils.create_folders', 'create_folders', ({(230, 23, 232, 40): '[train_path, val_path, train_im_path, train_label_path, val_im_path,\n    val_label_path, test_path, test_im_path, test_label_path]'}, {}), '([train_path, val_path, train_im_path, train_label_path,\n    val_im_path, val_label_path, test_path, test_im_path, test_label_path])', False, 'from mpunet.utils import create_folders\n'), ((251, 8, 251, 33), 'random.shuffle', 'random.shuffle', ({(251, 23, 251, 32): 'remaining'}, {}), '(remaining)', False, 'import random\n'), ((53, 15, 53, 35), 'os.path.exists', 'os.path.exists', ({(53, 30, 53, 34): '_dir'}, {}), '(_dir)', False, 'import os\n'), ((66, 24, 66, 61), 'os.path.join', 'os.path.join', ({(66, 37, 66, 44): 'out_dir', (66, 46, 66, 60): "'split_%i' % i"}, {}), "(out_dir, 'split_%i' % i)", False, 'import os\n'), ((68, 12, 68, 31), 'os.mkdir', 'os.mkdir', ({(68, 21, 68, 30): 'split_dir'}, {}), '(split_dir)', False, 'import os\n'), ((102, 24, 102, 59), 'os.path.relpath', 'os.path.relpath', ({(102, 40, 102, 42): 'im', (102, 44, 102, 58): 'im_folder_path'}, {}), '(im, im_folder_path)', False, 'import os\n'), ((103, 24, 103, 63), 'os.path.relpath', 'os.path.relpath', ({(103, 40, 103, 43): 'lab', (103, 45, 103, 62): 'label_folder_path'}, {}), '(lab, label_folder_path)', False, 'import os\n'), ((181, 21, 181, 49), 'numpy.ceil', 'np.ceil', ({(181, 29, 181, 48): 'N_total * test_frac'}, {}), '(N_total * test_frac)', True, 'import numpy as np\n'), ((202, 25, 202, 62), 'os.path.join', 'os.path.join', ({(202, 38, 202, 45): 'out_dir', (202, 47, 202, 61): "'split_%i' % i"}, {}), "(out_dir, 'split_%i' % i)", False, 'import os\n'), ((220, 23, 220, 54), 'os.path.join', 'os.path.join', ({(220, 36, 220, 46): 'split_path', (220, 48, 220, 53): '"""val"""'}, {}), "(split_path, 'val')", False, 'import os\n'), ((221, 26, 221, 70), 'os.path.join', 'os.path.join', ({(221, 39, 221, 47): 'val_path', (221, 49, 221, 69): "parser['im_sub_dir']"}, {}), "(val_path, parser['im_sub_dir'])", False, 'import os\n'), ((222, 29, 222, 74), 'os.path.join', 'os.path.join', ({(222, 42, 222, 50): 'val_path', (222, 52, 222, 73): "parser['lab_sub_dir']"}, {}), "(val_path, parser['lab_sub_dir'])", False, 'import os\n'), ((75, 13, 75, 29), 'os.path.split', 'os.path.split', ({(75, 27, 75, 28): 'i'}, {}), '(i)', False, 'import os\n'), ((80, 18, 80, 34), 'numpy.array', 'np.array', ({(80, 27, 80, 33): 'images'}, {}), '(images)', True, 'import numpy as np\n'), ((90, 24, 90, 41), 'os.path.split', 'os.path.split', ({(90, 38, 90, 40): 'im'}, {}), '(im)', False, 'import os\n'), ((95, 19, 95, 38), 'os.path.exists', 'os.path.exists', ({(95, 34, 95, 37): 'lab'}, {}), '(lab)', False, 'import os\n')]