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


# print(os.pwd())

species_to_imgpath = {'bird': './descendent_specific_topk_heatmap_withbb_ep=last_024+051',
                      'fish': './descendent_specific_topk_heatmap_withbb_ep=last_024+051',
                      'butterfly': './descendent_specific_topk_heatmap_withbb_ep=last_024+051',
                     }

# this has to be there for each species
imgname_to_filepath = {} # this ignores the extension such as .png
nodename_to_protoIDs = defaultdict(list)

for species, imgpath in species_to_imgpath.items():
    for foldername in os.listdir(imgpath):
        if os.path.isdir(os.path.join(imgpath, foldername)):
            folderpath = os.path.join(imgpath, foldername)
            for filename in os.listdir(folderpath):
                if filename.endswith('png') or filename.endswith('jpg'):
                    filepath = os.path.join(folderpath, filename)
                    imgname_to_filepath[filename] = filepath
                    nodename = filename.split('.')[0].split('-')[0]
                    protoID = filename.split('.')[0].split('-')[1]
                    nodename_to_protoIDs[nodename].append(protoID)


class Node():
    id = 0
    def __init__(self, name):
        self.id = Node.id
        Node.id += 1
        self.name = name
        self.parent = None
        self.children = [] # list of type Node

    def add_child(child):
            self.children.append(child)


name_to_node = {}
id_to_node = {}

def get_root(node):
    root = node
    while node:
        root = node
        node = node.parent

    return root


def get_tree(imgpath):
    
    for foldername in os.listdir(imgpath):
        if os.path.isdir(os.path.join(imgpath, foldername)):
            folderpath = os.path.join(imgpath, foldername)
            node_name = foldername
            child_names = list(set([filename.split('.')[0].split('-')[0] for filename in os.listdir(folderpath)]))

            if node_name in name_to_node:
                node = name_to_node[node_name]
            else:
                node = Node(node_name)
                name_to_node[node_name] = node
                id_to_node[node.id] = node
                
            child_nodes = []
            for child_name in child_names:
                if child_name in name_to_node:
                    child_node = name_to_node[child_name]
                else:
                    child_node = Node(child_name)
                    name_to_node[child_name] = child_node
                    id_to_node[child_node.id] = child_node
                child_node.parent = node
                child_nodes.append(child_node)

            node.children = child_nodes
            
        # To be finished
    return get_root(node)
                 


ROOT = None


def create_binary_tree_edges(root):
    edges = []
    prev = [root]
    while len(prev) > 0:
        new_prev = []
        for node in prev:
            # print(node.children, '\n')
            edges = edges + [(node.id, child.id) for child in node.children]
            new_prev = new_prev + [child for child in node.children if (len(child.children) > 0)]
        prev = new_prev

        # print(edges)
        # print('-*'*20, '\n')

    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)]

    root = ROOT
    edges = create_binary_tree_edges(root)
    # 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=[id_to_node[i].name for i in range(len(layout.coords))],
        # 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

def plot_tree_from_species(species_name):
    # Define the edges in a tree structure
    # edges = [(0, 1), (0, 2), (1, 3), (1, 4), (2, 5), (2, 6)]

    imgpath = species_to_imgpath[species_name]
    print(imgpath)
    root = get_tree(imgpath)

    # root = ROOT
    edges = create_binary_tree_edges(root)
    # 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=[id_to_node[i].name for i in range(len(layout.coords))],
        # 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 gr.Plot(fig)

def set_nodename_to_protoIDs(species_name):
    global nodename_to_protoIDs
    imgpath = species_to_imgpath[species_name]
    for foldername in os.listdir(imgpath):
        if os.path.isdir(os.path.join(imgpath, foldername)):
            folderpath = os.path.join(imgpath, foldername)
            for filename in os.listdir(folderpath):
                if filename.endswith('png') or filename.endswith('jpg'):
                    filepath = os.path.join(folderpath, filename)
                    imgname_to_filepath[filename] = filepath
                    nodename = filename.split('.')[0].split('-')[0]
                    protoID = filename.split('.')[0].split('-')[1]
                    nodename_to_protoIDs[nodename].append(protoID)


def get_protoIDs(nodename):
    return gr.Dropdown(choices=nodename_to_protoIDs[nodename], interactive=True)

def get_nodenames(species_name):
    return gr.Dropdown(choices=list(nodename_to_protoIDs.keys()), interactive=True)

def get_image(nodename, protoID):
    imgname = '-'.join([nodename, protoID]) + '.png'
    filepath = imgname_to_filepath[imgname]
    return gr.Image(filepath)

def species_change(species_name):
    set_nodename_to_protoIDs(species_name)
    return [plot_tree_from_species(species_name), get_nodenames(species_name), get_nodenames(species_name)]

with gr.Blocks() as demo:

    imgpath = species_to_imgpath['bird']
    print(imgpath)
    ROOT = get_tree(imgpath)
    print(ROOT.name)
    gr.Markdown("## Interactive Tree and Image Display")
        
    # with gr.Row():
    #     tree_output = gr.Plot(plot_tree_using_igraph)  # Connect the function directly
    
    # with gr.Row():
    #     with gr.Column():
    #         dropdown_1_nodename = gr.Dropdown(label="Select a node name", choices=list(nodename_to_protoIDs.keys()))
    #         dropdown_1_protos = gr.Dropdown(label="Select a prototype ID", choices=[], allow_custom_value=True)
    #         image_output_1 = gr.Image()
    #     with gr.Column():
    #         dropdown_2_nodename = gr.Dropdown(label="Select a node name", choices=list(nodename_to_protoIDs.keys()))
    #         dropdown_2_protos = gr.Dropdown(label="Select a prototype ID", choices=[], allow_custom_value=True)
    #         image_output_2 = gr.Image()

    # dropdown_1_nodename.change(get_protoIDs, dropdown_1_nodename, dropdown_1_protos)
    # dropdown_1_protos.change(get_image, [dropdown_1_nodename, dropdown_1_protos], image_output_1)
    # dropdown_2_nodename.change(get_protoIDs, dropdown_2_nodename, dropdown_2_protos)
    # dropdown_2_protos.change(get_image, [dropdown_2_nodename, dropdown_2_protos], image_output_2)

    with gr.Row():
        dropdown_species = gr.Dropdown(label="Select a species", choices=list(species_to_imgpath.keys()))
    with gr.Row():
        tree_output = gr.Plot()  # Connect the function directly
    with gr.Row():
        with gr.Column():
            dropdown_1_nodename = gr.Dropdown(label="Select a node name", choices=[])
            dropdown_1_protos = gr.Dropdown(label="Select a prototype ID", choices=[], allow_custom_value=True)
            image_output_1 = gr.Image()
        with gr.Column():
            dropdown_2_nodename = gr.Dropdown(label="Select a node name", choices=[])
            dropdown_2_protos = gr.Dropdown(label="Select a prototype ID", choices=[], allow_custom_value=True)
            image_output_2 = gr.Image()
    # dropdown_species.change(plot_tree_from_species, dropdown_species, tree_output)
    dropdown_species.change(species_change, dropdown_species, [tree_output, dropdown_1_nodename, dropdown_1_nodename])
    # dropdown_species.change(set_nodename_to_protoIDs)
    # dropdown_species.change(get_nodenames, dropdown_species, dropdown_1_nodename)
    # dropdown_species.change(get_nodenames, dropdown_species, dropdown_2_nodename)
    dropdown_1_nodename.change(get_protoIDs, dropdown_1_nodename, dropdown_1_protos)
    dropdown_1_protos.change(get_image, [dropdown_1_nodename, dropdown_1_protos], image_output_1)
    dropdown_2_nodename.change(get_protoIDs, dropdown_2_nodename, dropdown_2_protos)
    dropdown_2_protos.change(get_image, [dropdown_2_nodename, dropdown_2_protos], image_output_2)

    # imgpath = species_to_imgpath['bird']
    # print(imgpath)
    # ROOT = get_tree(imgpath)
    # print(ROOT.name)
    # gr.Markdown("## Interactive Tree and Image Display")
        
    # with gr.Row():
    #     tree_output = gr.Plot(plot_tree_using_igraph)  # Connect the function directly
    
    # with gr.Row():
    #     with gr.Column():
    #         dropdown_3_nodename = gr.Dropdown(label="Select a node name", choices=list(nodename_to_protoIDs.keys()))
    #         dropdown_3_protos = gr.Dropdown(label="Select a prototype ID", choices=[], allow_custom_value=True)
    #         image_output_3 = gr.Image()
    #     with gr.Column():
    #         dropdown_4_nodename = gr.Dropdown(label="Select a node name", choices=list(nodename_to_protoIDs.keys()))
    #         dropdown_4_protos = gr.Dropdown(label="Select a prototype ID", choices=[], allow_custom_value=True)
    #         image_output_4 = gr.Image()

    # dropdown_3_nodename.change(get_protoIDs, dropdown_3_nodename, dropdown_3_protos)
    # dropdown_3_protos.change(get_image, [dropdown_3_nodename, dropdown_3_protos], image_output_3)
    # dropdown_4_nodename.change(get_protoIDs, dropdown_4_nodename, dropdown_4_protos)
    # dropdown_4_protos.change(get_image, [dropdown_4_nodename, dropdown_4_protos], image_output_4)

    
        

# Initialize with placeholder images
# image_output_1.update(display_image_based_on_dropdown_1)
# image_output_2.update(display_image_based_on_dropdown_2)

demo.launch()