import gradio as gr

from utils import load_example, run_nn, load_model_, next_, prev_

demo = gr.Blocks()
import models

with demo:
    headline = gr.Markdown("## Raven resolver ")
    markdown = gr.Markdown("Below we show all 9 images from raven matrix. "
                           "Model gets 8 images and predicts the properties of last one. "
                           "Based on this properties the answer image is render in the  right panel. <br />"
                           "Note that angle rotation is only used as a noise. "
                           "There are not rules applied to angle property, so angle rotation of final output do not need to be the same as in example. "
                           "Additionally there are cases that other properties could be used as noise.")
    with gr.Row():
        with gr.Column():
            with gr.Row():
                text = gr.Textbox(models.START_IMAGE,
                                  label="Write the example number from validation dataset (0, 14,000). You can also paste here matrix representation from generator.")
            with gr.Row():
                prev = gr.Button("Prev")
                show = gr.Button("Show")
                next = gr.Button("Next")
                # button = gr.Button("Run")
            with gr.Row():
                image = gr.Image(value=load_example(models.START_IMAGE)[0], label="Raven matrix")
                desc = gr.Markdown(value=load_example(models.START_IMAGE)[1])
                # desc = gr.Text(value=load_example(models.START_IMAGE)[1])

        with gr.Column():
            with gr.Row():
                output = gr.Image(label="Generated image", shape=(200, 200))
    with gr.Row():
        button = gr.Button("Run")

    # text.change(load_example, inputs=text, outputs=[image, desc])
    show.click(load_example, inputs=text, outputs=[image, desc])
    # button.click(run_nn, inputs=image, outputs=output)
    button.click(run_nn, inputs=text, outputs=output)

    # next.click(next_, inputs=text, outputs=text)
    # next.click(load_example, inputs=text, outputs=[image, desc])
    next.click(next_, inputs=text, outputs=[text, image, desc])

    # prev.click(prev_, inputs=text, outputs=text)
    # prev.click(load_example, inputs=text, outputs=[image, desc])
    prev.click(prev_, inputs=text, outputs=[text, image, desc])

demo.launch(debug=True)