import gradio as gr
from datasets import load_from_disk
from PIL import Image
import io
import base64
import json

from graph_visualization import visualize_graph

# Load the dataset
dataset = load_from_disk(
    "Zaherrr/OOP_KG_Dataset"
)


def reshape_json_data_to_fit_visualize_graph(graph_data):

    nodes = graph_data["nodes"]
    edges = graph_data["edges"]

    transformed_nodes = [
        {"id": nodes["id"][idx], "label": nodes["label"][idx]}
        for idx in range(len(nodes["id"]))
    ]

    transformed_edges = [
        {"source": edges["source"][idx], "target": edges["target"][idx], "type": "->"}
        for idx in range(len(edges["source"]))
    ]

    # print(f"transformed nodes = {transformed_nodes}")

    graph_data = {"nodes": transformed_nodes, "edges": transformed_edges}

    return graph_data


def display_example(index):
    example = dataset[index]
    # print("This is the example: ")
    # print(example)
    # Get the image
    img = example["image"]

    # Prepare the graph data
    graph_data = {"nodes": example["nodes"], "edges": example["edges"]}

    # # Convert graph_data to JSON string
    # json_data = json.dumps(graph_data)
    transformed_graph_data = reshape_json_data_to_fit_visualize_graph(graph_data)

    # print(json_data)
    # Generate the graph visualization
    graph_html = visualize_graph(transformed_graph_data)

    return img, graph_html


def create_interface():
    with gr.Blocks() as demo:
        gr.Markdown("# Knowledge Graph Visualizer")

        with gr.Row():
            index_slider = gr.Slider(
                minimum=0, maximum=len(dataset) - 1, step=1, label="Example Index"
            )

        with gr.Row():
            image_output = gr.Image(type="pil", label="Image")
            graph_output = gr.HTML(label="Knowledge Graph")

        index_slider.change(
            fn=display_example,
            inputs=[index_slider],
            outputs=[image_output, graph_output],
        )

    return demo


# Create and launch the interface
if __name__ == "__main__":
    demo = create_interface()
    demo.launch(debug=True)