Sandbox.py

import pygame
import random
import math
import sys

# --- Constants & Configuration ---
SIM_WIDTH, SIM_HEIGHT = 800, 600
UI_WIDTH = 250
SCREEN_WIDTH = SIM_WIDTH + UI_WIDTH
SCREEN_HEIGHT = SIM_HEIGHT
CELL_SIZE = 4
FPS = 60

GRID_W = SIM_WIDTH // CELL_SIZE
GRID_H = SIM_HEIGHT // CELL_SIZE

# Spatial partitioning for performance
CHUNK_SIZE = 8
W_CHUNKS = math.ceil(GRID_W / CHUNK_SIZE)
H_CHUNKS = math.ceil(GRID_H / CHUNK_SIZE)

# --- Material Definitions ---
# id, name, grav, dens, disp, is_liq, flam, hex, variance, is_conductor
MAT_DEFS = [
    (0,  "Empty",   0,  0.0,  0, False, False, 0x0C0F12, 0,  False),
    (1,  "Sand",    1,  5.0,  0, False, False, 0xE6C27A, 15, False),
    (2,  "Water",   1,  1.0,  5, True,  False, 0x2980B9, 5,  True),
    (3,  "Stone",   0,  10.0, 0, False, False, 0x7F8C8D, 8,  False),
    (4,  "Wood",    0,  10.0, 0, False, True,  0x5C4033, 10, False),
    (5,  "Fire",   -1,  0.3,  0, False, False, 0xE74C3C, 0,  False),   
    (6,  "Smoke",  -1,  0.1,  0, False, False, 0x34495E, 5,  False),
    (7,  "Steam",  -1,  0.2,  0, False, False, 0xBDC3C7, 10, False),
    (8,  "Lava",    1,  2.0,  1, True,  False, 0xD35400, 5,  False),
    (9,  "Glass",   0,  10.0, 0, False, False, 0x85C1E9, 2,  False),
    (10, "Oil",     1,  0.8,  4, True,  True,  0x2C3E50, 4,  False),
    (11, "Acid",    1,  1.1,  2, True,  False, 0x2ECC71, 5,  True),
    (12, "Ice",     0,  10.0, 0, False, False, 0xA9CCE3, 2,  False),
    (13, "Snow",    1,  4.0,  0, False, False, 0xECF0F1, 5,  False),
    (14, "C4",      0,  10.0, 0, False, True,  0x8B0000, 2,  False),   
    (15, "Methane",-1,  0.2,  0, False, True,  0x48C9B0, 5,  False),   
    (16, "Plant",   0,  10.0, 0, False, True,  0x27AE60, 10, False),  
    (17, "Plasma", -1,  0.4,  0, False, False, 0x9B59B6, 0,  True),   
    (18, "Cloner",  0,  10.0, 0, False, False, 0xF39C12, 0,  False),   
    (19, "Void",    0,  10.0, 0, False, False, 0x17202A, 0,  False),
    (20, "Metal",   0,  10.0, 0, False, False, 0x95A5A6, 5,  True),
    (21, "Spark",   0,  0.0,  0, False, False, 0xFFFF00, 0,  False), 
    (22, "Battery", 0,  10.0, 0, False, False, 0x22A7F0, 0,  True),
    (23, "Mud",     1,  3.0,  1, True,  False, 0x5C4A3D, 5,  False),
    (24, "Ash",     1,  1.0,  0, False, False, 0xB2BABB, 5,  False),
    (25, "Virus",   1,  4.0,  0, False, False, 0x8E44AD, 10, False),
    (26, "Poison", -1,  0.3,  0, False, False, 0x8A2BE2, 5,  False),
    (27, "Gunpwdr", 1,  4.5,  0, False, True,  0x4B5320, 5,  False),
    (28, "Dirt",    1,  4.8,  0, False, False, 0x6E2C00, 10, False)
]

# Material Enums
(EMPTY, SAND, WATER, STONE, WOOD, FIRE, SMOKE, STEAM, LAVA, GLASS, 
 OIL, ACID, ICE, SNOW, C4, METHANE, PLANT, PLASMA, CLONER, VOID, 
 METAL, SPARK, BATTERY, MUD, ASH, VIRUS, POISON, GUNPOWDER, DIRT) = range(29)

# Fast-access lists
NAMES = [m[1] for m in MAT_DEFS]
GRAVITY = [m[2] for m in MAT_DEFS]
DENSITY = [m[3] for m in MAT_DEFS]
DISPERSION = [m[4] for m in MAT_DEFS]
IS_LIQUID = [m[5] for m in MAT_DEFS]
FLAMMABLE = [m[6] for m in MAT_DEFS]
BASE_COLORS = [m[7] for m in MAT_DEFS]
VARIANCES = [m[8] for m in MAT_DEFS]
CONDUCTOR = [m[9] for m in MAT_DEFS]

def get_varied_color(type_id):
    base_color = BASE_COLORS[type_id]
    variance = VARIANCES[type_id]
    if variance == 0: return base_color
    r, g, b = (base_color >> 16) & 0xFF, (base_color >> 8) & 0xFF, base_color & 0xFF
    noise = random.randint(-variance, variance)
    return (max(0, min(255, r + noise)) << 16) | (max(0, min(255, g + noise)) << 8) | max(0, min(255, b + noise))

class ParticleEngine:
    def __init__(self):
        self.g_type = bytearray(GRID_W * GRID_H)
        self.g_base = bytearray(GRID_W * GRID_H) # Store states (like conductive material under spark)
        self.g_color = [BASE_COLORS[EMPTY]] * (GRID_W * GRID_H)
        self.g_life = bytearray(GRID_W * GRID_H)
        self.g_updated = bytearray(GRID_W * GRID_H)
        
        self.chunks_active = [False] * (W_CHUNKS * H_CHUNKS)
        self.chunks_next = [False] * (W_CHUNKS * H_CHUNKS)
        self.frame_count = 1
        self.clear_grid()

    def clear_grid(self):
        self.g_type[:] = bytearray(GRID_W * GRID_H)
        self.g_base[:] = bytearray(GRID_W * GRID_H)
        self.g_color[:] = [BASE_COLORS[EMPTY]] * (GRID_W * GRID_H)
        self.g_life[:] = bytearray(GRID_W * GRID_H)
        self.g_updated[:] = bytearray(GRID_W * GRID_H)
        self.chunks_active = [False] * (W_CHUNKS * H_CHUNKS)
        self.chunks_next = [False] * (W_CHUNKS * H_CHUNKS)
        # Create a basic floor of metal
        for x in range(GRID_W):
            self.set_particle(x, GRID_H - 1, METAL, force=True)
            self.set_particle(x, GRID_H - 2, METAL, force=True)

    def wake_chunk(self, x, y):
        cx, cy = x // CHUNK_SIZE, y // CHUNK_SIZE
        if 0 <= cx < W_CHUNKS and 0 <= cy < H_CHUNKS:
            self.chunks_next[cy * W_CHUNKS + cx] = True
            if x % CHUNK_SIZE == 0 and cx > 0: self.chunks_next[cy * W_CHUNKS + cx - 1] = True
            elif x % CHUNK_SIZE == CHUNK_SIZE - 1 and cx < W_CHUNKS - 1: self.chunks_next[cy * W_CHUNKS + cx + 1] = True
            if y % CHUNK_SIZE == 0 and cy > 0: self.chunks_next[(cy - 1) * W_CHUNKS + cx] = True
            elif y % CHUNK_SIZE == CHUNK_SIZE - 1 and cy < H_CHUNKS - 1: self.chunks_next[(cy + 1) * W_CHUNKS + cx] = True

    def set_particle(self, x, y, type_id, force=False):
        if not (0 <= x < GRID_W and 0 <= y < GRID_H): return
        idx = y * GRID_W + x
        curr = self.g_type[idx]
        if curr == type_id: return
        
        if not force and type_id != EMPTY:
            if curr != EMPTY:
                if DENSITY[type_id] > DENSITY[curr] and GRAVITY[curr] < 0: pass 
                else: return
        
        self.g_type[idx] = type_id
        self.g_color[idx] = get_varied_color(type_id)
        self.g_base[idx] = 0
        
        if type_id == FIRE: self.g_life[idx] = random.randint(20, 60)
        elif type_id == SMOKE: self.g_life[idx] = random.randint(50, 100)
        elif type_id == STEAM: self.g_life[idx] = random.randint(60, 120)
        elif type_id == PLASMA: self.g_life[idx] = random.randint(30, 60)
        elif type_id == VIRUS: self.g_life[idx] = random.randint(150, 255)
        elif type_id == SPARK: self.g_life[idx] = 2
        else: self.g_life[idx] = 0
        
        self.wake_chunk(x, y)

    def swap_particles(self, x1, y1, x2, y2):
        idx1, idx2 = y1 * GRID_W + x1, y2 * GRID_W + x2
        self.g_type[idx1], self.g_type[idx2] = self.g_type[idx2], self.g_type[idx1]
        self.g_color[idx1], self.g_color[idx2] = self.g_color[idx2], self.g_color[idx1]
        self.g_life[idx1], self.g_life[idx2] = self.g_life[idx2], self.g_life[idx1]
        self.g_base[idx1], self.g_base[idx2] = self.g_base[idx2], self.g_base[idx1]
        self.g_updated[idx1] = self.g_updated[idx2] = self.frame_count
        self.wake_chunk(x1, y1)
        self.wake_chunk(x2, y2)

    def explode(self, cx, cy, radius):
        for dy in range(-radius, radius + 1):
            for dx in range(-radius, radius + 1):
                if dx*dx + dy*dy <= radius*radius:
                    ex, ey = cx + dx, cy + dy
                    if 0 <= ex < GRID_W and 0 <= ey < GRID_H:
                        idx = ey * GRID_W + ex
                        curr = self.g_type[idx]
                        if curr in (VOID, CLONER): continue
                        
                        dist = math.sqrt(dx*dx + dy*dy)
                        # Inner Core: Total destruction
                        if dist < radius * 0.5:
                            if curr in (C4, METHANE, OIL, GUNPOWDER) and dist > 1:
                                self.set_particle(ex, ey, FIRE, force=True)
                            else:
                                if random.random() < 0.2: self.set_particle(ex, ey, FIRE, force=True)
                                elif random.random() < 0.2: self.set_particle(ex, ey, PLASMA, force=True) 
                                else: self.set_particle(ex, ey, EMPTY, force=True)
                        # Mid Area
                        elif dist < radius * 0.8:
                            if curr in (C4, METHANE, OIL, GUNPOWDER):
                                self.set_particle(ex, ey, FIRE, force=True)
                            elif curr not in (STONE, METAL): 
                                if random.random() < 0.3: self.set_particle(ex, ey, FIRE, force=True)
                                else: self.set_particle(ex, ey, SMOKE, force=True)
                        # Outer Edge
                        else:
                            if curr == EMPTY or FLAMMABLE[curr]:
                                if random.random() < 0.4: self.set_particle(ex, ey, SMOKE, force=True)
                                elif random.random() < 0.1: self.set_particle(ex, ey, FIRE, force=True)

    def update(self):
        self.frame_count = (self.frame_count + 1) % 255
        if self.frame_count == 0: self.frame_count = 1
        step_x_dir = 1 if self.frame_count % 2 == 0 else -1
        
        g_type, g_updated = self.g_type, self.g_updated
        
        for cy in range(H_CHUNKS - 1, -1, -1):
            cx_range = range(W_CHUNKS) if step_x_dir == 1 else range(W_CHUNKS - 1, -1, -1)
            for cx in cx_range:
                if not self.chunks_active[cy * W_CHUNKS + cx]: continue
                
                for dy in range(CHUNK_SIZE - 1, -1, -1):
                    y = cy * CHUNK_SIZE + dy
                    if y >= GRID_H: continue
                    
                    dx_range = range(CHUNK_SIZE) if step_x_dir == 1 else range(CHUNK_SIZE - 1, -1, -1)
                    for dx in dx_range:
                        x = cx * CHUNK_SIZE + dx
                        if x >= GRID_W: continue
                        
                        idx = y * GRID_W + x
                        if g_updated[idx] == self.frame_count: continue
                        
                        t = g_type[idx]
                        if t == EMPTY: continue
                        self.update_particle(x, y, idx, t, step_x_dir)
        
        self.chunks_active = self.chunks_next[:]
        self.chunks_next = [False] * (W_CHUNKS * H_CHUNKS)

    def update_particle(self, x, y, idx, t, step_x_dir):
        # 1. Lifespan and Cooldowns
        if self.g_life[idx] > 0:
            if t == SPARK:
                self.g_life[idx] -= 1
                if self.g_life[idx] == 0:
                    base_mat = self.g_base[idx]
                    self.g_type[idx] = base_mat
                    self.g_color[idx] = get_varied_color(base_mat)
                    self.g_life[idx] = 10 # Post-spark electrical Cooldown
                    self.wake_chunk(x, y)
                    return
            elif t in (FIRE, STEAM, SMOKE, PLASMA, VIRUS):
                self.g_life[idx] -= 1
                if self.g_life[idx] == 0:
                    if t == FIRE:
                        if random.random() < 0.2: self.set_particle(x, y, SMOKE, force=True)
                        else: self.set_particle(x, y, EMPTY, force=True)
                    elif t == STEAM:
                        if random.random() < 0.1: self.set_particle(x, y, WATER, force=True)
                        else: self.set_particle(x, y, EMPTY, force=True)
                    elif t == VIRUS:
                        self.set_particle(x, y, ASH, force=True)
                    else: 
                        self.set_particle(x, y, EMPTY, force=True)
                    return
            else:
                self.g_life[idx] -= 1

        # 2. Electricity Generation & Arcing
        if t == BATTERY:
            if random.random() < 0.1:
                for dx, dy in ((-1,0), (1,0), (0,-1), (0,1)):
                    nx, ny = x + dx, y + dy
                    if 0 <= nx < GRID_W and 0 <= ny < GRID_H:
                        n_idx = ny * GRID_W + nx
                        nt = self.g_type[n_idx]
                        if CONDUCTOR[nt] and nt != SPARK and self.g_life[n_idx] == 0:
                            self.g_base[n_idx] = nt
                            self.g_type[n_idx] = SPARK
                            self.g_color[n_idx] = get_varied_color(SPARK)
                            self.g_life[n_idx] = 2
                            self.wake_chunk(nx, ny)

        if t == SPARK:
            for dx, dy in ((-1,0), (1,0), (0,-1), (0,1)):
                nx, ny = x + dx, y + dy
                if 0 <= nx < GRID_W and 0 <= ny < GRID_H:
                    n_idx = ny * GRID_W + nx
                    nt = self.g_type[n_idx]
                    if CONDUCTOR[nt] and nt != SPARK and self.g_life[n_idx] == 0:
                        self.g_base[n_idx] = nt
                        self.g_type[n_idx] = SPARK
                        self.g_color[n_idx] = get_varied_color(SPARK)
                        self.g_life[n_idx] = 2
                        self.wake_chunk(nx, ny)
                    elif nt in (C4, METHANE, OIL, GUNPOWDER):
                        self.explode(nx, ny, 8 if nt != GUNPOWDER else 5)
            return # Spark electricity halts movement for 2 frames
            
        # 3. Gravity with natural slope inertial sliding
        grav = GRAVITY[t]
        if t == SNOW and random.random() < 0.5: grav = 0 
        
        moved = False
        if grav != 0:
            ny = y + grav
            if 0 <= ny < GRID_H:
                nt = self.g_type[ny * GRID_W + x]
                if DENSITY[nt] < DENSITY[t]:
                    self.swap_particles(x, y, x, ny)
                    moved = True
                    y = ny
                else:
                    dl, dr = x - step_x_dir, x + step_x_dir
                    
                    can_go_l = (0 <= dl < GRID_W and DENSITY[self.g_type[y * GRID_W + dl]] < DENSITY[t]) and \
                               (0 <= dl < GRID_W and DENSITY[self.g_type[ny * GRID_W + dl]] < DENSITY[t])
                    can_go_r = (0 <= dr < GRID_W and DENSITY[self.g_type[y * GRID_W + dr]] < DENSITY[t]) and \
                               (0 <= dr < GRID_W and DENSITY[self.g_type[ny * GRID_W + dr]] < DENSITY[t])

                    if can_go_l and can_go_r:
                        nx = dl if random.random() < 0.5 else dr
                        self.swap_particles(x, y, nx, ny)
                        moved = True
                        x, y = nx, ny
                    elif can_go_l:
                        self.swap_particles(x, y, dl, ny)
                        moved = True
                        x, y = dl, ny
                    elif can_go_r:
                        self.swap_particles(x, y, dr, ny)
                        moved = True
                        x, y = dr, ny
            else:
                self.set_particle(x, y, EMPTY, force=True)
                return

        # 4. Liquid Dispersion
        if not moved and IS_LIQUID[t]:
            disp = DISPERSION[t]
            target_x = x
            for check_dir in (step_x_dir, -step_x_dir):
                for dx in range(1, disp + 1):
                    nx = x + dx * check_dir
                    if 0 <= nx < GRID_W:
                        if DENSITY[self.g_type[y * GRID_W + nx]] < DENSITY[t]: target_x = nx
                        else: break
                    else: break
                if target_x != x:
                    self.swap_particles(x, y, target_x, y)
                    moved = True
                    x = target_x
                    break

        # 5. Gas Lateral Drift
        if not moved and grav < 0:
            if random.random() < 0.6:
                nx = x + random.choice([-1, 1])
                if 0 <= nx < GRID_W and DENSITY[self.g_type[y * GRID_W + nx]] < DENSITY[t]:
                    self.swap_particles(x, y, nx, y)
                    x = nx
            if random.random() < 0.2:
                ny = y + random.choice([-1, 1])
                if 0 <= ny < GRID_H and DENSITY[self.g_type[ny * GRID_W + x]] < DENSITY[t]:
                    self.swap_particles(x, y, x, ny)
                    y = ny
                        
        # 6. Biologicals & Reactions
        if t in (FIRE, LAVA, ACID, WATER, ICE, PLASMA, CLONER, VOID, VIRUS, POISON, PLANT):
            self.process_reactions(x, y, t)

    def process_reactions(self, x, y, t):
        dirs = [(-1,0), (1,0), (0,-1), (0,1)]
        random.shuffle(dirs)
        nx, ny = x + dirs[0][0], y + dirs[0][1]
        
        if not (0 <= nx < GRID_W and 0 <= ny < GRID_H): return
        
        idx = y * GRID_W + x
        n_idx = ny * GRID_W + nx
        nt = self.g_type[n_idx]
        if nt == EMPTY and t != CLONER: return
        
        if t == FIRE:
            if nt == C4: self.explode(nx, ny, 14)
            elif nt == METHANE: self.explode(nx, ny, 6)
            elif nt == GUNPOWDER: self.explode(nx, ny, 5)
            elif nt == OIL: self.explode(nx, ny, 4)
            elif nt == WOOD:
                if random.random() < 0.05: self.set_particle(nx, ny, FIRE, force=True)
            elif FLAMMABLE[nt]:
                if random.random() < 0.15: self.set_particle(nx, ny, FIRE, force=True)
            elif nt in (ICE, SNOW):
                self.set_particle(nx, ny, WATER, force=True)
            elif nt == WATER:
                if random.random() < 0.3:
                    self.set_particle(x, y, EMPTY, force=True)
                    self.set_particle(nx, ny, STEAM, force=True)
            elif nt == SAND:
                if random.random() < 0.01: self.set_particle(nx, ny, GLASS, force=True)
            elif nt == MUD:
                if random.random() < 0.2: # Bake Mud into Dirt
                    self.set_particle(nx, ny, DIRT, force=True)
                    self.set_particle(x, y, EMPTY, force=True)
                    
        elif t == LAVA:
            if nt == WATER:
                self.set_particle(x, y, STONE, force=True)
                self.set_particle(nx, ny, STEAM, force=True)
            elif nt in (C4, METHANE, GUNPOWDER, OIL): self.explode(nx, ny, 10)
            elif nt == SAND:
                if random.random() < 0.1: self.set_particle(nx, ny, GLASS, force=True)
            elif FLAMMABLE[nt]: self.set_particle(nx, ny, FIRE, force=True)
            elif nt in (ICE, SNOW):
                self.set_particle(nx, ny, WATER, force=True)
                self.set_particle(x, y, STONE, force=True)
                
        elif t == ACID:
            if nt not in (EMPTY, ACID, GLASS, VOID, CLONER, SPARK, BATTERY):
                if nt == STONE and random.random() < 0.3:
                    self.set_particle(nx, ny, POISON, force=True)
                else:
                    if random.random() < 0.2: self.set_particle(nx, ny, SMOKE, force=True)
                if random.random() < 0.15: self.set_particle(x, y, EMPTY, force=True) 
        
        elif t == WATER:
            if nt == LAVA:
                self.set_particle(x, y, STEAM, force=True)
                self.set_particle(nx, ny, STONE, force=True)
            elif nt == DIRT or nt == SAND:
                if random.random() < 0.1: 
                    self.set_particle(nx, ny, MUD, force=True)
                    self.set_particle(x, y, EMPTY, force=True)
            elif nt == POISON:
                if random.random() < 0.1: self.set_particle(x, y, ACID, force=True)
                
        elif t == ICE:
            if nt == WATER:
                if random.random() < 0.05: self.set_particle(nx, ny, ICE, force=True) 
            elif nt in (FIRE, LAVA, PLASMA):
                self.set_particle(x, y, WATER, force=True)
                
        elif t == PLASMA:
            if nt not in (VOID, CLONER, EMPTY, PLASMA):
                self.explode(nx, ny, 3) 
                if random.random() < 0.1: self.set_particle(x, y, EMPTY, force=True)

        elif t == VIRUS:
            if nt not in (VOID, CLONER, EMPTY, VIRUS, FIRE, PLASMA, ASH, GLASS, SPARK, BATTERY):
                if random.random() < 0.1: self.set_particle(nx, ny, VIRUS, force=True)

        elif t == POISON:
            if nt == PLANT:
                if random.random() < 0.2: self.set_particle(nx, ny, ASH, force=True)
            elif nt == WATER:
                if random.random() < 0.1:
                    self.set_particle(nx, ny, ACID, force=True)
                    self.set_particle(x, y, EMPTY, force=True)

        elif t == PLANT:
            if nt == WATER:
                if random.random() < 0.3: self.set_particle(nx, ny, PLANT, force=True)
            elif nt == DIRT:
                if random.random() < 0.05: self.set_particle(nx, ny, PLANT, force=True)
            elif nt == WOOD:
                if random.random() < 0.01: self.set_particle(nx, ny, PLANT, force=True) # Slow Overgrowth
                
        elif t == CLONER:
            if self.g_base[idx] == 0:
                if nt not in (CLONER, EMPTY, SPARK):
                    self.g_base[idx] = nt
                    self.g_color[idx] = self.g_color[n_idx] 
                    self.wake_chunk(x, y)
            else:
                clone_mat = self.g_base[idx]
                for dx, dy in ((-1,0), (1,0), (0,-1), (0,1)):
                    cx, cy = x + dx, y + dy
                    if 0 <= cx < GRID_W and 0 <= cy < GRID_H:
                        if self.g_type[cy * GRID_W + cx] == EMPTY:
                            if random.random() < 0.1:
                                self.set_particle(cx, cy, clone_mat, force=True)

        elif t == VOID:
            if nt != VOID: self.set_particle(nx, ny, EMPTY, force=True)

# --- Drawing and UI ---
def draw_ui(screen, font, current_mat, brush_size, fps, engine, paused):
    pygame.draw.rect(screen, (15, 18, 22), (SIM_WIDTH, 0, UI_WIDTH, SCREEN_HEIGHT))
    
    title = font.render("PARTICLE SANDBOX", True, (240, 240, 255))
    screen.blit(title, (SIM_WIDTH + 20, 15))
    
    fps_text = font.render(f"FPS: {int(fps)}", True, (150, 180, 200))
    screen.blit(fps_text, (SIM_WIDTH + 20, 40))
    
    brush_text = font.render(f"Brush: {brush_size} px   ( [ / ] )", True, (150, 180, 200))
    screen.blit(brush_text, (SIM_WIDTH + 20, 65))
    
    status_text = font.render(f"Paused" if paused else "Running", True, (255, 100, 100) if paused else (100, 255, 100))
    screen.blit(status_text, (SIM_WIDTH + 140, 40))
    
    info = ["L-Click: Draw", "R-Click: Erase", "SPACE: Play/Pause", "C: Clear Grid"]
    for i, inf in enumerate(info):
        t = font.render(inf, True, (100, 110, 120))
        screen.blit(t, (SIM_WIDTH + 20, 95 + i * 18))

    # Materials Palette Matrix
    start_y = 180
    btn_w, btn_h = 66, 28
    gap_x, gap_y = 8, 8
    
    valid_mats = [m for m in MAT_DEFS if m[0] != EMPTY]
    for i, mat in enumerate(valid_mats):
        col, row = i % 3, i // 3
        bx = SIM_WIDTH + 15 + col * (btn_w + gap_x)
        by = start_y + row * (btn_h + gap_y)
        
        mat_rect = pygame.Rect(bx, by, btn_w, btn_h)
        bg_col = (50, 60, 70) if current_mat == mat[0] else (25, 30, 35)
        pygame.draw.rect(screen, bg_col, mat_rect, border_radius=6)
        if current_mat == mat[0]:
            pygame.draw.rect(screen, (255, 215, 0), mat_rect, 2, border_radius=6)
            
        color = (mat[7] >> 16, (mat[7] >> 8) & 0xFF, mat[7] & 0xFF)
        pygame.draw.circle(screen, color, (bx + 12, by + btn_h//2), 5)
        
        name_text = font.render(mat[1], True, (200, 210, 220))
        screen.blit(name_text, (bx + 22, by + int(btn_h/2 - name_text.get_height()/2)))

# --- Main Application ---
def main():
    pygame.init()
    screen = pygame.display.set_mode((SCREEN_WIDTH, SCREEN_HEIGHT))
    pygame.display.set_caption("Extended Particle Sandbox")
    clock = pygame.time.Clock()
    
    try: font = pygame.font.SysFont("Segoe UI, Arial", 13, bold=True)
    except: font = pygame.font.Font(None, 18)
    
    engine = ParticleEngine()
    sim_surf = pygame.Surface((GRID_W, GRID_H))
    
    current_mat = SAND
    brush_size = 4
    paused = False
    last_mx, last_my = 0, 0

    def apply_brush(gx, gy, radius, mat, force):
        for dy in range(-radius, radius + 1):
            for dx in range(-radius, radius + 1):
                if dx*dx + dy*dy <= radius*radius:
                    engine.set_particle(gx + dx, gy + dy, mat, force)

    def draw_line(x0, y0, x1, y1, radius, mat, force):
        dx, dy = abs(x1 - x0), -abs(y1 - y0)
        sx = 1 if x0 < x1 else -1
        sy = 1 if y0 < y1 else -1
        err = dx + dy
        while True:
            apply_brush(x0, y0, radius, mat, force)
            if x0 == x1 and y0 == y1: break
            e2 = 2 * err
            if e2 >= dy: err += dy; x0 += sx
            if e2 <= dx: err += dx; y0 += sy

    running = True
    while running:
        mouse_x, mouse_y = pygame.mouse.get_pos()
        grid_mx, grid_my = mouse_x // CELL_SIZE, mouse_y // CELL_SIZE
        mouse_pressed = pygame.mouse.get_pressed()
        
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                running = False
            elif event.type == pygame.MOUSEBUTTONDOWN:
                if mouse_x >= SIM_WIDTH and event.button == 1:
                    start_y, btn_w, btn_h, gap_x, gap_y = 180, 66, 28, 8, 8
                    valid_mats = [m for m in MAT_DEFS if m[0] != EMPTY]
                    for i, mat in enumerate(valid_mats):
                        col, row = i % 3, i // 3
                        bx = SIM_WIDTH + 15 + col * (btn_w + gap_x)
                        by = start_y + row * (btn_h + gap_y)
                        if pygame.Rect(bx, by, btn_w, btn_h).collidepoint(mouse_x, mouse_y):
                            current_mat = mat[0]
                            break
                else:
                    last_mx, last_my = grid_mx, grid_my
                    
            elif event.type == pygame.KEYDOWN:
                if event.key == pygame.K_SPACE or event.key == pygame.K_p: paused = not paused
                elif event.key == pygame.K_c: engine.clear_grid()
                elif event.key == pygame.K_LEFTBRACKET: brush_size = max(0, brush_size - 1)
                elif event.key == pygame.K_RIGHTBRACKET: brush_size = min(20, brush_size + 1)
        
        if mouse_x < SIM_WIDTH:
            if mouse_pressed[0]: draw_line(last_mx, last_my, grid_mx, grid_my, brush_size, current_mat, force=False)
            elif mouse_pressed[2]: draw_line(last_mx, last_my, grid_mx, grid_my, brush_size, EMPTY, force=True)
        last_mx, last_my = grid_mx, grid_my
        
        if not paused: engine.update()
            
        px = pygame.PixelArray(sim_surf)
        for cy in range(H_CHUNKS):
            for cx in range(W_CHUNKS):
                if engine.chunks_active[cy * W_CHUNKS + cx]:
                    for dy in range(CHUNK_SIZE):
                        y = cy * CHUNK_SIZE + dy
                        if y >= GRID_H: continue
                        idx = y * GRID_W + cx * CHUNK_SIZE
                        for dx in range(CHUNK_SIZE):
                            x = cx * CHUNK_SIZE + dx
                            if x >= GRID_W: break
                            px[x, y] = engine.g_color[idx]
                            idx += 1
        px.close()

        screen.blit(pygame.transform.scale(sim_surf, (SIM_WIDTH, SCREEN_HEIGHT)), (0, 0))
        draw_ui(screen, font, current_mat, brush_size, clock.get_fps(), engine, paused)
        
        if mouse_x < SIM_WIDTH:
            pygame.draw.circle(screen, (200, 200, 200), (mouse_x, mouse_y), (brush_size * CELL_SIZE) + max(1, CELL_SIZE//2), 1)

        pygame.display.flip()
        clock.tick(FPS)
        
    pygame.quit()
    sys.exit()

if __name__ == "__main__":
    main()