import modal
import gradio as gr 
import numpy as np 

f = modal.Cls.lookup("casa-interior-hf-v3", "DesignModel")

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

#     if (room_type is None or interior_style is None or objects is None): 
#         print('Please select all three (Room type, Interior Style and Objects')
#         return None

#     objects_str = ", ".join(objects)
#     text = f"A {room_type} in style of {interior_style} with {objects_str}."
    
#     result_image = f.inference.remote("tab1", image, text)
#     return result_image

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

    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
    
    result_image = f.inference.remote("tab2", image, text, mask)
    return result_image

def casa_ai_run_tab3(dict=None):
    selected_crop = dict["layers"][0].convert("RGB")
    
    if selected_crop is None: 
        print('Please provide cropped object')
        return None
    
    result_image = f.inference.remote("tab3", None, None, None, selected_crop)
    return result_image

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

    with gr.Tab("ReImagine"):
        with gr.Row():
            with gr.Column():
                inputs = [
                            gr.Image(sources='upload', type="pil", label="Upload"), 
                            gr.Textbox(label="Prompt for redesigning masked object")
                            # gr.Dropdown(
                            #             label="Room type",
                            #             choices=["Living Room", "Bedroom", "Home Office", "Dining Room", "Kitchen", "Bathroom"],
                            #             multiselect=False
                            #         ), 
                            # gr.Dropdown(
                            #             label="Interior Style",
                            #             choices=["Modern", "Traditional", "Contemporary", "Minimalist", "Rustic", "Industrial", "Scandinavian"],
                            #             multiselect=False
                            #         ), 
                            # gr.Dropdown(
                            #             label="Objects",
                            #             choices=["Sofa", "Bed", "Table", "Chair", "Fridge", "Dinning table", "Coffee table"],
                            #             multiselect=True
                            #         ),
                        ]
            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)

        
    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="Prompt 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="pil", 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()

# demo = gr.Interface(casa_ai_run,
#                     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="Prompt to design room"), 
#                               ],
#                     outputs = [
#                               gr.Image(label="Generated room image"),
#                               ],
#                     title = title,
#                     description = description,
#                     )
# demo.launch()