import pygame
import random

# Define particle IDs
EMPTY = 0
SAND = 1
WATER = 2
OIL = 3
FIRE = 4
SMOKE = 5
PLANT = 6
ACID = 7
PLASMA = 8
STONE = 9
STEAM = 10

LAVA = 11
WOOD = 12
LIGHTNING = 13
DIRT = 14
GLASS = 15
STEEL = 16
FOAM = 17
GAS = 18
PLASMA_BALL = 19
EXPLOSIVE = 20
ACID_WATER = 30
WET_SAND = 31

ICE = 32
SALT = 33
SPARK = 34
SLIME = 35
MUD = 36
METAL_POWDER = 37
ELECTRIC_CHARGE = 38
VINE = 39

# Colors for particles
COLORS = {
    EMPTY: (0, 0, 0),
    SAND: (194, 178, 128),
    WATER: (0, 0, 255),
    OIL: (20, 20, 20),
    FIRE: (255, 100, 0),
    SMOKE: (105, 105, 105),
    PLANT: (34, 139, 34),
    ACID: (0, 255, 0),
    PLASMA: (255, 0, 255),
    STONE: (128, 128, 128),
    STEAM: (211, 211, 211),
    LAVA: (255, 69, 0),
    WOOD: (139, 69, 19),
    LIGHTNING: (255, 255, 0),
    DIRT: (101, 67, 33),
    GLASS: (135, 206, 250),
    STEEL: (192, 192, 192),
    FOAM: (255, 250, 250),
    GAS: (169, 169, 169),
    PLASMA_BALL: (138, 43, 226),
    EXPLOSIVE: (255, 0, 0),
    ACID_WATER: (0, 100, 0),
    WET_SAND: (150, 130, 90),
    ICE: (180, 220, 255),
    SALT: (230, 230, 240),
    SPARK: (255, 255, 100),
    SLIME: (100, 255, 100),
    MUD: (101, 67, 33),
    METAL_POWDER: (150, 150, 150),
    ELECTRIC_CHARGE: (255, 255, 0),
    VINE: (34, 139, 34),       # dark brown rock cooled from lava
}

# Map keys to particles
PARTICLE_KEYS = {
    pygame.K_1: SAND,
    pygame.K_2: WATER,
    pygame.K_3: OIL,
    pygame.K_4: FIRE,
    pygame.K_5: SMOKE,
    pygame.K_6: PLANT,
    pygame.K_7: ACID,
    pygame.K_8: PLASMA,
    pygame.K_9: STONE,
    pygame.K_0: STEAM,

    pygame.K_q: LAVA,
    pygame.K_w: WOOD,
    pygame.K_e: LIGHTNING,
    pygame.K_r: DIRT,
    pygame.K_t: GLASS,
    pygame.K_y: STEEL,
    pygame.K_u: FOAM,
    pygame.K_i: GAS,
    pygame.K_o: PLASMA_BALL,
    pygame.K_p: EXPLOSIVE,
    pygame.K_a: ICE,
    pygame.K_s: SALT,
    pygame.K_d: SPARK,
    pygame.K_f: SLIME,
    pygame.K_g: MUD,
    pygame.K_h: METAL_POWDER,
    pygame.K_j: ELECTRIC_CHARGE,
    pygame.K_k: VINE,


    
}

def main():
    pygame.init()
    WIDTH, HEIGHT = 1800, 990
    screen = pygame.display.set_mode((WIDTH, HEIGHT))
    pygame.display.set_caption("Extended Particle Sandbox with Full Key Mapping")
    clock = pygame.time.Clock()
    FPS = 60

    PIXEL_SIZE = 4
    COLUMNS = WIDTH // PIXEL_SIZE
    ROWS = HEIGHT // PIXEL_SIZE

    grid = [[EMPTY for _ in range(COLUMNS)] for _ in range(ROWS)]
    updated = [[False for _ in range(COLUMNS)] for _ in range(ROWS)]

    selected_particle = SAND
    brush_size = 3

    font = pygame.font.SysFont("Arial", 18)

    def in_bounds(x, y):
        return 0 <= x < COLUMNS and 0 <= y < ROWS
    
    def neighbors(x, y):
        # Returns list of adjacent cell coordinates (left, right, up, down)
        return [(x - 1, y), (x + 1, y), (x, y - 1), (x, y + 1)]

    def swap(x1, y1, x2, y2):
        grid[y1][x1], grid[y2][x2] = grid[y2][x2], grid[y1][x1]
        updated[y1][x1], updated[y2][x2] = True, True

    def add_particles(mx, my, ptype, size):
        cx = mx // PIXEL_SIZE
        cy = my // PIXEL_SIZE
        for dx in range(-size, size + 1):
            for dy in range(-size, size + 1):
                x, y = cx + dx, cy + dy
                if in_bounds(x, y):
                    if grid[y][x] == EMPTY:
                        grid[y][x] = ptype



    def update_sand(x, y):
        below = (x, y + 1)

        # Fall through empty or water
        if in_bounds(*below):
            target = grid[below[1]][below[0]]
            if target == EMPTY or target == WATER:
                if target == WATER:
                    # Swap with water and become wet
                    grid[y][x] = WET_SAND
                swap(x, y, *below)
                updated[below[1]][below[0]] = True
                return

        # Diagonal movement
        options = []
        for dx in [-1, 1]:
            nx, ny = x + dx, y + 1
            if in_bounds(nx, ny):
                target = grid[ny][nx]
                if target == EMPTY or target == WATER:
                    options.append((nx, ny))

        if options:
            nx, ny = random.choice(options)
            if grid[ny][nx] == WATER:
                grid[y][x] = WET_SAND
            swap(x, y, nx, ny)
            updated[ny][nx] = True
            return

        # Wet if touching water
        for dx, dy in [(-1,0), (1,0), (0,-1), (0,1)]:
            nx, ny = x + dx, y + dy
            if in_bounds(nx, ny) and grid[ny][nx] == WATER:
                grid[y][x] = WET_SAND
                updated[y][x] = True
                return


        # Check if water is directly above this sand block
        if in_bounds(x, y - 1) and grid[y - 1][x] == WATER:
            # Turn this sand into wet sand and remove the water
            grid[y][x] = WET_SAND
            grid[y - 1][x] = EMPTY
            updated[y][x] = True
            updated[y - 1][x] = True

            # Propagate wet sand downward max 5 pixels from this point ONLY once here
            for depth in range(1, 6):
                ny = y + depth
                if in_bounds(x, ny) and grid[ny][x] == SAND:
                    grid[ny][x] = WET_SAND
                    updated[ny][x] = True
                else:
                    break


    def update_wet_sand(x, y):
        # Wet sand falls like sand
        if in_bounds(x, y + 1) and grid[y + 1][x] == EMPTY:
            swap(x, y, x, y + 1)
            updated[y + 1][x] = True
            return
        
        else:
            options = []
            if in_bounds(x - 1, y + 1) and grid[y + 1][x - 1] == EMPTY:
                options.append((x - 1, y + 1))
            if in_bounds(x + 1, y + 1) and grid[y + 1][x + 1] == EMPTY:
                options.append((x + 1, y + 1))
            if options:
                nx, ny = random.choice(options)
                swap(x, y, nx, ny)
                updated[ny][nx] = True
                return

        # Check if water is adjacent in all 8 directions
        adjacent_water = False
        for dx in [-1, 0, 1]:
            for dy in [-1, 0, 1]:
                if dx == 0 and dy == 0:
                    continue
                nx, ny = x + dx, y + dy
                if in_bounds(nx, ny) and grid[ny][nx] == WATER:
                    adjacent_water = True
                    break
            if adjacent_water:
                break

        # Only dry if no adjacent water
        if not adjacent_water:
            # Dry slowly with low chance per frame
            if random.random() < 0.001:
                grid[y][x] = SAND
                updated[y][x] = True

    

    def update_water(x, y):
        if in_bounds(x, y + 1) and grid[y + 1][x] == EMPTY:
            swap(x, y, x, y + 1)
        else:
            options = []
            if in_bounds(x - 1, y) and grid[y][x - 1] == EMPTY:
                options.append((x - 1, y))
            if in_bounds(x + 1, y) and grid[y][x + 1] == EMPTY:
                options.append((x + 1, y))
            if options:
                nx, ny = random.choice(options)
                swap(x, y, nx, ny)

    def update_fire(x, y):
        if random.random() < 0.1:
            grid[y][x] = SMOKE
        else:
            if in_bounds(x, y - 1) and grid[y - 1][x] == EMPTY:
                swap(x, y, x, y - 1)
            for dx, dy in [(-1,0),(1,0),(0,-1),(0,1)]:
                nx, ny = x + dx, y + dy
                if in_bounds(nx, ny) and grid[ny][nx] == WOOD:
                    grid[ny][nx] = FIRE

    def update_smoke(x, y):
        if random.random() < 0.02:
            grid[y][x] = EMPTY
            return
        if in_bounds(x, y - 1) and grid[y - 1][x] == EMPTY:
            swap(x, y, x, y - 1)
        else:
            options = []
            if in_bounds(x - 1, y) and grid[y][x - 1] == EMPTY:
                options.append((x - 1, y))
            if in_bounds(x + 1, y) and grid[y][x + 1] == EMPTY:
                options.append((x + 1, y))
            if options:
                nx, ny = random.choice(options)
                swap(x, y, nx, ny)

    def update_oil(x, y):
        if in_bounds(x, y + 1) and grid[y + 1][x] == EMPTY:
            swap(x, y, x, y + 1)
        else:
            options = []
            if in_bounds(x - 1, y) and grid[y][x - 1] == EMPTY:
                options.append((x - 1, y))
            if in_bounds(x + 1, y) and grid[y][x + 1] == EMPTY:
                options.append((x + 1, y))
            if options:
                nx, ny = random.choice(options)
                swap(x, y, nx, ny)
        for dx, dy in [(-1,0),(1,0),(0,-1),(0,1)]:
            nx, ny = x + dx, y + dy
            if in_bounds(nx, ny) and grid[ny][nx] in (FIRE,):
                if random.random() < 0.1:
                    grid[y][x] = FIRE

    def update_plant(x, y):
        for dx, dy in [(-1, 0), (1, 0), (0, -1), (0, 1)]:
            nx, ny = x + dx, y + dy
            if 0 <= nx < COLUMNS and 0 <= ny < ROWS:
                neighbor = grid[ny][nx]
                # Burn if next to fire
                if neighbor == FIRE:
                    grid[y][x] = FIRE
                    return
                # Grow into empty space with low probability
                elif neighbor == EMPTY and random.random() < 0.002:
                    grid[ny][nx] = PLANT
                    updated[ny][nx] = True


    def update_acid(x, y):
        # Gravity fall
        if in_bounds(x, y + 1) and grid[y + 1][x] == EMPTY:
            swap(x, y, x, y + 1)
        else:
            options = []
            if in_bounds(x - 1, y) and grid[y][x - 1] == EMPTY:
                options.append((x - 1, y))
            if in_bounds(x + 1, y) and grid[y][x + 1] == EMPTY:
                options.append((x + 1, y))
            if options:
                nx, ny = random.choice(options)
                swap(x, y, nx, ny)

        # Reaction with neighbors
        for dx, dy in [(-1, 0), (1, 0), (0, -1), (0, 1)]:
            nx, ny = x + dx, y + dy
            if in_bounds(nx, ny):
                neighbor = grid[ny][nx]

                # Corrode soft/organic materials only
                if neighbor in (WOOD, SAND, PLANT, DIRT, WET_SAND):
                    if random.random() < 0.05:
                        grid[ny][nx] = EMPTY
                        updated[ny][nx] = True

                # Mix with water to form acid water
                elif neighbor == WATER:
                    if random.random() < 0.1:
                        grid[ny][nx] = ACID_WATER
                        updated[ny][nx] = True

                        # Chance to emit smoke nearby
                        if random.random() < 0.2:
                            for sdx, sdy in [(-1, 0), (1, 0), (0, -1), (0, 1)]:
                                sx, sy = x + sdx, y + sdy
                                if in_bounds(sx, sy) and grid[sy][sx] == EMPTY:
                                    grid[sy][sx] = SMOKE
                                    updated[sy][sx] = True
                                    break                          

    def update_acid_water(x, y):
        # Flow downward or sideways like a fluid
        if in_bounds(x, y + 1) and grid[y + 1][x] == EMPTY:
            swap(x, y, x, y + 1)
        else:
            options = []
            if in_bounds(x - 1, y) and grid[y][x - 1] == EMPTY:
                options.append((x - 1, y))
            if in_bounds(x + 1, y) and grid[y][x + 1] == EMPTY:
                options.append((x + 1, y))
            if options:
                nx, ny = random.choice(options)
                swap(x, y, nx, ny)

        # Mild corrosive behavior
        for dx, dy in [(-1, 0), (1, 0), (0, -1), (0, 1)]:
            nx, ny = x + dx, y + dy
            if in_bounds(nx, ny):
                neighbor = grid[ny][nx]

                # Slowly corrode soft materials
                if neighbor in (WOOD, SAND, PLANT, DIRT, WET_SAND) and random.random() < 0.02:
                    grid[ny][nx] = EMPTY
                    updated[ny][nx] = True

        # Occasionally emit smoke
        if random.random() < 0.02:
            for dx, dy in [(-1, 0), (1, 0), (0, -1), (0, 1)]:
                nx, ny = x + dx, y + dy
                if in_bounds(nx, ny) and grid[ny][nx] == EMPTY:
                    grid[ny][nx] = SMOKE
                    updated[ny][nx] = True
                    break
                          


    def update_stone(x, y):
        # stone acts like sand but heavier
        if in_bounds(x, y + 1) and grid[y + 1][x] == EMPTY:
            swap(x, y, x, y + 1)

    def update_steam(x, y):
        # steam rises and dissipates
        if random.random() < 0.02:
            grid[y][x] = EMPTY
            return
        if in_bounds(x, y - 1) and grid[y - 1][x] == EMPTY:
            swap(x, y, x, y - 1)

    # Add more update functions for lava, lightning, dirt, glass, steel, foam, gas, plasma_ball, explosive etc. as needed.
    # For brevity, these will just act like smoke or empty behavior here.

    def update_generic_gas(x, y):
        # gas rises like smoke
        if random.random() < 0.02:
            grid[y][x] = EMPTY
            return
        if in_bounds(x, y - 1) and grid[y - 1][x] == EMPTY:
            swap(x, y, x, y - 1)

    def update_explosive(x, y):
        # Explode randomly and clear neighbors
        if random.random() < 0.005:
            for dx in range(-1, 2):
                for dy in range(-1, 2):
                    nx, ny = x + dx, y + dy
                    if in_bounds(nx, ny):
                        grid[ny][nx] = EMPTY
        else:
            pass

    def update_lightning(x, y):
        # Moves quickly sideways randomly
        if random.random() < 0.1:
            options = []
            if in_bounds(x - 1, y) and grid[y][x - 1] == EMPTY:
                options.append((x - 1, y))
            if in_bounds(x + 1, y) and grid[y][x + 1] == EMPTY:
                options.append((x + 1, y))
            if options:
                nx, ny = random.choice(options)
                swap(x, y, nx, ny)

    def update_ice(x, y):
        # Ice behaves like solid but can slowly melt if near fire or warmth
        if any(grid[ny][nx] == FIRE for nx, ny in neighbors(x, y)):
            if random.random() < 0.02:
                grid[y][x] = WATER
                updated[y][x] = True

    def update_salt(x, y):
        # Salt slowly dissolves near water turning to EMPTY
        if any(grid[ny][nx] == WATER for nx, ny in neighbors(x, y)):
            if random.random() < 0.03:
                grid[y][x] = EMPTY
                updated[y][x] = True

    def update_spark(x, y):
        # Spark moves randomly and dissipates fast
        directions = [(1,0), (-1,0), (0,1), (0,-1)]
        if random.random() < 0.5:
            dx, dy = random.choice(directions)
            nx, ny = x + dx, y + dy
            if in_bounds(nx, ny) and grid[ny][nx] == EMPTY:
                swap(x, y, nx, ny)
        if random.random() < 0.1:
            grid[y][x] = EMPTY
            updated[y][x] = True

    def update_slime(x, y):
        # Slime moves slowly like a viscous liquid
        if in_bounds(x, y + 1) and grid[y + 1][x] == EMPTY:
            swap(x, y, x, y + 1)
        else:
            if random.random() < 0.3:
                if in_bounds(x - 1, y) and grid[y][x - 1] == EMPTY:
                    swap(x, y, x - 1, y)
                elif in_bounds(x + 1, y) and grid[y][x + 1] == EMPTY:
                    swap(x, y, x + 1, y)

    def update_mud(x, y):
        # Mud behaves like dirt, but can slowly dry into dirt when no water nearby
        if not any(grid[ny][nx] == WATER for nx, ny in neighbors(x, y)):
            if random.random() < 0.02:
                grid[y][x] = DIRT
                updated[y][x] = True

    def update_metal_powder(x, y):
        # Metal powder falls down like sand but settles on top of solid stuff
        if in_bounds(x, y + 1) and grid[y + 1][x] == EMPTY:
            swap(x, y, x, y + 1)
        else:
            # Try sideways
            options = []
            if in_bounds(x - 1, y + 1) and grid[y + 1][x - 1] == EMPTY:
                options.append((x - 1, y + 1))
            if in_bounds(x + 1, y + 1) and grid[y + 1][x + 1] == EMPTY:
                options.append((x + 1, y + 1))
            if options:
                nx, ny = random.choice(options)
                swap(x, y, nx, ny)

    def update_electric_charge(x, y):
        # Electric charge randomly jumps to adjacent conductive elements (metal powder, steel)
        for nx, ny in neighbors(x, y):
            if in_bounds(nx, ny) and grid[ny][nx] in (METAL_POWDER, STEEL):
                if random.random() < 0.2:
                    grid[ny][nx] = ELECTRIC_CHARGE
                    updated[ny][nx] = True
        # Electric charge dissipates quickly
        if random.random() < 0.1:
            grid[y][x] = EMPTY
            updated[y][x] = True

    def update_vine(x, y):
        # Vine grows upward slowly if empty above, otherwise spreads sideways on wood or dirt
        if in_bounds(x, y - 1) and grid[y - 1][x] == EMPTY:
            if random.random() < 0.05:
                grid[y - 1][x] = VINE
                updated[y - 1][x] = True
        else:
            for dx in [-1, 1]:
                nx = x + dx
                if in_bounds(nx, y) and grid[y][nx] in (WOOD, DIRT) and random.random() < 0.03:
                    grid[y][nx] = VINE
                    updated[y][nx] = True

    def update_lava(x, y):
        # Lava flows like a slow, viscous liquid
        moved = False
        if in_bounds(x, y + 1) and grid[y + 1][x] == EMPTY and random.random() < 0.4:
            swap(x, y, x, y + 1)
            moved = True
        elif random.random() < 0.2:
            directions = []
            if in_bounds(x - 1, y) and grid[y][x - 1] == EMPTY:
                directions.append((x - 1, y))
            if in_bounds(x + 1, y) and grid[y][x + 1] == EMPTY:
                directions.append((x + 1, y))
            if directions:
                nx, ny = random.choice(directions)
                swap(x, y, nx, ny)
                moved = True

        # Heat interactions
        for dx, dy in [(-1, 0), (1, 0), (0, -1), (0, 1)]:
            nx, ny = x + dx, y + dy
            if not in_bounds(nx, ny):
                continue
            neighbor = grid[ny][nx]

            # Water to steam
            if neighbor == WATER:
                grid[ny][nx] = STEAM
                updated[ny][nx] = True

            # Ice to steam (instantly)
            elif neighbor == ICE:
                grid[ny][nx] = STEAM
                updated[ny][nx] = True

            # Sand to glass
            elif neighbor == SAND:
                grid[ny][nx] = GLASS
                updated[ny][nx] = True

            # Burnable materials to fire
            elif neighbor in (PLANT, WOOD, OIL, FOAM, SLIME, VINE, GAS, SPARK):
                grid[ny][nx] = FIRE
                updated[ny][nx] = True

            # Acid water evaporates (or turns to gas)
            elif neighbor == ACID_WATER and random.random() < 0.1:
                grid[ny][nx] = GAS
                updated[ny][nx] = True

           





    def update_particle(x, y):
        p = grid[y][x]
        if p == SAND:
            update_sand(x, y)
        elif p == WATER:
            update_water(x, y)
        elif p == FIRE:
            update_fire(x, y)
        elif p == SMOKE:
            update_smoke(x, y)
        elif p == OIL:
            update_oil(x, y)
        elif p == PLANT:
            update_plant(x, y)
        elif p == ACID:
            update_acid(x, y)
        elif p == STONE:
            update_stone(x, y)
        elif p == STEAM:
            update_steam(x, y)
        elif p == LAVA:
            update_lava(x, y)
        elif p == WOOD:
            # wood mostly static
            pass
        elif p == LIGHTNING:
            update_lightning(x, y)
        elif p == DIRT:
            # dirt like sand but no movement here
            pass
        elif p == GLASS:
            # glass static
            pass
        elif p == STEEL:
            # steel static
            pass
        elif p == FOAM:
            update_generic_gas(x, y)
        elif p == GAS:
            update_generic_gas(x, y)
        elif p == PLASMA:
            update_generic_gas(x, y)
        elif p == PLASMA_BALL:
            update_generic_gas(x, y)
        elif p == EXPLOSIVE:
            update_explosive(x, y)
        elif p == ACID_WATER:
            update_acid_water(x,y)
        elif p == WET_SAND:
            update_wet_sand(x,y)
        elif p == ICE:
            update_ice(x, y)
        elif p == SALT:
            update_salt(x, y)
        elif p == SPARK:
            update_spark(x, y)
        elif p == SLIME:
            update_slime(x, y)
        elif p == MUD:
            update_mud(x, y)
        elif p == METAL_POWDER:
            update_metal_powder(x, y)
        elif p == ELECTRIC_CHARGE:
            update_electric_charge(x, y)
        elif p == VINE:
            update_vine(x, y)
        else:
            pass

    running = True
    while running:
        screen.fill((0, 0, 0))

        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                running = False

            elif event.type == pygame.KEYDOWN:
                if event.key in PARTICLE_KEYS:
                    selected_particle = PARTICLE_KEYS[event.key]

            elif event.type == pygame.MOUSEWHEEL:
                brush_size += event.y
                if brush_size < 1:
                    brush_size = 1
                elif brush_size > 10:
                    brush_size = 10

        mouse_pressed = pygame.mouse.get_pressed()
        if mouse_pressed[0]:
            mx, my = pygame.mouse.get_pos()
            add_particles(mx, my, selected_particle, brush_size)

        # Reset updated grid
        for y in range(ROWS):
            for x in range(COLUMNS):
                updated[y][x] = False

        # Update particles bottom to top
        for y in reversed(range(ROWS)):
            for x in range(COLUMNS):
                if grid[y][x] != EMPTY and not updated[y][x]:
                    update_particle(x, y)

        # Draw particles
        for y in range(ROWS):
            for x in range(COLUMNS):
                p = grid[y][x]
                if p != EMPTY:
                    color = COLORS.get(p, (255, 255, 255))
                    rect = pygame.Rect(x * PIXEL_SIZE, y * PIXEL_SIZE, PIXEL_SIZE, PIXEL_SIZE)
                    pygame.draw.rect(screen, color, rect)

        # Display selected particle and brush size
        particle_names = {
            SAND: "Sand",
            WATER: "Water",
            OIL: "Oil",
            FIRE: "Fire",
            SMOKE: "Smoke",
            PLANT: "Plant",
            ACID: "Acid",
            PLASMA: "Plasma",
            STONE: "Stone",
            STEAM: "Steam",
            LAVA: "Lava",
            WOOD: "Wood",
            LIGHTNING: "Lightning",
            DIRT: "Dirt",
            GLASS: "Glass",
            STEEL: "Steel",
            FOAM: "Foam",
            GAS: "Gas",
            PLASMA_BALL: "Plasma Ball",
            EXPLOSIVE: "Explosive",
            EMPTY: "Erase",
            ICE: "Ice",
            SALT: "Salt",
            SPARK: "Spark",
            SLIME: "Slime",
            MUD: "Mud",
            METAL_POWDER: "Metal Powder",
            ELECTRIC_CHARGE: "Electric Charge",
            VINE: "Vine"
        }
        selected_name = particle_names.get(selected_particle, "Unknown")
        info_text = f"Selected: {selected_name} (1-0, Q-P keys) | Brush size: {brush_size} (Mouse wheel)"
        text_surf = font.render(info_text, True, (255, 255, 255))
        screen.blit(text_surf, (10, HEIGHT - 30))

        pygame.display.flip()
        clock.tick(FPS)

    pygame.quit()

if __name__ == "__main__":
    main()