import gradio as gr
import json
import numpy as np
import datasets

bugs_ds = datasets.load_dataset("asgaardlab/SampleDataset", split="validation")


def generate_annotations(image_index):
    image_index = int(image_index)
    objects_json = bugs_ds[image_index]["Objects JSON (Correct)"]
    objects = json.loads(objects_json)

    segmentation_image_rgb = bugs_ds[image_index]["Segmentation Image (Correct)"]
    segmentation_image_rgb = np.array(segmentation_image_rgb)

    annotations = []
    for obj in objects:
        color = tuple(obj["color"])[:-1]
        mask = np.all(segmentation_image_rgb[:, :, :3] == np.array(color), axis=-1).astype(np.float32)

        annotations.append((mask, obj["labelName"]))

    object_count = 0 # bugs_ds[image_index]["Object Count"]
    victim_name = bugs_ds[image_index]["Victim Name"]
    bug_type = bugs_ds[image_index]["Tag"]

    return (
        (bugs_ds[image_index]["Correct Image"], annotations),
        objects,
        object_count,
        victim_name,
        bug_type,
    )


# Setting up the Gradio interface using blocks API
with gr.Blocks() as demo:
    gr.Markdown(
        "Enter the image index and click **Submit** to view the segmentation annotations."
    )
    with gr.Row():
        inp = gr.Slider(
            minimum=0, maximum=len(bugs_ds) - 1, step=1, label="Image Index"
        )
        btn = gr.Button("Submit")
    with gr.Row():
        with gr.Column():
            object_count = gr.Number(label="Object Count")
            victim_name = gr.Textbox(label="Victim Name")
            bug_type = gr.Textbox(label="Bug Type")

        seg_img = gr.AnnotatedImage()

    with gr.Row():
        json_data = gr.JSON()

    btn.click(
        fn=generate_annotations,
        inputs=inp,
        outputs=[seg_img, json_data, object_count, victim_name, bug_type],
    )

demo.launch()