def run_simulation(num_workers, batch_time, network_latency, mode, num_batches):
    """Wrapper for Gradio interface."""
    timelines, metrics = simulate_batches(
        num_workers=int(num_workers),
        batch_time=float(batch_time),
        network_latency=float(network_latency),
        mode=mode,
        num_batches=int(num_batches)
    )
    fig = plot_timeline(timelines, metrics, num_workers)
    return fig, (
        f"Mode: {mode.capitalize()}\n"
        f"Epoch Time: {metrics['epoch_time_ms']:.2f} ms\n"
        f"Idle Time: {metrics['idle_percent']} %\n"
        f"Throughput: {metrics['throughput']} batches/sec"
    )

# Define Gradio UI
interface = gr.Interface(
    fn=run_simulation,
    inputs=[
        gr.Slider(1, 8, value=4, step=1, label="Number of Workers"),
        gr.Slider(100, 1000, value=500, step=50, label="Batch Processing Time (ms)"),
        gr.Slider(50, 500, value=200, step=25, label="Network Latency (ms)"),
        gr.Radio(["synchronous", "asynchronous"], value="synchronous", label="Mode"),
        gr.Slider(5, 30, value=10, step=1, label="Number of Batches per Epoch"),
    ],
    outputs=[
        gr.Plot(label="Timeline Visualization"),
        gr.Textbox(label="Simulation Summary", lines=6, max_lines=8, show_copy_button=True)
    ],
    title="🧠 Batch Scheduler Simulator",
    description="Visualize how synchronous vs. asynchronous batch scheduling affects throughput, idle time, and epoch duration.",
    examples=[
        [4, 500, 200, "synchronous", 10],
        [8, 400, 150, "asynchronous", 15]
    ]
)

# Launch interface
interface.launch(share=True)