import os
import pickle
from typing import List
from PIL import Image, ImageFile
import numpy as np
from multiprocessing import Pool, Manager

ImageFile.LOAD_TRUNCATED_IMAGES = True

class PictureInfos:
    def __init__(self, image_path, color, used):
        self.image_path = image_path
        self.color = color
        self.used = used


def get_average_color(image_path):
    with Image.open(image_path) as img:
        img = img.convert('RGB')
        np_image = np.array(img)
        avg_color = np.mean(np_image, axis=(0,1))
        return avg_color


def get_color_space_array(target_image_path):
    with Image.open(target_image_path) as img:
        global PIXELS_TO_FILL_WIDTH
        global PIXELS_TO_FILL_HEIGHT
        PIXELS_TO_FILL_WIDTH = int(np.floor(img.size[0]/width_segments_count))
        PIXELS_TO_FILL_HEIGHT = int(np.floor(img.size[1]/height_segments_count))
        img = img.convert('RGB')
        matrix = np.zeros((height_segments_count, width_segments_count, 3))
        for y in range(height_segments_count):
            for x in range(width_segments_count):
                pixels = []
                for i in range(PIXELS_TO_FILL_HEIGHT):
                    for j in range(PIXELS_TO_FILL_WIDTH):
                        r,g,b = (img.getpixel((x*PIXELS_TO_FILL_WIDTH + j,y*PIXELS_TO_FILL_HEIGHT + i)))
                        pixels.append([r,g,b]) 
                matrix[y][x]= np.mean(pixels, axis=0)
        return matrix


def get_picture_names(file_path):
    return os.listdir(file_path)


def get_closest(target, allPictures):
    image: PictureInfos
    closest = np.inf
    closest_path = ''
    for image in allPictures:
        if not image.used:
            color = image.color
            delta = np.sqrt(((target[0]-color[0])**2)+((target[2]-color[1])**2)+((target[2]-color[2])**2))
            if(delta  == 0):
                # image.used = True
                return image.image_path
            if (delta<closest and delta>0) or (delta>closest and delta<0):
                closest = delta
                closest_path = image.image_path
    return closest_path

def create_picture_segment(target, allPictures, starting_layer, end_layer, width_segments_count, PIXELS_TO_FILL_WIDTH, PIXELS_TO_FILL_HEIGHT, FILLING_QUALITY):
    picture_piece = Image.new('RGB', (int(width_segments_count*PIXELS_TO_FILL_WIDTH*FILLING_QUALITY), int((end_layer-starting_layer)*PIXELS_TO_FILL_HEIGHT*FILLING_QUALITY)))
    layer_offset = 0
    for layer in target[int(starting_layer):int(end_layer)]:
        segment_offset = 0
        for segment in layer:
            closest = get_closest(segment, allPictures)
            closest_img = Image.open(closest)
            closest_img = closest_img.resize((PIXELS_TO_FILL_WIDTH*FILLING_QUALITY, PIXELS_TO_FILL_HEIGHT*FILLING_QUALITY))
            picture_piece.paste(closest_img, (segment_offset*PIXELS_TO_FILL_WIDTH*FILLING_QUALITY, layer_offset*PIXELS_TO_FILL_HEIGHT*FILLING_QUALITY))
            segment_offset +=1
        layer_offset += 1
        print(layer_offset,'Layer done')
    return picture_piece




if __name__ == '__main__':
    ### Starting sequence ###
    target_path = input('Input relative path for target image: ')
    Image.open(target_path)
    print('For best results put segments at same value')
    user_input = input('Input segment width count (Higher=more pictures used. Default = 100. Musn\'t exceed pixel of original): ')
    try:
        width_segments_count = int(user_input)
    except:
        width_segments_count = 100

    user_input = input('Input segment height count (Higher=more pictures used. Default = 100. Musn\'t exceed pixel of original): ')
    try:
        height_segments_count = int(user_input)
    except:
        height_segments_count = 100

    user_input = input('Input quality for filling pictures (1=default, 2=double the size): ')
    try:
        FILLING_QUALITY = int(user_input)
    except:
        FILLING_QUALITY = 1
    
    user_input = input('Input worker count: (default 1) Recomended PC\'s cores if you dont want to kill your PC: ')
    try:
        working_size = int(user_input)
    except:
        working_size = 1

    ### Getting all pictures ###

    if 'allPictures.pkl' not in os.listdir():
        source_categories = os.listdir('./101_ObjectCategories')
        allPictures = list()
        print('Getting all pictures')
        for category in source_categories:
            imgs_path = './101_ObjectCategories/'+category
            category_images = get_picture_names(imgs_path)
            for img in category_images:
                img_path = imgs_path +'/'+img
                allPictures.append(PictureInfos(img_path, get_average_color(img_path), False))
            print(category,'done')

        with open('allPictures.pkl', 'wb') as file:
            pickle.dump(allPictures, file)
        print('done')
    else:
        with open('allPictures.pkl', 'rb') as file:
            allPictures = pickle.load(file)

    print('Segmenting target')
    target = get_color_space_array(target_path)
    print('Segmenting done')

    ### Start to create image ###
    # print((int(width_segments_count*PIXELS_TO_FILL_WIDTH*FILLING_QUALITY), int(height_segments_count*PIXELS_TO_FILL_HEIGHT*FILLING_QUALITY)))


    
    working_length = np.floor(len(target)/working_size)

    pool = Pool()
    picture_segments = pool.starmap(create_picture_segment, [(target, allPictures, working_length*i, working_length*(i+1), width_segments_count, PIXELS_TO_FILL_WIDTH, PIXELS_TO_FILL_HEIGHT, FILLING_QUALITY) for i in range(working_size)])
    pool.close()
    pool.join()
    
    i = 0
    new_im = Image.new('RGB', (int(width_segments_count*PIXELS_TO_FILL_WIDTH*FILLING_QUALITY), int(height_segments_count*PIXELS_TO_FILL_HEIGHT*FILLING_QUALITY)))
    for segment in picture_segments:
        new_im.paste(segment, (0, int(i*PIXELS_TO_FILL_HEIGHT*FILLING_QUALITY*working_length)))
        i +=1

    new_im.show()
    new_im.save('final.jpg')