"""Stitch and build a mosaic of all images in a directory"""
import argparse
import cv2 as cv
import glob
import numpy as np
import math
import itertools
import matplotlib.pyplot as plt


def mosaic(images_in_path, tile_size, fill_value=0, display=True):
    """
    Mosaic each image in "images_in_path" after resizing them to tile_size,
    all the pixels that don't fit in by the resized tiles
    are filled with fill_value

    :param str images_in_path: Directory with images to mosaic with
    :param tuple tile_size: tuple that each image in the mosaic is resized to
    :param int fill_value: Bounding box line width
    :param boolean display: display the overlaid roi before returning the array
    """
    x_tiles = y_tiles = math.ceil(np.sqrt(len(images_in_path)))
    if tile_size is not None:
        tile_size_x = tile_size[0]
        tile_size_y = tile_size[1]
    else:
        shape = cv.imread(images_in_path[0], cv.IMREAD_ANYDEPTH | cv.IMREAD_ANYCOLOR).shape
        tile_size_x = shape[0]
        tile_size_y = shape[1]
    mosaiced_im = np.ones((x_tiles * tile_size_x, y_tiles * tile_size_y, 3), dtype=np.uint8)
    mosaiced_im = mosaiced_im * fill_value

    indices = list(itertools.product(range(x_tiles), range(y_tiles)))
    for index, im_path in enumerate(images_in_path):
        im_rgb = cv.imread(im_path, cv.IMREAD_ANYDEPTH | cv.IMREAD_ANYCOLOR)
        if tile_size is not None:
            resized = cv.resize(im_rgb, (tile_size_x, tile_size_y))
        else:
            resized = im_rgb
        x, y = indices[index]
        mosaiced_im[
            x * tile_size_x: tile_size_x * (x + 1),
            y * tile_size_y: tile_size_y * (y + 1),
            :] = resized

    plt.imshow(mosaiced_im)
    plt.axis('off')
    if display:
        plt.show()
    return mosaiced_im


def main():
    parser = argparse.ArgumentParser(
        description="Save the mosaic from images")
    parser.add_argument(
        "--im_dir", help="Absolute path to folder name containing images", required=True, type=str)
    parser.add_argument(
        "--tile_size",
        help="[x,y] list of tile size in x, y", nargs='+', type=int, required=False)
    parser.add_argument(
        "--fill_value",
        help="fill the image with zeros or the first intensity at [0,0] in the image",
        type=str, required=False)
    parser.add_argument(
        "--output_png",
        help="Absolute path to folder name to save the roi overlaid images to", required=True, type=str)
    parser.add_argument(
        "--fmt",
        help="Image format", required=True, type=str)
    parser.add_argument(
        "--display",
        help="Display overlaid images, Default False", action='store_true')

    args = parser.parse_args()
    display = args.display
    im_dir = args.im_dir
    output_png = args.output_png
    tile_size = args.tile_size
    fill_value = args.fill_value
    fmt = args.fmt
    images_in_path = glob.glob(im_dir + "*." + fmt)
    if fill_value == "first":
        fill_value = cv.imread(images_in_path[0], cv.IMREAD_GRAYSCALE)[0, 0]
    else:
        fill_value = int(fill_value)
    im_rgb = mosaic(images_in_path, tile_size, fill_value, display=display)
    cv.imwrite(output_png, im_rgb)

if __name__ == '__main__':
    main()
