import gradio as gr
import numpy as np
import matplotlib.pyplot as plt
from scipy.ndimage import label

def create_initial_plot(grid_size):
    grid = np.zeros((grid_size, grid_size))
    fig = plt.figure(figsize=(15, 6))
    grid_ax = fig.add_subplot(121)
    graph_ax = fig.add_subplot(122)

    grid_ax.imshow(grid, cmap='Greys', alpha=0.3)
    grid_ax.set_title(f'Site Occupation Probability p = 0.00')

    graph_ax.plot([0], [0], '-b', label='Largest Cluster Size')
    graph_ax.set_xlabel('Occupation Probability p')
    graph_ax.set_ylabel('Spanning Cluster Size Ratio')
    graph_ax.set_title('Phase Transition in 2D Percolation')
    graph_ax.grid(True)
    graph_ax.set_xlim(0, 1)
    graph_ax.set_ylim(0, 1)

    # plt.tight_layout()
    return fig

def get_largest_cluster_spanning_size(grid):
    labeled_array, num_features = label(grid)
    if num_features == 0:
        return np.zeros_like(grid), 0

    sizes = [np.sum(labeled_array == i) for i in range(1, num_features + 1)]

    max_cluster_index = np.argmax(sizes) + 1
    max_cluster_mask = labeled_array == max_cluster_index
    max_cluster_size = sizes[max_cluster_index - 1]

    perc_x = np.intersect1d(labeled_array[0,:],labeled_array[-1,:])
    perc_x = perc_x[np.where(perc_x>0)]
    # perc_y = np.intersect1d(labeled_array[0,:],labeled_array[-1,:])
    # perc_y = perc_y[np.where(perc_y>0)]

    if len(perc_x)>0:# or len(perc_y)>0: # spanning cond.
        return max_cluster_mask, max_cluster_size
    else:
        return max_cluster_mask, 0

def run_percolation_simulation(grid_size, num_samples=10):
    p_step = 0.02
    fine_step = 0.01

    p1 = np.arange(0, 0.56, p_step*2)
    p2 = np.arange(0.56, 0.60 + fine_step, fine_step)
    p3 = np.arange(0.60 + p_step, 0.98 + p_step, p_step)
    steps = np.concatenate((p1, p2, p3))

    p_values = []
    spanning_clusters = []

    fig = plt.figure(figsize=(15, 6))
    grid_ax = fig.add_subplot(121)
    graph_ax = fig.add_subplot(122)

    np.random.seed(42)
    for p in steps:
        total_spanning_cluster_size = 0
        for _ in range(num_samples):
            grid = np.random.random((grid_size, grid_size)) < p
            largest_cluster_mask, spanning_cluster_size = get_largest_cluster_spanning_size(grid)
            total_spanning_cluster_size += spanning_cluster_size
        
        spanning_cluster_ratio = total_spanning_cluster_size / (num_samples * grid_size * grid_size)

        p_values.append(p)
        spanning_clusters.append(spanning_cluster_ratio)

        grid_ax.clear()
        graph_ax.clear()

        grid_ax.imshow(grid, cmap='Greys', alpha=0.3)
        grid_ax.imshow(largest_cluster_mask, cmap='Blues', alpha=0.7)

        grid_ax.set_title(f'Site Occupation Probability p = {p:.2f}')

        graph_ax.plot(p_values, spanning_clusters, '-b', label='Largest Cluster Size')
        graph_ax.set_xlabel('Occupation Probability p')
        graph_ax.set_ylabel('Spanning Cluster Size Ratio')
        graph_ax.set_title('Phase Transition in 2D Percolation')
        graph_ax.grid(True)
        graph_ax.set_xlim(0, 1)
        graph_ax.set_ylim(0, 1)

        # plt.tight_layout()

        yield fig

with gr.Blocks() as demo:
    gr.Markdown("# 2D Site Percolation Simulation")
    gr.Markdown("""
    This simulation shows the formation of clusters in a 2D percolation system as the occupation probability increases.
    Watch how the system undergoes a phase transition around p ≈ 0.593 (critical point).

    - Gray dots: Occupied sites
    - Blue region: Largest connected cluster
    - Spanning Cluster: A cluster spanning from one side to another 
    """)

    with gr.Row():
        plot = gr.Plot(value=create_initial_plot(150))
    with gr.Row():
        grid_size = gr.Slider(
            minimum=20, maximum=200, step=10, value=150,
            label="Grid Size", info="Size of the simulation grid"
        )
    with gr.Row():
        start_btn = gr.Button("Start Simulation")

    start_btn.click(
        fn=run_percolation_simulation,
        inputs=[grid_size],
        outputs=plot,
        show_progress=False
    )
    grid_size.change(
        fn=create_initial_plot,
        inputs=[grid_size],
        outputs=plot
    )

if __name__ == "__main__":
    demo.queue().launch(
        debug=True
    )