import gradio as gr
import numpy as np

from src.part1 import create_Gaussian_kernel, create_hybrid_image

input_pictures_examples = [
    ["./data/1a_dog.bmp", "./data/1b_cat.bmp", 5],
    ["./data/2a_motorcycle.bmp", "./data/2b_bicycle.bmp", 5],
    ["./data/3a_plane.bmp", "./data/3b_bird.bmp", 5],
    ["./data/4a_einstein.bmp", "./data/4b_marilyn.bmp", 5],
    ["./data/5a_submarine.bmp", "./data/5b_fish.bmp", 5],
    ["./data/Message+1.png", "./data/Message+2.png", 5],
]


def action_hybrid_image(a, b, c):
    # check input is None
    if a is None or b is None or c is None:
        return None, None, None

    l, h, lh = create_hybrid_image(a / 255, b / 255, c / 255)

    l, h, lh = (l * 255).astype(np.uint8), ((h + 0.5) * 255).astype(np.uint8), (lh * 255).astype(np.uint8),

    return l, h, lh, lh, lh, lh, lh, lh


with gr.Blocks(theme=gr.themes.Soft()) as app:
    gr.Markdown("## Input")
    with gr.Row() as uploader:
        # allow user to upload 2 images TODO: Limit picture size, button: exchange pictures
        image1 = gr.Image(label="Image 1 (Low Frequency)", interactive=True)

        # Gaussian Kernel
        with gr.Column() as k:
            # allow user to select a cutoff frequency
            cutoff = gr.Slider(1, 7, 5, step=1, label="Cutoff Frequency")
            # display the kernel
            kernel = gr.Image(label="Kernel", image_mode="L")

        image2 = gr.Image(label="Image 2 (High Frequency)", interactive=True)

        cutoff.release(create_Gaussian_kernel,
                       inputs=cutoff, outputs=kernel, api_name="draw_gaussian_kernel")

    submit = gr.Button("Submit", variant="primary")

    gr.Markdown("## Output")
    # output
    with gr.Row() as output:
        low_frequencies = gr.Image(label="Low Frequencies")
        hybrid_image = gr.Image(label="Hybrid Image")
        high_frequencies = gr.Image(label="High Frequencies")

    gr.Markdown("## Outputs In Different Sizes")
    with gr.Row() as output:
        hybrid_large = gr.Image(label="Large", show_label=False).style(height=400, width=400)
        hybrid_medium = gr.Image(label="Medium", show_label=False).style(height=300, width=300)
        hybrid_small = gr.Image(label="Small", show_label=False).style(height=200, width=200)
        hybrid_tiny = gr.Image(label="Tiny", show_label=False).style(height=100, width=100)
        hybrid_super_tiny = gr.Image(label="Super Tiny", show_label=False).style(height=50, width=50)

    submit.click(
        action_hybrid_image,
        inputs=[image1, image2, kernel],
        outputs=[
            low_frequencies, high_frequencies, hybrid_image,
            hybrid_large, hybrid_medium, hybrid_small, hybrid_tiny, hybrid_super_tiny
        ]
    )

    gr.Markdown("## Use Examples")

    gr.Examples(
        examples=input_pictures_examples,
        inputs=[image1, image2, cutoff],
        # outputs=txt_3,
        # fn=combine,
        # cache_examples=True, # cache examples to local storage
    )

app.launch()