from cv2 import transpose
import numpy as np
import gradio as gr
from Segmentation.segmentation import get_mask,replace_sofa
from StyleTransfer.styleTransfer import create_styledSofa
from PIL import Image
from random import randint


#https://colab.research.google.com/drive/11CtQpSeRBGAuw4TtE_rL470tRo-1X-p2#scrollTo=edGukUHXyymr
#https://colab.research.google.com/drive/1xq33YKf0LVKCkbbUZIoNPzgpR_4Kd0qL#scrollTo=sPuM8Xypjs-c
#https://github.com/dhawan98/Post-Processing-of-Image-Segmentation-using-CRF

def resize_sofa(img):
    """
    This function adds padding to make the orignal image square and 640by640.
    It also returns the orignal ratio of the image, such that it can be reverted later.
    Parameters:
        img = original image
    Return:
        im1 = squared image
        box = parameters to later crop the image to it original ratio
    """
    width, height = img.size 
    idx = np.argmin([width,height])
    newsize = (640, 640) # parameters from test script

    if idx==0:
        img1 = Image.new(img.mode, (height, height), (255, 255, 255))
        img1.paste(img, ((height-width)//2, 0))
        box = ( newsize[0]*(1-width/height)//2,
                0, 
                newsize[0]-newsize[0]*(1-width/height)//2,
                newsize[1])
    else:
        img1 = Image.new(img.mode, (width, width), (255, 255, 255))
        img1.paste(img, (0, (width-height)//2)) 
        box = (0, 
                newsize[1]*(1-height/width)//2, 
                newsize[0], 
                newsize[1]-newsize[1]*(1-height/width)//2)
    im1 = img1.resize(newsize)
    return im1,box

def resize_style(img):
    """
    This function generates a zoomed out version of the style image and resizes it to a 640by640 square.
    Parameters:
        img = image containing the style/pattern
    Return:
        dst = a zoomed-out and resized version of the pattern
    """
    width, height = img.size
    idx = np.argmin([width,height])

    # Makes the image square by cropping
    if idx==0:
        top= (height-width)//2
        bottom= height-(height-width)//2
        left = 0
        right= width
    else:
        left = (width-height)//2
        right = width - (width-height)//2
        top = 0
        bottom = height
    newsize = (640, 640) # parameters from test script
    im1 = img.crop((left, top, right, bottom))

    # Constructs a zoomed-out version
    copies = 8
    resize = (newsize[0]//copies,newsize[1]//copies)
    dst = Image.new('RGB', (resize[0]*copies,resize[1]*copies))
    im2 = im1.resize((resize))
    for row in range(copies):
        im2 = im2.transpose(Image.FLIP_LEFT_RIGHT)
        for column in range(copies):
            im2 = im2.transpose(Image.FLIP_TOP_BOTTOM)
            dst.paste(im2, (resize[0]*row, resize[1]*column))
    dst = dst.resize((newsize))
    return dst

def style_sofa(input_img: np.ndarray, style_img: np.ndarray):
    """
    Styles (all) the sofas in the image to the given style.
    This function uses a transformer to combine the image with the desired style according 
    to a generated mask of the sofas in the image.
    Input:
        input_img = image containing a sofa
        style_img = image containing a style
    Return:
        new_sofa  = image containing the styled sofa
    """
    id = randint(0, 10)
    print('Starting job ', id)
    # preprocess input images to be (640,640) squares to fit requirements of the segmentation model
    input_img,style_img = Image.fromarray(input_img),Image.fromarray(style_img)
    resized_img,box = resize_sofa(input_img)
    resized_style = resize_style(style_img)
    #resized_style.save('resized_style.jpg')
    # generate mask for image
    print('generating mask...')
    mask = get_mask(resized_img)
    #mask.save('mask.jpg')
    # Created a styled sofa
    print('Styling sofa...')
    styled_sofa = create_styledSofa(resized_img,resized_style)
    #styled_sofa.save('styled_sofa.jpg')
    # postprocess the final image
    print('Replacing sofa...')
    new_sofa = replace_sofa(resized_img,mask,styled_sofa)
    new_sofa = new_sofa.crop(box)
    print('Finishing job', id)
    return new_sofa

demo = gr.Interface(
    style_sofa,
    inputs = [gr.inputs.Image(),gr.inputs.Image()],
    outputs = 'image',
    examples= [['figures/sofa_example1.jpg','figures/style_example1.jpg'],
            ['figures/sofa_example1.jpg','figures/style_example2.jpg'],
            ['figures/sofa_example1.jpg','figures/style_example3.jpg'],
            ['figures/sofa_example1.jpg','figures/style_example4.jpg'],
            ['figures/sofa_example1.jpg','figures/style_example5.jpg']],
    title="🛋 Style your sofa 🛋 ",
    description="Customize your sofa to your wildest dreams 💭!\
                \nProvide a picture of your sofa and a desired pattern\
                 or choose one of the examples below. ⬇",
    theme="huggingface",
    # article="**References**\n\n"
    # "<a href='https://www.tensorflow.org/hub/tutorials/tf2_arbitrary_image_stylization' target='_blank'>1. Tutorial to implement Fast Neural Style Transfer using the pretrained model from TensorFlow Hub</a>  \n"
    # "<a href='https://huggingface.co/spaces/luca-martial/neural-style-transfer' target='_blank'>2. The idea to build a neural style transfer application was inspired from this Hugging Face Space </a>"
)

if __name__ == "__main__":
    demo.launch()