import gradio as gr
import igraph as ig
import plotly.graph_objects as go

PLOTLY_TREE = None

def get_plotly_tree():
    return PLOTLY_TREE

def create_plotly_tree():
    # Create an igraph tree
    g = ig.Graph.Tree(7, 2)  # Example tree
    lay = g.layout('rt')  # Reingold-Tilford layout

    # Plotly setup
    edge_x = []
    edge_y = []
    for edge in g.get_edgelist():
        x0, y0 = lay[edge[0]]
        x1, y1 = lay[edge[1]]
        edge_x.extend([x0, x1, None])
        edge_y.extend([y0, y1, None])

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

    node_x = [lay[k][0] for k in range(len(lay))]
    node_y = [lay[k][1] for k in range(len(lay))]

    node_trace = go.Scatter(
        x=node_x, y=node_y,
        mode='markers',
        hoverinfo='text',
        marker=dict(showscale=False, size=10, color='#850', line_width=2))

    fig = go.FigureWidget(data=[edge_trace, node_trace],
                    layout=go.Layout(
                        showlegend=False,
                        hovermode='closest',
                        margin=dict(b=0, l=0, t=0, r=0),
                        xaxis=dict(showgrid=False, zeroline=False, showticklabels=False),
                        yaxis=dict(showgrid=False, zeroline=False, showticklabels=False))
                    )

    return fig

def print_on_click(trace, points, selector):
    print('Clicked')

def update_point(trace, points, selector):
    c = list(scatter.marker.color)
    s = list(scatter.marker.size)
    for i in points.point_inds:
        c[i] = '#bae2be'
        s[i] = 20
        with f.batch_update():
            scatter.marker.color = c
            scatter.marker.size = s
    

PLOTLY_TREE = create_plotly_tree()
node_trace = PLOTLY_TREE.data[1]
print('node_trace', node_trace)
print(node_trace.on_click)
node_trace.on_click(update_point)

print('PLOTLY_TREE', PLOTLY_TREE)

def setup_interface():
    iface = gr.Interface(
        fn=get_plotly_tree,
        inputs=[],
        outputs=gr.Plot(),
        title="Interactive Tree Visualization"
    )
    return iface

with gr.Blocks() as demo:
    with gr.Column():
        tree_interface = setup_interface()
        textbox = gr.Textbox()

# Launch the interface
if __name__ == "__main__":
    # iface = setup_interface()
    # iface.launch()
    print('before launch')
    demo.launch()