import os
import subprocess

def install(package):
    subprocess.check_call(["pip", "install", package])

# Manually install each required library
install("numpy")
install("networkx")
install("matplotlib")
install("gradio")

import math
import itertools
import numpy as np
import networkx as nx
import matplotlib.pyplot as plt
import gradio as gr

# --- Topological Index Functions ---

def wiener_index(graph):
    """
    Wiener Index: Sum of shortest path distances between all pairs of vertices.
    """
    sp = dict(nx.all_pairs_shortest_path_length(graph))
    total = 0
    for u in sp:
        for v in sp[u]:
            if u < v:
                total += sp[u][v]
    return total

def compute_indices(graph, index_type):
    if index_type == "Wiener Index":
        return wiener_index(graph)
    
    elif index_type == "Randić Index":
        # Randić Index = Σ[1/√(d(u)*d(v))] for every edge (u,v)
        return sum(1 / math.sqrt(graph.degree(u) * graph.degree(v)) for u, v in graph.edges())
    
    elif index_type == "Balaban Index":
        n = graph.number_of_nodes()
        m = graph.number_of_edges()
        if m == 0 or n <= 1:
            return 0
        return (m / (n - 1)) * sum(1 / math.sqrt(graph.degree(u) * graph.degree(v)) for u, v in graph.edges())
    
    elif index_type == "Zagreb Index M1":
        # M1 = Σ[d(v)]² over all vertices
        return sum(d**2 for _, d in graph.degree())
    
    elif index_type == "Zagreb Index M2":
        # M2 = Σ[d(u)*d(v)] for every edge (u,v)
        return sum(graph.degree(u) * graph.degree(v) for u, v in graph.edges())
    
    elif index_type == "Harary Index":
        # H = Σ[1 / d(u,v)] for all distinct vertex pairs
        return sum(1 / nx.shortest_path_length(graph, u, v) 
                   for u, v in itertools.combinations(graph.nodes(), 2))
    
    elif index_type == "Schultz Index":
        # Schultz Index = Σ[(d(u)+d(v))*d(u,v)] over all edges (simplified version)
        return sum((graph.degree(u) + graph.degree(v)) * nx.shortest_path_length(graph, u, v)
                   for u, v in graph.edges())
    
    elif index_type == "Gutman Index":
        # Gutman Index = Σ[d(u)*d(v)*d(u,v)] over all edges
        return sum(graph.degree(u) * graph.degree(v) * nx.shortest_path_length(graph, u, v)
                   for u, v in graph.edges())
    
    elif index_type == "Estrada Index":
        # Estrada Index = Σ(exp(λ)) over all eigenvalues of the adjacency matrix.
        A = nx.adjacency_matrix(graph).todense()
        eigenvalues = np.linalg.eigvals(A)
        return sum(math.exp(ev) for ev in eigenvalues)
    
    elif index_type == "Hosoya Index":
        # Hosoya Index is the number of matchings; here we use the number of edges.
        return graph.number_of_edges()
    
    else:
        return "Invalid Index Type"

# --- Graph Visualization Function ---

def draw_graph(graph, index_type, index_value, is_regular=False, expected_degree=None):
    """
    Draws the graph using a spring layout.
    If is_regular is True, it checks each node:
      - Nodes with degree equal to expected_degree receive a red marker.
      - The plot title also indicates whether the graph is regular or not.
    """
    plt.figure(figsize=(6, 6))
    pos = nx.spring_layout(graph, seed=42)
    
    # Draw the edges
    nx.draw_networkx_edges(graph, pos, edge_color="gray")
    
    # Set up default node color (light blue)
    node_colors = ['lightblue' for _ in graph.nodes()]
    
    regular_flag = None
    if is_regular and expected_degree is not None:
        # Check each node if it meets the expected degree
        regular_flag = all(graph.degree(n) == expected_degree for n in graph.nodes())
        # Draw a red dot on nodes that meet the expected degree.
        for n in graph.nodes():
            if graph.degree(n) == expected_degree:
                x, y = pos[n]
                plt.scatter(x, y, c="red", s=100, zorder=3)
    
    # Construct title text
    title_text = f"{index_type}: {round(index_value, 3)}"
    if is_regular and expected_degree is not None:
        if regular_flag:
            title_text += " | Regular Graph"
        else:
            title_text += " | Not Regular"
    
    plt.title(title_text, fontsize=14)
    
    filename = "graph.png"
    plt.savefig(filename)
    plt.close()
    return filename

# --- Extended Main Processing Function with Regular Graph Feature ---

def process_graph(node_count, edge_count, index_type, custom_edges, is_regular, degree):
    G = nx.Graph()
    if is_regular:
        try:
            n = int(node_count)
            d = int(degree)
            # Validate that the degree is less than the number of nodes.
            if d >= n:
                return "Error: 'Degree per Node' must be less than 'Number of Nodes'.", None
            # Validate that (n*d) is even.
            if (n * d) % 2 != 0:
                return "Error: (Nodes × Degree) must be even for a valid regular graph.", None
            G = nx.random_regular_graph(d, n)
        except Exception as e:
            return f"Error generating regular graph: {e}", None
    elif not custom_edges.strip():
        G = nx.gnm_random_graph(int(node_count), int(edge_count))
    else:
        try:
            edges = [tuple(map(int, e.strip().split("-"))) for e in custom_edges.split(",")]
            all_nodes = set()
            for u, v in edges:
                all_nodes.update([u, v])
            n = max(all_nodes) + 1
            G = nx.Graph()
            G.add_nodes_from(range(n))
            G.add_edges_from(edges)
        except Exception as e:
            return f"Error in custom edges input: {e}", None

    index_value = compute_indices(G, index_type)
    
    # If regular graph mode, pass the expected degree to the drawing function.
    if is_regular:
        graph_img = draw_graph(G, index_type, index_value, is_regular=True, expected_degree=int(degree))
    else:
        graph_img = draw_graph(G, index_type, index_value)
        
    return index_value, graph_img

# --- Gradio Interface Setup ---

with gr.Blocks() as demo:
    gr.Markdown("# Topological Index Calculator with Graph Visualization")
    
    with gr.Row():
        node_count = gr.Number(label="Number of Nodes", value=5, minimum=1)
        edge_count = gr.Number(label="Number of Edges", value=5, minimum=0)
    
    index_type = gr.Dropdown(
        choices=["Wiener Index", "Randić Index", "Balaban Index", "Zagreb Index M1", "Zagreb Index M2", 
                 "Harary Index", "Schultz Index", "Gutman Index", "Estrada Index", "Hosoya Index"],
        label="Select Topological Index"
    )
    
    custom_edges = gr.Textbox(label="Custom Edges (e.g., 0-1,1-2,2-3)", placeholder="Leave blank for random graph")

    with gr.Row():
        regular_graph_checkbox = gr.Checkbox(label="Generate Regular Graph?", value=False)
        degree_input = gr.Number(label="Degree per Node", value=2, minimum=1, visible=False)

    def toggle_degree_input(is_checked):
        return gr.update(visible=is_checked)

    regular_graph_checkbox.change(
        toggle_degree_input,
        inputs=regular_graph_checkbox,
        outputs=degree_input
    )

    calc_button = gr.Button("Calculate & Visualize")
    result_box = gr.Textbox(label="Computed Index Value", interactive=False)
    graph_output = gr.Image(label="Graph Visualization", interactive=False)
    
    calc_button.click(
        fn=process_graph,
        inputs=[node_count, edge_count, index_type, custom_edges, regular_graph_checkbox, degree_input],
        outputs=[result_box, graph_output]
    )

# --- Run the App ---

if __name__ == "__main__":
    demo.launch()