import streamlit as st
import matplotlib.pyplot as plt
import networkx as nx
import numpy as np

# Sidebar for selecting an option
sidebar_option = st.sidebar.radio("Select an option", 
                                 ["Select an option", "Basic: Properties", 
                                  "Basic: Read and write graphs", "Basic: Simple graph", 
                                  "Basic: Simple graph Directed", "Drawing: Custom Node Position",
                                  "Drawing: Cluster Layout", "Drawing: Degree Analysis"])

# Helper function to draw and display graph
def draw_graph(G, pos=None, title="Graph Visualization"):
    plt.figure(figsize=(8, 6))
    nx.draw(G, pos=pos, with_labels=True, node_color='lightblue', node_size=500, font_size=10, font_weight='bold')
    st.pyplot(plt)

# Function to display properties and graph for Basic: Properties
def display_graph_properties(G):
    pathlengths = []
    st.write("### Source vertex {target:length, }")
    for v in G.nodes():
        spl = dict(nx.single_source_shortest_path_length(G, v))
        st.write(f"Vertex {v}: {spl}")
        for p in spl:
            pathlengths.append(spl[p])

    avg_path_length = sum(pathlengths) / len(pathlengths)
    st.write(f"### Average shortest path length: {avg_path_length}")

    dist = {}
    for p in pathlengths:
        dist[p] = dist.get(p, 0) + 1
    st.write("### Length #paths")
    for d in sorted(dist.keys()):
        st.write(f"Length {d}: {dist[d]} paths")

    st.write("### Properties")
    st.write(f"Radius: {nx.radius(G)}")
    st.write(f"Diameter: {nx.diameter(G)}")
    st.write(f"Eccentricity: {nx.eccentricity(G)}")
    st.write(f"Center: {nx.center(G)}")
    st.write(f"Periphery: {nx.periphery(G)}")
    st.write(f"Density: {nx.density(G)}")

    # Visualize the graph
    st.write("### Graph Visualization")
    pos = nx.spring_layout(G, seed=3068)  # Seed layout for reproducibility
    draw_graph(G, pos)

# Function to display graph for Basic: Read and write graphs
def display_read_write_graph(G):
    st.write("### Adjacency List:")
    for line in nx.generate_adjlist(G):
        st.write(line)
    
    # Write the graph's edge list to a file
    st.write("### Writing Edge List to 'grid.edgelist' file:")
    nx.write_edgelist(G, path="grid.edgelist", delimiter=":")  # Save edge list
    st.write("Edge list written to 'grid.edgelist'")

    # Read the graph from the edge list
    st.write("### Reading Edge List from 'grid.edgelist' file:")
    H = nx.read_edgelist(path="grid.edgelist", delimiter=":")
    st.write("Edge list read into graph H")

    # Visualize the graph
    st.write("### Graph Visualization:")
    pos = nx.spring_layout(H, seed=200)  # Seed for reproducibility
    draw_graph(H, pos)

# Function to display Simple Graphs for Basic: Simple graph
def display_simple_graph(G, pos=None):
    options = {
        "font_size": 36,
        "node_size": 3000,
        "node_color": "white",
        "edgecolors": "black",
        "linewidths": 5,
        "width": 5,
    }
    
    # Draw the network
    nx.draw_networkx(G, pos, **options)

    # Set margins for the axes so that nodes aren't clipped
    ax = plt.gca()
    ax.margins(0.20)
    plt.axis("off")
    st.pyplot(plt)

# Function to display Simple Directed Graphs for Basic: Simple graph Directed
def display_simple_directed_graph(G, pos=None):
    options = {
        "node_size": 500,
        "node_color": "lightblue",
        "arrowsize": 20,
        "width": 2,
        "edge_color": "gray",
    }
    
    # Draw the directed graph with the given positions and options
    nx.draw_networkx(G, pos, **options)

    # Set margins for the axes so that nodes aren't clipped
    ax = plt.gca()
    ax.margins(0.20)
    plt.axis("off")
    st.pyplot(plt)

# Function to display Custom Node Position Graphs for Drawing: Custom Node Position
def display_custom_node_position():
    st.title("Drawing: Custom Node Position")
    
    # Default example graph (path graph with custom node position)
    G = nx.path_graph(20)
    center_node = 5
    edge_nodes = set(G) - {center_node}
    
    # Ensure the nodes around the circle are evenly distributed
    pos = nx.circular_layout(G.subgraph(edge_nodes))
    pos[center_node] = np.array([0, 0])  # Manually specify node position
    
    # Draw the graph
    draw_graph(G, pos)

# Function to display Cluster Layout for Drawing: Cluster Layout
def display_cluster_layout():
    st.title("Drawing: Cluster Layout")

    G = nx.davis_southern_women_graph()  # Example graph
    communities = nx.community.greedy_modularity_communities(G)

    # Compute positions for the node clusters as if they were themselves nodes in a supergraph using a larger scale factor
    supergraph = nx.cycle_graph(len(communities))
    superpos = nx.spring_layout(G, scale=50, seed=429)

    # Use the "supernode" positions as the center of each node cluster
    centers = list(superpos.values())
    pos = {}
    for center, comm in zip(centers, communities):
        pos.update(nx.spring_layout(nx.subgraph(G, comm), center=center, seed=1430))

        # Nodes colored by cluster
    for nodes, clr in zip(communities, ("tab:blue", "tab:orange", "tab:green")):
        nx.draw_networkx_nodes(G, pos=pos, nodelist=nodes, node_color=clr, node_size=100)
    nx.draw_networkx_edges(G, pos=pos)

    plt.tight_layout()
    st.pyplot(plt)

# Function to display Degree Analysis for Drawing: Degree Analysis
def display_degree_analysis():
    st.title("Drawing: Degree Analysis")

    option = st.radio("Choose a graph type:", ("Default Example", "Create your own"))

    if option == "Default Example":
        G = nx.gnp_random_graph(100, 0.02, seed=10374196)
        
        degree_sequence = sorted((d for n, d in G.degree()), reverse=True)
        dmax = max(degree_sequence)

        fig = plt.figure("Degree of a random graph", figsize=(8, 8))
        # Create a gridspec for adding subplots of different sizes
        axgrid = fig.add_gridspec(5, 4)

        ax0 = fig.add_subplot(axgrid[0:3, :])
        Gcc = G.subgraph(sorted(nx.connected_components(G), key=len, reverse=True)[0])
        pos = nx.spring_layout(Gcc, seed=10396953)
        nx.draw_networkx_nodes(Gcc, pos, ax=ax0, node_size=20)
        nx.draw_networkx_edges(Gcc, pos, ax=ax0, alpha=0.4)
        ax0.set_title("Connected components of G")
        ax0.set_axis_off()

        ax1 = fig.add_subplot(axgrid[3:, :2])
        ax1.plot(degree_sequence, "b-", marker="o")
        ax1.set_title("Degree Rank Plot")
        ax1.set_ylabel("Degree")
        ax1.set_xlabel("Rank")

        ax2 = fig.add_subplot(axgrid[3:, 2:])
        ax2.bar(*np.unique(degree_sequence, return_counts=True))
        ax2.set_title("Degree histogram")
        ax2.set_xlabel("Degree")
        ax2.set_ylabel("# of Nodes")

        fig.tight_layout()
        st.pyplot(fig)

    elif option == "Create your own":
        n_nodes = st.number_input("Number of nodes:", min_value=2, max_value=500, value=100)
        p_edge = st.slider("Edge probability:", min_value=0.0, max_value=1.0, value=0.02)
        
        if st.button("Generate"):
            if n_nodes >= 2:
                G_custom = nx.gnp_random_graph(n_nodes, p_edge, seed=10374196)
                degree_sequence = sorted((d for n, d in G_custom.degree()), reverse=True)
                dmax = max(degree_sequence)

                fig = plt.figure("Degree of a random graph", figsize=(8, 8))
                # Create a gridspec for adding subplots of different sizes
                axgrid = fig.add_gridspec(5, 4)

                ax0 = fig.add_subplot(axgrid[0:3, :])
                Gcc = G_custom.subgraph(sorted(nx.connected_components(G_custom), key=len, reverse=True)[0])
                pos = nx.spring_layout(Gcc, seed=10396953)
                nx.draw_networkx_nodes(Gcc, pos, ax=ax0, node_size=20)
                nx.draw_networkx_edges(Gcc, pos, ax=ax0, alpha=0.4)
                ax0.set_title("Connected components of G")
                ax0.set_axis_off()

                ax1 = fig.add_subplot(axgrid[3:, :2])
                ax1.plot(degree_sequence, "b-", marker="o")
                ax1.set_title("Degree Rank Plot")
                ax1.set_ylabel("Degree")
                ax1.set_xlabel("Rank")

                ax2 = fig.add_subplot(axgrid[3:, 2:])
                ax2.bar(*np.unique(degree_sequence, return_counts=True))
                ax2.set_title("Degree histogram")
                ax2.set_xlabel("Degree")
                ax2.set_ylabel("# of Nodes")

                fig.tight_layout()
                st.pyplot(fig)

# Display Basic: Properties if selected
if sidebar_option == "Basic: Properties":
    st.title("Basic: Properties")
    option = st.radio("Choose a graph type:", ("Default Example", "Create your own"))

    if option == "Default Example":
        G = nx.lollipop_graph(4, 6)
        display_graph_properties(G)

    elif option == "Create your own":
        num_nodes = st.number_input("Number of nodes:", min_value=2, max_value=50, value=5)
        num_edges = st.number_input("Number of edges per group (for lollipop graph):", min_value=1, max_value=10, value=3)

        if st.button("Generate"):
            if num_nodes >= 2 and num_edges >= 1:
                G_custom = nx.lollipop_graph(num_nodes, num_edges)
                display_graph_properties(G_custom)

# Display Basic: Read and write graphs if selected
elif sidebar_option == "Basic: Read and write graphs":
    st.title("Basic: Read and write graphs")
    option = st.radio("Choose a graph type:", ("Default Example", "Create your own"))

    if option == "Default Example":
        G = nx.grid_2d_graph(5, 5)
        display_read_write_graph(G)

    elif option == "Create your own":
        rows = st.number_input("Number of rows:", min_value=2, max_value=20, value=5)
        cols = st.number_input("Number of columns:", min_value=2, max_value=20, value=5)

        if st.button("Generate"):
            if rows >= 2 and cols >= 2:
                G_custom = nx.grid_2d_graph(rows, cols)
                display_read_write_graph(G_custom)

# Display Basic: Simple Graph if selected
elif sidebar_option == "Basic: Simple graph":
    st.title("Basic: Simple graph")
    option = st.radio("Choose a graph type:", ("Default Example", "Create your own"))

    if option == "Default Example":
        G = nx.Graph()
        G.add_edge(1, 2)
        G.add_edge(1, 3)
        G.add_edge(1, 5)
        G.add_edge(2, 3)
        G.add_edge(3, 4)
        G.add_edge(4, 5)

        pos = {1: (0, 0), 2: (-1, 0.3), 3: (2, 0.17), 4: (4, 0.255), 5: (5, 0.03)}
        display_simple_graph(G, pos)

    elif option == "Create your own":
        edges = []
        edge_input = st.text_area("Edges:", value="1,2\n1,3\n2,3")
        if edge_input:
            edge_list = edge_input.split("\n")
            for edge in edge_list:
                u, v = map(int, edge.split(","))
                edges.append((u, v))

        if st.button("Generate"):
            G_custom = nx.Graph()
            G_custom.add_edges_from(edges)
            pos = nx.spring_layout(G_custom, seed=42)
            display_simple_graph(G_custom, pos)

# Display Basic: Simple Directed Graph if selected
elif sidebar_option == "Basic: Simple graph Directed":
    st.title("Basic: Simple graph Directed")
    option = st.radio("Choose a graph type:", ("Default Example", "Create your own"))

    if option == "Default Example":
        G = nx.DiGraph([(0, 3), (1, 3), (2, 4), (3, 5), (3, 6), (4, 6), (5, 6)])

        left_nodes = [0, 1, 2]
        middle_nodes = [3, 4]
        right_nodes = [5, 6]

        pos = {n: (0, i) for i, n in enumerate(left_nodes)}
        pos.update({n: (1, i + 0.5) for i, n in enumerate(middle_nodes)})
        pos.update({n: (2, i + 0.5) for i, n in enumerate(right_nodes)})

        display_simple_directed_graph(G, pos)

    elif option == "Create your own":
        edges = []
        edge_input = st.text_area("Edges:", value="1,2\n1,3\n2,3")
        if edge_input:
            edge_list = edge_input.split("\n")
            for edge in edge_list:
                u, v = map(int, edge.split(","))
                edges.append((u, v))

        if st.button("Generate"):
            G_custom = nx.DiGraph()
            G_custom.add_edges_from(edges)
            pos = nx.spring_layout(G_custom, seed=42)
            display_simple_directed_graph(G_custom, pos)

# Display Drawing: Custom Node Position if selected
elif sidebar_option == "Drawing: Custom Node Position":
    display_custom_node_position()

# Display Drawing: Cluster Layout if selected
elif sidebar_option == "Drawing: Cluster Layout":
    display_cluster_layout()

# Display Drawing: Degree Analysis if selected
elif sidebar_option == "Drawing: Degree Analysis":
    display_degree_analysis()