'''
1. Accept image via Gradio 
    * image will come out an a numpy array
    Done!
2. resize image 
    * Solved with PIL 
3. split image into blocks 
    * Solved with PIL
4. find color of the square 
    * potentially solved with PIL 
6. Find image with similar color profile
    *need database 
7. Place image in the new output image
    * can be done with PIL
'''

import gradio as gr
import numpy as np
from PIL import Image
from SSIM_PIL import compare_ssim

# Define a function that applies a sepia filter to an image.
def mosiac(input_img):
    pil_img = Image.fromarray(input_img)

    width, height = pil_img.size

    #resizing the image
    if width != height: 
        if width < height:
            #make the dimensions the width
            dims = np.pow(np.floor(np.sqrt(np.sqrt(width))), 4)
        elif height < width: 
            #make dimensions the height
            dims = np.pow(np.floor(np.sqrt(np.sqrt(height))), 4)

    dims = int(dims)
    box = (dims, dims)
    cropped_original = pil_img.resize(box)
    pil_img = cropped_original.copy()

    quant_img = pil_img.quantize(colors = 100)
    
    #putting a grid on the image

    #gridsize
    grid_size = 16

    quant_img = pil_img.quantize(colors = 100)
    
    #putting a grid on the image

    #gridsize
    grid_size = 150

    #size of each image
    img_size = int(width/grid_size)
    
    for column in range(0, grid_size):
        for row in range(0, grid_size):
            box = (int((row - 1) * img_size), int((column - 1) * img_size), int(row * img_size), int(column * img_size))
            tile = quant_img.crop(box)
            tile = tile.quantize(colors=3, method=2)
            tile_rgb = tile.getpalette()[:9]
            first_flag = True
            with open("imageRGBs.txt", "r") as database:
                for line in database: 
                    rgb, img_num = line.split(';')
                    rgb = rgb.split(' ')
                    rgb = list(map(int, rgb))
                    
                    difference = np.abs((tile_rgb[0] - rgb[0]) + (tile_rgb[1] - rgb[1]) + (tile_rgb[2] - rgb[2]))
                    difference = difference + np.abs((tile_rgb[3] - rgb[3]) + (tile_rgb[4] - rgb[4]) + (tile_rgb[5] - rgb[5])) * .5
                    difference = difference + np.abs((tile_rgb[6] - rgb[6]) + (tile_rgb[7] - rgb[7]) + (tile_rgb[8] - rgb[8])) * .333
                    
                    if first_flag:
                        min_diff = difference
                        min_img_num = img_num.replace("\n", "")
                        first_flag = False
                    else: 
                        if (min_diff > difference): 
                            min_diff = difference 
                            min_img_num = img_num.replace("\n", "")
            replace_image = Image.open("images/" + min_img_num)
            replace_image = replace_image.resize((img_size, img_size))
            pil_img.paste(replace_image, box)
    
    ssim_score = compare_ssim(pil_img, cropped_original)
    print(ssim_score)

    return pil_img

# Create a Gradio interface that takes an image input, runs it through the sepia function, and returns the filtered image.
demo = gr.Interface(mosiac, gr.Image(), "image")

# Launch the interface.
if __name__ == "__main__":
    demo.launch(show_error=True)