import igraph as ig
import plotly.graph_objects as go

import gradio as gr
import plotly.graph_objects as go
import os
from collections import defaultdict
import igraph as ig

# Let's create a list of edges with a binary tree structure up to depth 6
# A binary tree has a structure where each node has two children, often referred to as left and right child.
# We will assume that the topmost node (root) has an index 0.

class binNode():
    def __init__(self, id) -> None:
        self.id = id
        self.child1 = None
        self.child2 = None

def create_binary_tree_edges(depth):
    edges = []

    # build binary tree
    id = 0
    root = binNode(id)
    prev = [root]
    for _ in range(depth):
        new_prev = []
        for node in prev:
            id += 1
            node.child1 = binNode(id)
            edges.append((node.id, node.child1.id))
            id += 1
            node.child2 = binNode(id)
            edges.append((node.id, node.child2.id))
            new_prev += [node.child1, node.child2]
        prev = new_prev

    return edges



def plot_tree_using_igraph():
    # Define the edges in a tree structure
    # edges = [(0, 1), (0, 2), (1, 3), (1, 4), (2, 5), (2, 6)]
    edges = create_binary_tree_edges(8)
    # edges = [(str(n1), str(n2)) for (n1, n2) in edges]

    print(edges)
    
    # Create an igraph Graph from the edge list
    g = ig.Graph(edges, directed=True)
    
    # Validate the root index
    if g.vcount() > 0:  # Check if the graph has any vertices
        root_vertex_id = 0  # This assumes that vertex '0' is the root
    else:
        print("The graph has no vertices.")
        return None

    # Use the Reingold-Tilford layout to position the nodes
    try:
        layout = g.layout_reingold_tilford(root=None)  # Correct root specification
    except Exception as e:
        print(f"Error computing layout: {e}")
        return None

    # Edge traces
    edge_x = []
    edge_y = []
    for edge in edges:
        start_idx, end_idx = edge
        x0, y0 = layout.coords[start_idx]
        x1, y1 = layout.coords[end_idx]
        edge_x.extend([x0, x1, None])
        edge_y.extend([-y0, -y1, None])  # y values are inverted to make the tree top-down

    edge_trace = go.Scatter(
        x=edge_x, y=edge_y,
        line=dict(width=0.5, color='#888'),
        hoverinfo='none',
        mode='lines')

    # Node traces
    node_x = [pos[0] for pos in layout.coords]
    node_y = [-pos[1] for pos in layout.coords]  # y values are inverted

    node_trace = go.Scatter(
        x=node_x, y=node_y,
        text=["Node {}".format(i) for i in range(len(layout.coords))],
        mode='markers+text',
        hoverinfo='text',
        marker=dict(
            showscale=False,
            size=10,
            color='LightSkyBlue'
        ),
        textposition="bottom center"
    )

    # Create a Plotly figure
    fig = go.Figure(data=[edge_trace, node_trace],
                    layout=go.Layout(
                        title='<b>Tree Layout with iGraph and Plotly</b>',
                        titlefont_size=16,
                        showlegend=False,
                        hovermode='closest',
                        margin=dict(b=0, l=0, r=0, t=50),
                        xaxis=dict(showgrid=False, zeroline=False, showticklabels=False),
                        yaxis=dict(showgrid=False, zeroline=False, showticklabels=False),
                        # height=600,
                        # width=600,
                        annotations=[dict(
                            showarrow=False,
                            xref="paper", yref="paper",
                            x=0.005, y=-0.002 )]
                    ))

    return fig

# Try generating the figure
# fig = plot_tree_using_igraph()
# print(fig)
# # if fig is not None:
# #     fig.show()

# from plotly.offline import plot
# plot(fig)

with gr.Blocks() as demo:

    gr.Markdown("## Interactive Tree and Image Display")
    with gr.Row():
        tree_output = gr.Plot(plot_tree_using_igraph, scale=2)  # Connect the function directly


demo.launch()