import time
import gradio as gr
import matplotlib.pyplot as plt
import networkx as nx

# ---- Import generators (NO circular imports) ----
from graphGen3 import NetworkGenerator as NetworkGenerator3
from graphGen4 import NetworkGenerator as NetworkGenerator4
from graphGen5 import NetworkGenerator as NetworkGenerator5


# ---- Registry of available generators ----
GENERATOR_MAP = {
    "graphGen3": NetworkGenerator3,
    "graphGen4": NetworkGenerator4,
    "graphGen5": NetworkGenerator5,
    
}


def generate_network(generator_name, size, variant, topology):
    """
    Gradio callback: generate a network using the selected generator.
    """
    GeneratorClass = GENERATOR_MAP[generator_name]

    generator = GeneratorClass(
        size=size,
        variant=variant,
        topology=topology
    )

    start = time.time()
    graph = generator.generate()
    elapsed = time.time() - start

    stats = (
        f"Generator: {generator_name}\n"
        f"Operation Time: {elapsed:.4f} seconds\n"
        f"Nodes: {len(graph.nodes())}\n"
        f"Edges: {len(graph.edges())}"
    )

    # ---- Plot ----
    fig, ax = plt.subplots(figsize=(8, 8))
    pos = {node: (node[1], -node[0]) for node in graph.nodes()}
    nx.draw(
        graph,
        pos,
        ax=ax,
        with_labels=True,
        node_size=300,
        font_size=8
    )
    ax.set_title(f"{generator_name} | {size}, {variant}, {topology}")
    ax.grid(True)

    return fig, stats


# ---- Gradio UI ----
with gr.Blocks(title="Network Generator") as demo:
    gr.Markdown("# Network Generator")

    with gr.Row():
        generator_choice = gr.Dropdown(
            choices=list(GENERATOR_MAP.keys()),
            value="graphGen3",
            label="Generator Logic"
        )

    with gr.Row():
        size = gr.Dropdown(
            choices=["S", "M", "L"],
            value="S",
            label="Size"
        )
        variant = gr.Dropdown(
            choices=["F", "R"],
            value="F",
            label="Variant"
        )
        topology = gr.Dropdown(
            choices=["highly_connected", "bottlenecks", "linear"],
            value="highly_connected",
            label="Topology"
        )

    generate_btn = gr.Button("Generate Network")

    with gr.Row():
        plot_out = gr.Plot(label="Generated Graph")
        stats_out = gr.Textbox(
            label="Statistics",
            lines=6,
            interactive=False
        )

    generate_btn.click(
        fn=generate_network,
        inputs=[generator_choice, size, variant, topology],
        outputs=[plot_out, stats_out]
    )


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