import numpy as np
import matplotlib.pyplot as plt
from numba import njit, prange
from matplotlib.cm import ScalarMappable
from matplotlib.colors import Normalize

# Parameters
circle_radius = 100  # Maximum boundary of the circle
num_particles = 1000000  # Total number of particles
step_size = 1 # Size of each random walk step
max_attempts = 100000  # Maximum attempts per particle before discarding

# Initialize the seed at the center
seed = (0, 0)
aggregated_particles = [seed]

# Initialize the figure for visualization
plt.ion()
fig, ax = plt.subplots()
ax.set_xlim(-circle_radius, circle_radius)
ax.set_ylim(-circle_radius, circle_radius)
ax.set_aspect('equal')

# Set a dark background color
ax.set_facecolor('black')  # Dark background
fig.patch.set_facecolor('black')  # Dark background for the figure

# Create a colormap (e.g., 'hot' for heatmap effect)
cmap = plt.get_cmap('hot')
norm = Normalize(vmin=0, vmax=circle_radius)  # Normalize distances to [0, circle_radius]
scatter = ax.scatter(*zip(*aggregated_particles), c=[0], cmap=cmap, norm=norm, s=1)

# Add a colorbar for reference
cbar = plt.colorbar(ScalarMappable(norm=norm, cmap=cmap), ax=ax)
cbar.set_label('Distance from Center', color='white')  # White text for visibility
cbar.ax.yaxis.set_tick_params(color='white')  # White ticks for visibility
plt.setp(plt.getp(cbar.ax.axes, 'yticklabels'), color='white')  # White tick labels

# JIT-compiled functions for performance
@njit
def is_inside_circle(x, y, radius):
    """Check if a point is inside the circle."""
    return x**2 + y**2 < radius**2

@njit
def random_walk(x, y, step_size):
    """Perform a random walk in one of four cardinal directions."""
    direction_idx = np.random.randint(0, 4)  # Random index for direction
    directions = np.array([(1, 0), (-1, 0), (0, 1), (0, -1)], dtype=np.float64)
    direction = directions[direction_idx]
    return x + direction[0] * step_size, y + direction[1] * step_size

@njit
def is_near_aggregated(x, y, aggregated_particles_np, step_size):
    """Check if the particle is near any aggregated particle."""
    threshold = step_size**2 * 1.3  # Slightly larger threshold to ensure connection
    for particle in aggregated_particles_np:
        if (x - particle[0])**2 + (y - particle[1])**2 <= threshold:
            return True
    return False

# Function to compute the dynamic spawn radius
@njit
def compute_spawn_radius(aggregated_particles_np, margin=10):
    """Determine the spawn radius based on the furthest aggregated particle."""
    max_dist = 0.0
    for particle in aggregated_particles_np:
        dist = particle[0]**2 + particle[1]**2
        if dist > max_dist:
            max_dist = dist
    return np.sqrt(max_dist) + margin

# Convert aggregated_particles to a NumPy array for compatibility with Numba
aggregated_particles_np = np.array(aggregated_particles, dtype=np.float64)

# Main simulation loop
for _ in range(num_particles):
    # Dynamically compute the spawn radius
    spawn_radius = compute_spawn_radius(aggregated_particles_np)
    angle = np.random.uniform(0, 2 * np.pi)
    x, y = spawn_radius * np.cos(angle), spawn_radius * np.sin(angle)
 
    attempts = 0
    while attempts < max_attempts:
        # If the particle exits the boundary, discard it and spawn a new one
        if not is_inside_circle(x, y, circle_radius):
            break
 
        # Check if the particle is near the aggregation cluster
        if is_near_aggregated(x, y, aggregated_particles_np, step_size):
            # Add the particle to the aggregated list
            aggregated_particles_np = np.vstack((aggregated_particles_np, np.array([x, y])))
 
            # Update the visualization
            distances = np.sqrt(aggregated_particles_np[:, 0]**2 + aggregated_particles_np[:, 1]**2)
            scatter.set_offsets(aggregated_particles_np)
            scatter.set_array(distances)  # Update colors based on distance
            plt.draw()
            plt.pause(0.001)
            break
 
        # Perform the random walk
        x, y = random_walk(x, y, step_size)
        attempts += 1

# Finalize the visualization
plt.ioff()
plt.show()

import numpy as np
import matplotlib.pyplot as plt
from numba import njit
from matplotlib.animation import FuncAnimation
import time
from tqdm import tqdm

# Parameters
radius = 150
grid_size = 2 * radius + 1
num_particles = 15000
max_attempts = 150000
spawn_radius_margin = 20

# Grid setup
dla_grid = np.full((grid_size, grid_size), np.nan, dtype=np.float32)
center = grid_size // 2
dla_grid[center, center] = 0  # Seed in the center

# Figure setup
fig, ax = plt.subplots()
ax.set_xticks([])
ax.set_yticks([])
ax.set_facecolor("black")
fig.patch.set_facecolor("black")

# Colormap
plasma_cmap = plt.get_cmap("plasma")
plasma_cmap.set_bad(color="black")  # NaN areas are black
im = ax.imshow(dla_grid, cmap=plasma_cmap, origin="lower", vmin=0, vmax=grid_size // 2)
cbar = plt.colorbar(im, ax=ax)
cbar.set_label("Euclidean Distance from Center", color="white")
cbar.ax.yaxis.set_tick_params(color="white")
plt.setp(plt.getp(cbar.ax.axes, "yticklabels"), color="white")

@njit(fastmath=True)
def clamp(value, min_value, max_value):
    return max(min_value, min(value, max_value))

@njit(fastmath=True)
def is_inside_circle(x, y, radius):
    return x * x + y * y < radius * radius

@njit(fastmath=True)
def random_walk(x, y, grid_size):
    # Precompute movement directions
    direction_table = np.array([
        (1, 0), (-1, 0), (0, 1), (0, -1),
        (1, 1), (-1, 1), (1, -1), (-1, -1)
    ], dtype=np.int8)
    dx, dy = direction_table[np.random.randint(8)]
    return clamp(x + dx, 0, grid_size - 1), clamp(y + dy, 0, grid_size - 1)

@njit
def is_near_aggregated(x, y, grid):
    """Vectorized check for neighboring aggregated pixels."""
    neighbors = grid[max(0, x - 1):min(grid.shape[0], x + 2),
                     max(0, y - 1):min(grid.shape[1], y + 2)]
    return np.any(neighbors >= 0)

@njit
def compute_spawn_radius(grid, center, margin):
    nonzero_x, nonzero_y = np.where(~np.isnan(grid))
    if len(nonzero_x) == 0:
        return margin
    max_dist_sq = np.max((nonzero_x - center) ** 2 + (nonzero_y - center) ** 2)
    return np.sqrt(max_dist_sq) + margin

@njit(fastmath=True)
def compute_euclidean_distance_sq(x, y, center):
    """Return squared Euclidean distance to avoid unnecessary sqrt calls."""
    return (x - center) ** 2 + (y - center) ** 2
    
@njit(fastmath=True)
def spawn_particle(center, spawn_radius, grid_size):
    """Spawn a particle near the computed radius."""
    angle = np.random.uniform(0, 2 * np.pi)
    x = int(center + spawn_radius * np.cos(angle))
    y = int(center + spawn_radius * np.sin(angle))
    return clamp(x, 0, grid_size - 1), clamp(y, 0, grid_size - 1)

@njit(fastmath=True)
def perform_random_walk(x, y, center, grid, max_attempts):
    """Perform a random walk until aggregation or max attempts reached."""
    attempts = 0
    while attempts < max_attempts:
        if not is_inside_circle(x - center, y - center, grid.shape[0] // 2):
            return None
        if is_near_aggregated(x, y, grid):
            return x, y
        x, y = random_walk(x, y, grid.shape[0])
        attempts += 1
    return None

# Initialize progress bar
pbar = tqdm(total=num_particles)

# Animation update function
def update(frame):
    # Compute dynamic spawn radius
    spawn_radius = compute_spawn_radius(dla_grid, center, spawn_radius_margin)
    spawn_radius = min(spawn_radius, grid_size // 2 - 1)

    # Spawn and simulate a particle
    x, y = spawn_particle(center, spawn_radius, grid_size)
    result = perform_random_walk(x, y, center, dla_grid, max_attempts)

    if result is not None:
        x, y = result
        color = np.sqrt(compute_euclidean_distance_sq(x, y, center))
        dla_grid[x, y] = color  # Update the data matrix

        # Directly update the pixel buffer instead of redrawing the whole image
        im_array = im.get_array()
        im_array[x, y] = color
        im.set_array(im_array)  # Apply the change

    # Update progress bar
    pbar.update(1)

    # Stop animation when complete
    if frame == num_particles - 1:
        pbar.close()
        print(f"Simulation complete! Time taken: {time.time() - start_time:.2f} seconds.")

start_time = time.time()

# Run animation
ani = FuncAnimation(fig, update, frames=num_particles, interval=1, repeat=False)
plt.show()