from utils_html import HTML_TEMPLATE
from io import BytesIO
import gradio as gr 
import numpy as np 
import requests
import modal
import PIL

f_gc = modal.Cls.lookup("casa-interior-gc-v2", "GetProduct")

def casa_ai_run_tab1(image=None, text=None): 
    
    if image is None: 
        print('Please provide image of empty room to design')
        return None

    if text is None: 
        print('Please provide a text prompt')
        return None

    f = modal.Cls.lookup("casa-interior-hf-v4", "DesignModel")
    result_image = f.inference.remote("tab1", image, text)
    return result_image

def casa_ai_run_tab2(dict=None, text=None):
    
    image = dict["background"].convert("RGB")
    mask = dict["layers"][0].convert('L')

    if np.sum(np.array(mask)) == 0: 
        mask = None 
        
    if mask is None: 
        print('Please provide a mask over the object you want to generate again.')
        
    if image is None and text is None: 
        print('Please provide context in form of image, text')
        return None
        
    f_tab2 = modal.Cls.lookup("casa-interior-hf-v3", "DesignModel")
    result_image = f_tab2.inference.remote("tab2", image, text, mask)
    return result_image

def casa_ai_run_tab3(dict=None):

    selected_crop = dict["composite"]
    
    if selected_crop is None: 
        print('Please provide cropped object')
        return None

    selected_crop = PIL.Image.fromarray(selected_crop).convert('RGB')
    results = f_gc.inference.remote(selected_crop)

    return results

def casa_ai_run_tab_sketch(image=None, room_type=None, room_style=None):

    if image is None: 
        print('Please provide a sketch or ketchup image')
        return None

    if room_type is None: 
        print('Please select a room type')
        return None

    if room_style is None: 
        print('Please select a room style')
        return None

    text = f"{room_type}, {room_style}"
    f = modal.Cls.lookup("casa-interior-hf-v6-sketch", "DesignModel")
    result_image = f.inference.remote(image, text)
    return result_image


with gr.Blocks() as casa:
    title = "Casa-AI Demo"
    description = "A Gradio interface to use CasaAI for virtual staging"
    gr.HTML(value=HTML_TEMPLATE, show_label=False)

    with gr.Tab("Reimagine"):
        with gr.Row():
            with gr.Column():
                inputs = [
                            gr.Image(sources='upload', type="pil", label="Upload"), 
                            gr.Textbox(label="Room description.")
                        ]
            with gr.Column():
                outputs = [gr.Image(label="Generated room image")]

        
        submit_btn = gr.Button("Generate!")
        submit_btn.click(casa_ai_run_tab1, inputs=inputs, outputs=outputs)
        gr.Examples(examples=[['example_images/image_0.jpg', 'Living room in bohemian style'], 
                              ['example_images/image_1.jpg', 'A minimalist and scandinavian style living room with black leather sofa'], 
                              ['example_images/image_2.jpg', 'Modern bedroom art deco style']], 
                    inputs=inputs, outputs=outputs, fn=casa_ai_run_tab1, cache_examples=True)

    
    with gr.Tab("Sketch Transform"):
        with gr.Row():
            with gr.Column():
                inputs = [
                            gr.Image(sources='upload', type="numpy", label="Upload"),
                            gr.Dropdown(["Living Room", "Bedroom", "Kitchen"], label="Room Type", info="Select Room Type"),
                            gr.Dropdown(["Modern", "Minimalist", "Scandinavian"], label="Style", info="Interior Style!"),
                ]
                    
            with gr.Column():
                outputs = [gr.Image(label="Image of sketch transformed")]
                
        submit_btn = gr.Button("Transform!")
        submit_btn.click(casa_ai_run_tab_sketch, inputs=inputs, outputs=outputs)
        gr.Examples(examples=[['example_images/sketch01.jpeg', 'Living Room',  'Modern'], 
                              ['example_images/sketch02.jpeg', 'Bedroom', 'Minimalist'], 
                              ['example_images/sketchup.jpg', 'Kitchen', 'Modern']], 
                    inputs=inputs, outputs=outputs, fn=casa_ai_run_tab_sketch, cache_examples=True)

        
    with gr.Tab("Redesign"):
        with gr.Row():
            with gr.Column():
                inputs = [
                            gr.ImageEditor(sources='upload', brush=gr.Brush(colors=["#FFFFFF"]), elem_id="image_upload", type="pil", label="Upload", layers=False, eraser=True, transforms=[]),
                            gr.Textbox(label="Description for redesigning masked object")]
            with gr.Column():
                outputs = [gr.Image(label="Image with new designed object")]
                
        submit_btn = gr.Button("Redesign!")
        submit_btn.click(casa_ai_run_tab2, inputs=inputs, outputs=outputs)

    with gr.Tab("Recommendation"):
        with gr.Row():
            with gr.Column():
                inputs = [
                            gr.ImageEditor(sources='upload', elem_id="image_upload", type="numpy", label="Upload", layers=False, eraser=False, brush=False, transforms=['crop']),
                            ]
            with gr.Column():
                outputs = [gr.Gallery(label="Similar products")]
                
        submit_btn = gr.Button("Find similar products!")
        submit_btn.click(casa_ai_run_tab3, inputs=inputs, outputs=outputs)


casa.launch()