import numpy as np
import cv2

import os

        
class MontageCreator():
    
    def __init__(self, sources_dir, main_image):
        self.sources_dir = sources_dir 
        self.main_image = cv2.imread(main_image)
        nb = input("How many images per axis do you want ?")
        self.grid_size = (int(nb),int(nb))
        self.selection_size = 50
        
    def createMosaic(self):
        """
        Decrease the resolution of the input image
        For every image in the list compute the average value of every channel (3 numbers - RGB)
        For every pixel in the input image with values (r,g,b) do the following: 
            Randomly sample 30 (just a number that works well) images from the list. 
            For every such random image in the sample, 
                compute the distance (*) between the rgb values and choose the image with smallest distance.

        """
        
        self.grid = cv2.resize(self.main_image, self.grid_size, interpolation =cv2.INTER_AREA)
        
        self.__fill_image_list()
        print("Fill image list : OK")
        self.__associate_images()
        print("Associate images : OK")
        self.result = self.__compose_image()
        print("Composition : OK")
        return self.result
    
    def save_image(self, im, path):
       cv2.imwrite(path, im)
        
        
    def __fill_image_list(self):
        height_tile, width_tile = int(self.main_image.shape[0]/self.grid_size[0]), int(self.main_image.shape[1]/self.grid_size[1])
        self.tile_format = (height_tile, width_tile, 3)
        self.im_list = []
        for file in os.listdir(self.sources_dir):
            im = cv2.imread(os.path.join(self.sources_dir, file))
            if im is not None:
                height, width, _ = im.shape
                if height < width:
                    self.im_list.append(cv2.resize(im, (width_tile, height_tile), interpolation =cv2.INTER_AREA))
        self.Nim = len(self.im_list)
        self.ims_mean_RGB = np.zeros((len(self.im_list), 3), dtype=np.uint8)
        for i in range(self.Nim):
            self.ims_mean_RGB[i,:] = np.mean(self.im_list[i],axis=(0, 1))
            
    def __associate_images(self):
        self.index_image = np.zeros((self.grid.shape[0], self.grid.shape[1]), np.int32)
        for i in range(self.grid.shape[0]):
            for j in range(self.grid.shape[1]):
                selection = np.random.randint(0, len(self.im_list), self.selection_size)
                d_min = 255
                ind_min = 0
                for k in range(self.selection_size):
                    d = self.__compute_distance(self.ims_mean_RGB[selection[k]], self.grid[i,j])
                    if d < d_min:
                        ind_min = k
                        d_min = d
                self.index_image[i, j] = selection[ind_min]
        
    def __compute_distance(self,rgb1, rgb2):
        return np.linalg.norm(rgb1.astype(np.int32)-rgb2.astype(np.int32))
    
    def __compose_image(self):
        height_tile, width_tile, _ = self.tile_format
        montage = np.zeros((height_tile*self.grid.shape[0], width_tile*self.grid.shape[1], 3), np.uint8)
        for i in range(self.grid.shape[0]):
            for j in range(self.grid.shape[1]):
                montage[i*height_tile:(i+1)*height_tile, j*width_tile:(j+1)*width_tile,:] = self.im_list[self.index_image[i, j]]
        return montage
                
        
        

            

        
        
    