import gradio as gr
import networkx as nx
import matplotlib.pyplot as plt
import random
import time
from graphGen7 import NetworkGenerator



def get_preset_dims(preset_mode, topology):
    """Calculates dimensions based on preset and topology."""
    if preset_mode == "Custom":
        # Enable inputs, keep current values (passed via state in UI, but here we just return interaction update)
        return gr.update(interactive=True), gr.update(interactive=True)
    
    # Preset Logic
    if topology == "linear":
        if preset_mode == "Small": dims = (4, 4)
        elif preset_mode == "Medium": dims = (6, 11)
        else: dims = (10, 26)
    else: # Highly Connected / Bottlenecks
        if preset_mode == "Small": dims = (4, 4)
        elif preset_mode == "Medium": dims = (8, 8)
        else: dims = (16, 16)
        
    return gr.update(value=dims[0], interactive=False), gr.update(value=dims[1], interactive=False)

def update_void_settings(variant, width, height):
    """Calculates void fraction based on variant logic."""
    if variant == "Custom":
        return gr.update(interactive=True)
    
    # Fixed Logic
    area = width * height
    val = 0.60 if area <= 20 else 0.35
    return gr.update(value=val, interactive=False)

def generate_network_viz(topology, width, height, variant, void_frac):
    try:
        variant_code = "F" if variant == "Fixed" else "R"
        
        generator = NetworkGenerator(
            width=width,
            height=height,
            variant=variant_code,
            topology=topology,
            node_drop_fraction=void_frac
        )
        
        start = time.time()
        graph = generator.generate()
        end = time.time()
        
        # Plotting
        fig, ax = plt.subplots(figsize=(8, 8))
        pos = {node: (node[0], node[1]) for node in graph.nodes()}
        
        nx.draw_networkx_edges(graph, pos, ax=ax, width=2, alpha=0.6, edge_color="#333")
        nx.draw_networkx_nodes(graph, pos, ax=ax, node_size=350, node_color="#4F46E5", edgecolors="white", linewidths=1.5)
        nx.draw_networkx_labels(graph, pos, ax=ax, font_size=7, font_color="white", font_weight="bold")
        
        ax.set_xlim(-1, width + 1)
        ax.set_ylim(-1, height + 1)
        ax.invert_yaxis() # Camera looks top-down (0 at top)
        ax.grid(True, linestyle=':', alpha=0.3)
        ax.set_axis_on()
        ax.tick_params(left=True, bottom=True, labelleft=False, labelbottom=False)
        ax.set_title(f"{topology.upper()} | {width}x{height} | Voids: {int(void_frac*100)}%")
        
        # Metrics
        info_text = (
            f"**Nodes:** {len(graph.nodes())} | "
            f"**Edges:** {len(graph.edges())} | "
            f"**Density:** {nx.density(graph):.2f} | "
            f"**Time:** {end - start:.3f}s"
        )
        
        return fig, info_text
    
    except Exception as e:
        return None, f"Error: {str(e)}"


with gr.Blocks(title="Spatial Network Generator") as demo:
    gr.Markdown("# Spatial Network Generator")
    gr.Markdown("Generate procedural room-like graphs with topological constraints.")
    
    with gr.Row():
        # --- LEFT COLUMN: CONTROLS ---
        with gr.Column(scale=1):
            topology_dd = gr.Dropdown(
                choices=["highly_connected", "bottlenecks", "linear"], 
                value="highly_connected", 
                label="Topology"
            )
            
            preset_radio = gr.Radio(
                choices=["Small", "Medium", "Large", "Custom"], 
                value="Medium", 
                label="Preset Size"
            )
            
            with gr.Row():
                width_num = gr.Number(value=8, label="Width (X)", precision=0, interactive=False)
                height_num = gr.Number(value=8, label="Height (Y)", precision=0, interactive=False)

            gr.Markdown("---")
            
            variant_dd = gr.Dropdown(
                choices=["Fixed", "Custom"], 
                value="Fixed", 
                label="Variant",
                info="Fixed = Standard density. Custom = Random density."
            )
            
            void_slider = gr.Slider(
                minimum=0.0, maximum=0.9, value=0.35, step=0.05, 
                label="Void Fraction", 
                interactive=False,
                info="Percentage of grid to leave empty."
            )
            
            gen_btn = gr.Button("Generate Network", variant="primary")

        # --- RIGHT COLUMN: VISUALIZATION ---
        with gr.Column(scale=2):
            metrics_out = gr.Markdown("Click Generate to see metrics...")
            plot_out = gr.Plot(label="Network Visualization")

    # ==========================================
    # EVENT LISTENERS
    # ==========================================
    
    # 1. Update Dimensions when Preset or Topology changes
    #    Note: We pass inputs to calculate new dims, and output to width/height inputs
    input_group_dims = [preset_radio, topology_dd]
    output_group_dims = [width_num, height_num]
    
    preset_radio.change(fn=get_preset_dims, inputs=input_group_dims, outputs=output_group_dims)
    topology_dd.change(fn=get_preset_dims, inputs=input_group_dims, outputs=output_group_dims)

    # 2. Update Void Settings when Variant or Dimensions change
    #    (If user switches to Fixed, we lock slider. If dimensions change while Fixed, we recalc 0.60 vs 0.35)
    input_group_void = [variant_dd, width_num, height_num]
    
    variant_dd.change(fn=update_void_settings, inputs=input_group_void, outputs=void_slider)
    width_num.change(fn=update_void_settings, inputs=input_group_void, outputs=void_slider)
    height_num.change(fn=update_void_settings, inputs=input_group_void, outputs=void_slider)
    
    # 3. Main Generation
    input_group_gen = [topology_dd, width_num, height_num, variant_dd, void_slider]
    output_group_gen = [plot_out, metrics_out]
    
    gen_btn.click(fn=generate_network_viz, inputs=input_group_gen, outputs=output_group_gen)

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