import gradio as gr
from PIL import Image, ImageOps
import numpy as np


def ui(controller):
    with gr.Blocks() as ui:
        with gr.Row():
            with gr.Column():
                sketch=gr.Sketchpad(sources = 'upload', label='Sketch', type='pil', image_mode="RGB", brush=gr.Brush()) #gr.Image(sources = 'upload', label='Sketch', type = 'pil')
                                
            with gr.Column():
                improved_sketch_view = gr.Image(type="pil", label="Improved Sketch") 
        
            with gr.Column():
                result=gr.Image(type="pil", label="Colored Image")
                

        with gr.Row():
            with gr.Column():
                first_prompt = gr.Textbox(label="Prompt to Improve Sketch", lines=3)
                first_negative_prompt = gr.Textbox(label="Negative prompt", lines=3, value="sketch, lowres, error, extra digit, fewer digits, cropped, worst quality,low quality, normal quality, jpeg artifacts, blurry")
                firts_prompt_helper = gr.Button(value="Prompt Helper", variant="primary")                
                improve_sketch = gr.Button(value="Improve Sketch", variant="primary")

            with gr.Column():
                second_prompt = gr.Textbox(label="Prompt to Coloring Image", lines=3)
                second_negative_prompt = gr.Textbox(label="Negative prompt", lines=3, value="disfigured, extra digit, fewer digits, cropped, worst quality, low quality")
                second_prompt_helper = gr.Button(value="Prompt Helper", variant="primary")
                result_button = gr.Button(value="Paint It", variant="primary")

        with gr.Row():
            gr.Examples(examples=[["./examples/sketch.png", 
                                   "solo, open mouth, white background, standing, tail, full body,bmonochrome, greyscale, from side, no humans, animal, dog, animal focus, realistic, walk on the street, add city on the backstage, city street",
                                   "solo, open mouth, simple background, grey background, standing, full body, monochrome, greyscale, from side, no humans, animal, dog, animal focus, lineart, wolf"],
                                  ["./examples/cat.png",
                                   "solo, simple background, animal ears, monochrome, greyscale, no humans, cat, white background",
                                   "solo, looking at viewer, simple background, white background, monochrome, greyscale, no humans, animal, cat, slit pupils, animal focus, lineart, whiskers"]],
                        inputs=[sketch, first_prompt, second_prompt],
                        label='Examples. Be concrete as much as possible.')
                
                      
        improve_sketch.click(fn=controller.get_first_result,
                             inputs=[sketch, first_prompt, first_negative_prompt], 
                             outputs=improved_sketch_view)

        result_button.click(fn=controller.get_second_result,
                             inputs=[improved_sketch_view, second_prompt, second_negative_prompt],
                             outputs=result)

        firts_prompt_helper.click(fn=controller.get_help_w_prompt,
                                 inputs=[sketch], 
                                 outputs=first_prompt)

        second_prompt_helper.click(fn=controller.get_help_w_prompt,
                                 inputs=[improved_sketch_view],
                                 outputs=second_prompt) 

    ui.queue().launch(debug=True)