import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation

# Parameters
N_PARTICLES = 40
VOLUME_FRACTION = 0.05
WIDTH, HEIGHT = 100, 100
K = 2.0  # Kinetic constant for growth rate
DT = 0.1 # Time step

# Initialize particles: random positions and sizes
# Small initial sizes to satisfy the "small fraction" requirement
radii = np.random.uniform(1.0, 3.0, N_PARTICLES)
positions = np.random.uniform(0, [WIDTH, HEIGHT], size=(N_PARTICLES, 2))

fig, ax = plt.subplots(figsize=(6, 6))
ax.set_xlim(0, WIDTH)
ax.set_ylim(0, HEIGHT)
ax.set_aspect('equal')
ax.set_title("Ostwald Coarsening Simulation (2D)")

# Create circle patches
circles = [plt.Circle(pos, r, color='royalblue', ec='black') for pos, r in zip(positions, radii)]
for circle in circles:
    ax.add_patch(circle)

def update(frame):
    global radii
    
    # Calculate critical radius (mean radius)
    # Only consider particles that haven't dissolved (r > 0)
    active = radii > 0.1
    if not np.any(active):
        return circles

    r_mean = np.mean(radii[active])
    
    # LSW-style growth law: dR/dt = K/R * (1/R_mean - 1/R)
    # This ensures larger particles grow and smaller ones shrink
    dr = (K / radii) * (1/r_mean - 1/radii) * DT
    
    # Update radii and ensure they don't go below zero
    radii += dr
    radii[radii < 0] = 0
    
    # Maintain (approximate) Volume Fraction conservation 
    # by normalizing if needed, though the law is naturally conservative
    
    # Update the visual circles
    for i, circle in enumerate(circles):
        circle.set_radius(radii[i])
        # If particle dissolved, make it invisible
        if radii[i] <= 0.1:
            circle.set_visible(False)
            
    return circles

ani = FuncAnimation(fig, update, frames=200, interval=50, blit=True)

# To save the movie, uncomment the line below:
# ani.save('ostwald_ripening.mp4', writer='ffmpeg')

plt.show()