import cv2
import glob
import numpy as np


def stitch_together(input_images, layout, resize_dim=None, off_x=None, off_y=None):
    """
    Stitch together N input images into a bigger frame, using a grid layout.
    Input images can be either color or grayscale, but must all have the same size.

    :param input_images: list of input images
    :param layout: grid layout expressed (rows, cols) of the stitch
    :param resize_dim: if not None, stitch is resized to this size
    :param off_x: offset between stitched images along x axis
    :param off_y: offset between stitched images along y axis
    :return: stitch of input images
    """

    if len(set([img.shape for img in input_images])) > 1:
        raise ValueError('All images must have the same shape')

    # determine if input images are color (3 channels) or grayscale (single channel)
    if len(input_images[0].shape) == 2:
        mode = 'grayscale'
        img_h, img_w = input_images[0].shape
    elif len(input_images[0].shape) == 3:
        mode = 'color'
        img_h, img_w, img_c = input_images[0].shape
    else:
        raise ValueError('Unknown shape for input images')

    # if no offset is provided, set to 10% of image size
    if off_x is None:
        off_x = img_w // 10
    if off_y is None:
        off_y = img_h // 10

    # create stitch mask
    rows, cols = layout
    stitch_h = rows * img_h + (rows + 1) * off_y
    stitch_w = cols * img_w + (cols + 1) * off_x
    if mode == 'color':
        stitch = np.zeros(shape=(stitch_h, stitch_w, img_c), dtype=np.uint8)
    elif mode == 'grayscale':
        stitch = np.zeros(shape=(stitch_h, stitch_w), dtype=np.uint8)

    for r in range(0, rows):
        for c in range(0, cols):

            list_idx =  r * cols + c

            if list_idx < len(input_images):
                if mode == 'color':
                    stitch[ r * (off_y + img_h) + off_y: r*(off_y+img_h) + off_y + img_h,
                            c * (off_x + img_w) + off_x: c * (off_x + img_w) + off_x + img_w,
                            :] = input_images[list_idx]
                elif mode == 'grayscale':
                    stitch[ r * (off_y + img_h) + off_y: r*(off_y+img_h) + off_y + img_h,
                            c * (off_x + img_w) + off_x: c * (off_x + img_w) + off_x + img_w]\
                        = input_images[list_idx]

    if resize_dim:
        stitch = cv2.resize(stitch, dsize=(resize_dim[::-1]))

    return stitch


if __name__ == '__main__':

    img_list = glob.glob('img/*.jpg')

    images = [cv2.imread(f, cv2.IMREAD_COLOR) for f in img_list]

    s = stitch_together(images, layout=(5, 5), resize_dim=(1000, 1000))

    cv2.imshow('s', s)
    cv2.waitKey()




@@ -295,22 +295,45 @@
 e.%0A%0A    
-:p
+P
 aram
+eters%0A    ----------%0A   
  input_i
@@ -341,14 +341,37 @@
 ages
+ 
 : list
+(ndarray)%0A        List
  of 
@@ -391,24 +391,32 @@
     
-:param 
 layout
+ 
 : 
-g
+tuple%0A        G
 rid 
@@ -463,23 +463,16 @@
 tch%0A    
-:param 
 resize_d
@@ -477,73 +477,160 @@
 _dim
+ 
 : 
-if not None, stitch is resized to this size%0A    :param
+tuple, optional%0A        If present, stitch is resized to this size following convention resize_dim=(new_h, new_w)%0A   
  off_x
+ 
 : 
-o
+int, optional%0A        O
 ffse
@@ -676,23 +676,39 @@
     
-:param 
 off_y
+ 
 : 
-o
+int, optional%0A        O
 ffse
@@ -750,30 +750,75 @@
 xis%0A
+%0A
     
-:r
+R
 eturn
-: stitch of
+s%0A    -------%0A    stitch : ndarray%0A        Stitch of the
  inp
@@ -822,24 +822,25 @@
 input images
+.
 %0A    %22%22%22%0A%0A  

Update stitching docs to clarify resizing usage