import os
import numpy as np

from skimage import io, img_as_ubyte

from shutil import copy
from collections import Container
from argparse import ArgumentParser


# I would like to take a minute to express that relative imports in Python are horrible.
# Although this function is implemented somewhere else, it cannot be imported, since its
# folder is in the parent folder of this. Relative imports result in ValueErrors. The
# design choice behind this decision eludes me. The only way to circumvent this is either
# make this package installable, add the parent folder to PATH or implement it again.
# I went with the latter one.
#
# If you are reading this and you also hate the relative imports in Python, cheers!
# You are not alone.
def chk_mkdir(*paths: Container) -> None:
    """
    Creates folders if they do not exist.

    Args:
        paths: Container of paths to be created.
    """
    for path in paths:
        if not os.path.exists(path):
            os.makedirs(path)


def merge_masks(masks_folder):
    masks = list()
    for mask_img_filename in os.listdir(masks_folder):
        mask_img = io.imread(os.path.join(masks_folder, mask_img_filename))
        masks.append(mask_img)

    merged_mask = np.sum(masks, axis=0)
    merged_mask[merged_mask > 0] = 1

    return img_as_ubyte(merged_mask)


if __name__ == '__main__':

    parser = ArgumentParser()
    parser.add_argument('--dataset_path', required=True, type=str)
    parser.add_argument('--export_path', required=True, type=str)
    args = parser.parse_args()

    new_images_folder = os.path.join(args.export_path, 'images')
    new_masks_folder = os.path.join(args.export_path, 'masks')

    chk_mkdir(args.export_path, new_images_folder, new_masks_folder)

    for image_name in os.listdir(args.dataset_path):
        images_folder = os.path.join(args.dataset_path, image_name, 'images')
        masks_folder = os.path.join(args.dataset_path, image_name, 'masks')
        # copy the image
        copy(src=os.path.join(images_folder, image_name + '.png'),
             dst=os.path.join(new_images_folder, image_name + '.png'))

        # convert and save the masks
        mask_img = merge_masks(masks_folder)
        io.imsave(os.path.join(new_masks_folder, image_name + '.png'), mask_img)
