slime_sim.py

"""
CPU port of the Unity slime simulation.

The original project uses Unity compute shaders to move agents, deposit trails,
diffuse/decay the trail map, and build a coloured display texture. This Python
version follows the same logic on the CPU with NumPy. Visualisation is done
with matplotlib; the core simulation is kept separate so it can also run
headless.
"""

import math
import random
from dataclasses import dataclass, field
from enum import Enum
from typing import List, Sequence, Tuple

import numpy as np

try:
    import taichi as ti  # type: ignore
except Exception:  # noqa: BLE001
    ti = None


class SpawnMode(Enum):
    RANDOM = "random"
    POINT = "point"
    INWARD_CIRCLE = "inward_circle"
    RANDOM_CIRCLE = "random_circle"


@dataclass
class SpeciesSettings:
    move_speed: float = 60.0
    turn_speed: float = 0.25  # revolutions per second
    sensor_angle_degrees: float = 22.5
    sensor_offset_dst: float = 15.0
    sensor_size: int = 1
    colour: Tuple[float, float, float, float] = (1.0, 1.0, 1.0, 1.0)


@dataclass
class SlimeSettings:
    # Simulation
    steps_per_frame: int = 1
    width: int = 640
    height: int = 360
    num_agents: int = 20000
    spawn_mode: SpawnMode = SpawnMode.RANDOM

    # Trail
    trail_weight: float = 1.0
    decay_rate: float = 0.25
    diffuse_rate: float = 15.0

    species_settings: Sequence[SpeciesSettings] = field(default_factory=lambda: [
        SpeciesSettings(colour=(1.0, 0.2, 0.2, 1.0)),
        SpeciesSettings(colour=(0.2, 1.0, 0.2, 1.0)),
        SpeciesSettings(colour=(0.2, 0.4, 1.0, 1.0)),
    ])


class Simulation:
    """
    CPU recreation of the Unity compute-shader simulation.
    Fallback for when Taichi is not available.
    """

    def __init__(self, settings: SlimeSettings, delta_time: float = 1.0 / 60.0):
        self.s = settings
        self.dt = float(delta_time)

        self.num_species = min(4, len(self.s.species_settings))
        self.positions = np.zeros((self.s.num_agents, 2), dtype=np.float32)
        self.angles = np.zeros((self.s.num_agents,), dtype=np.float32)
        self.species_index = np.zeros((self.s.num_agents,), dtype=np.int32)
        self.species_mask = np.zeros((self.s.num_agents, 4), dtype=np.float32)

        self.trail_map = np.zeros((self.s.height, self.s.width, 4), dtype=np.float32)
        self._init_agents()

    # ------------------------------------------------------------------ setup
    def _init_agents(self) -> None:
        centre = np.array([self.s.width / 2.0, self.s.height / 2.0], dtype=np.float32)
        for i in range(self.s.num_agents):
            angle = random.random() * math.tau
            pos = centre.copy()

            if self.s.spawn_mode == SpawnMode.POINT:
                pos = centre.copy()
            elif self.s.spawn_mode == SpawnMode.RANDOM:
                pos = np.array(
                    [random.uniform(0, self.s.width), random.uniform(0, self.s.height)],
                    dtype=np.float32,
                )
            elif self.s.spawn_mode == SpawnMode.INWARD_CIRCLE:
                r = random.random() * 0.5 * self.s.height
                phi = random.random() * math.tau
                pos = centre + np.array([math.cos(phi), math.sin(phi)]) * r
                to_centre = (centre - pos)
                angle = math.atan2(to_centre[1], to_centre[0])
            elif self.s.spawn_mode == SpawnMode.RANDOM_CIRCLE:
                r = random.random() * 0.15 * self.s.height
                phi = random.random() * math.tau
                pos = centre + np.array([math.cos(phi), math.sin(phi)]) * r
            else:
                pos = centre.copy()

            if self.num_species == 1:
                idx = 0
            else:
                idx = random.randrange(self.num_species)
            mask = np.zeros(4, dtype=np.float32)
            mask[idx] = 1.0

            self.positions[i] = pos
            self.angles[i] = angle
            self.species_index[i] = idx
            self.species_mask[i] = mask

    # ----------------------------------------------------------------- update
    def step(self) -> None:
        for _ in range(self.s.steps_per_frame):
            self._update_agents()
            self._diffuse_and_decay()

    def _update_agents(self) -> None:
        for i in range(self.s.num_agents):
            idx = int(self.species_index[i])
            settings = self.s.species_settings[idx]
            pos = self.positions[i]

            weight_forward = self._sense(i, settings, 0.0)
            sensor_angle = math.radians(settings.sensor_angle_degrees)
            weight_left = self._sense(i, settings, sensor_angle)
            weight_right = self._sense(i, settings, -sensor_angle)

            random_steer = random.random()
            turn_speed = settings.turn_speed * math.tau

            # Steer logic mirrors the compute shader
            if weight_forward > weight_left and weight_forward > weight_right:
                pass
            elif weight_forward < weight_left and weight_forward < weight_right:
                self.angles[i] += (random_steer - 0.5) * 2 * turn_speed * self.dt
            elif weight_right > weight_left:
                self.angles[i] -= random_steer * turn_speed * self.dt
            elif weight_left > weight_right:
                self.angles[i] += random_steer * turn_speed * self.dt

            direction = np.array([math.cos(self.angles[i]), math.sin(self.angles[i])])
            new_pos = pos + direction * self.dt * settings.move_speed

            if (
                new_pos[0] < 0
                or new_pos[0] >= self.s.width
                or new_pos[1] < 0
                or new_pos[1] >= self.s.height
            ):
                rand_angle = random.random() * math.tau
                new_pos[0] = min(self.s.width - 1, max(0.0, new_pos[0]))
                new_pos[1] = min(self.s.height - 1, max(0.0, new_pos[1]))
                self.angles[i] = rand_angle
            else:
                x = int(new_pos[0])
                y = int(new_pos[1])
                mask = self.species_mask[i]
                self.trail_map[y, x] = np.minimum(
                    1.0, self.trail_map[y, x] + mask * self.s.trail_weight * self.dt
                )

            self.positions[i] = new_pos

    def _sense(self, i: int, settings: SpeciesSettings, sensor_angle_offset: float) -> float:
        angle = self.angles[i] + sensor_angle_offset
        sensor_dir = np.array([math.cos(angle), math.sin(angle)])
        sensor_pos = self.positions[i] + sensor_dir * settings.sensor_offset_dst
        cx = int(sensor_pos[0])
        cy = int(sensor_pos[1])

        cx = max(0, min(self.s.width - 1, cx))
        cy = max(0, min(self.s.height - 1, cy))

        sense_weight = self.species_mask[i] * 2.0 - 1.0
        size = settings.sensor_size
        total = 0.0
        for dx in range(-size, size + 1):
            for dy in range(-size, size + 1):
                sx = max(0, min(self.s.width - 1, cx + dx))
                sy = max(0, min(self.s.height - 1, cy + dy))
                total += float(np.dot(sense_weight, self.trail_map[sy, sx]))
        return total

    # ---------------------------------------------------------- post-process
    def _diffuse_and_decay(self) -> None:
        # 3x3 blur implemented with padding + slicing to avoid scipy dependency
        padded = np.pad(self.trail_map, ((1, 1), (1, 1), (0, 0)), mode="edge")
        sum_map = np.zeros_like(self.trail_map)
        for dx in (-1, 0, 1):
            for dy in (-1, 0, 1):
                sum_map += padded[1 + dy : 1 + dy + self.s.height, 1 + dx : 1 + dx + self.s.width]

        blurred = sum_map / 9.0
        diffuse_weight = min(1.0, self.s.diffuse_rate * self.dt)
        blurred = self.trail_map * (1.0 - diffuse_weight) + blurred * diffuse_weight
        self.trail_map = np.maximum(0.0, blurred - self.s.decay_rate * self.dt)

    # --------------------------------------------------------------- display
    def build_colour_map(self) -> np.ndarray:
        colour_map = np.zeros((self.s.height, self.s.width, 3), dtype=np.float32)
        for i in range(self.num_species):
            mask = self.trail_map[:, :, i]
            colour = np.array(self.s.species_settings[i].colour[:3], dtype=np.float32)
            colour_map += colour * mask[..., None]
        return np.clip(colour_map, 0.0, 1.0)


# ------------------------------ Taichi backend ------------------------------
@ti.data_oriented
class TaichiSimulation:
    """High-performance version using Taichi kernels."""

    def __init__(self, settings: SlimeSettings, delta_time: float = 1.0 / 60.0):
        if ti is None:
            raise ImportError("taichi is not installed")
        self.s = settings
        self.dt = float(delta_time)

        self.num_species = min(4, len(self.s.species_settings))

        # Initialise Taichi with GPU fallback (copied from your working example)
        try:
            ti.init(arch=ti.cuda, device_memory_fraction=0.6)
        except Exception:
            ti.init(arch=ti.gpu, device_memory_fraction=0.6)

        self.positions = ti.Vector.field(2, dtype=ti.f32, shape=self.s.num_agents)
        self.angles = ti.field(dtype=ti.f32, shape=self.s.num_agents)
        self.species_index = ti.field(dtype=ti.i32, shape=self.s.num_agents)
        self.species_mask = ti.Vector.field(4, dtype=ti.f32, shape=self.s.num_agents)

        self.trail_map = ti.Vector.field(4, dtype=ti.f32, shape=(self.s.width, self.s.height))
        self.diffused_map = ti.Vector.field(4, dtype=ti.f32, shape=(self.s.width, self.s.height))
        self.colour_map = ti.Vector.field(3, dtype=ti.f32, shape=(self.s.width, self.s.height))

        # species parameters
        self.move_speed = ti.field(dtype=ti.f32, shape=self.num_species)
        self.turn_speed = ti.field(dtype=ti.f32, shape=self.num_species)
        self.sensor_angle = ti.field(dtype=ti.f32, shape=self.num_species)
        self.sensor_offset = ti.field(dtype=ti.f32, shape=self.num_species)
        self.sensor_size = ti.field(dtype=ti.i32, shape=self.num_species)
        self.species_colour = ti.Vector.field(4, dtype=ti.f32, shape=self.num_species)

        self.trail_weight = ti.field(dtype=ti.f32, shape=())
        self.decay_rate = ti.field(dtype=ti.f32, shape=())
        self.diffuse_rate = ti.field(dtype=ti.f32, shape=())

        self._init_constants()
        self._init_agents()

    def _init_constants(self) -> None:
        for i, sp in enumerate(self.s.species_settings[: self.num_species]):
            self.move_speed[i] = sp.move_speed
            self.turn_speed[i] = sp.turn_speed * math.tau
            self.sensor_angle[i] = math.radians(sp.sensor_angle_degrees)
            self.sensor_offset[i] = sp.sensor_offset_dst
            self.sensor_size[i] = sp.sensor_size
            self.species_colour[i] = ti.Vector(list(sp.colour[:4]))
        self.trail_weight[None] = self.s.trail_weight
        self.decay_rate[None] = self.s.decay_rate
        self.diffuse_rate[None] = self.s.diffuse_rate

    def _init_agents(self) -> None:
        centre = ti.Vector([self.s.width / 2.0, self.s.height / 2.0])
        positions_np = np.zeros((self.s.num_agents, 2), dtype=np.float32)
        angles_np = np.zeros((self.s.num_agents,), dtype=np.float32)
        idx_np = np.zeros((self.s.num_agents,), dtype=np.int32)
        mask_np = np.zeros((self.s.num_agents, 4), dtype=np.float32)

        for i in range(self.s.num_agents):
            angle = random.random() * math.tau
            pos = centre.to_numpy()
            if self.s.spawn_mode == SpawnMode.POINT:
                pos = centre.to_numpy()
            elif self.s.spawn_mode == SpawnMode.RANDOM:
                pos = np.array(
                    [random.uniform(0, self.s.width), random.uniform(0, self.s.height)],
                    dtype=np.float32,
                )
            elif self.s.spawn_mode == SpawnMode.INWARD_CIRCLE:
                r = random.random() * 0.5 * self.s.height
                phi = random.random() * math.tau
                pos = centre.to_numpy() + np.array([math.cos(phi), math.sin(phi)]) * r
                angle = math.atan2(centre[1] - pos[1], centre[0] - pos[0])
            elif self.s.spawn_mode == SpawnMode.RANDOM_CIRCLE:
                r = random.random() * 0.15 * self.s.height
                phi = random.random() * math.tau
                pos = centre.to_numpy() + np.array([math.cos(phi), math.sin(phi)]) * r
            if self.num_species == 1:
                idx = 0
            else:
                idx = random.randrange(self.num_species)
            mask = np.zeros(4, dtype=np.float32)
            mask[idx] = 1.0
            positions_np[i] = pos
            angles_np[i] = angle
            idx_np[i] = idx
            mask_np[i] = mask

        self.positions.from_numpy(positions_np)
        self.angles.from_numpy(angles_np)
        self.species_index.from_numpy(idx_np)
        self.species_mask.from_numpy(mask_np)
        self.trail_map.fill(0)
        self.diffused_map.fill(0)
        self.colour_map.fill(0)

    @ti.func
    def _sense(self, i, angle_offset):
        angle = self.angles[i] + angle_offset
        idx = self.species_index[i]
        sensor_dir = ti.Vector([ti.cos(angle), ti.sin(angle)])
        sensor_pos = self.positions[i] + sensor_dir * self.sensor_offset[idx]
        cx = ti.cast(sensor_pos[0], ti.i32)
        cy = ti.cast(sensor_pos[1], ti.i32)
        cx = ti.max(0, ti.min(self.s.width - 1, cx))
        cy = ti.max(0, ti.min(self.s.height - 1, cy))
        size = self.sensor_size[idx]
        sense_weight = self.species_mask[i] * 2.0 - 1.0
        total = 0.0
        for dx in range(-size, size + 1):
            for dy in range(-size, size + 1):
                sx = ti.max(0, ti.min(self.s.width - 1, cx + dx))
                sy = ti.max(0, ti.min(self.s.height - 1, cy + dy))
                total += sense_weight.dot(self.trail_map[sx, sy])
        return total

    # FIXED: Use float type hint (Python native) instead of ti.template or ti.f32
    @ti.kernel
    def _update_agents(self, dt: ti.f32):
        for i in range(self.s.num_agents):
            idx = self.species_index[i]
            weight_f = self._sense(i, 0.0)
            sensor_ang = self.sensor_angle[idx]
            weight_l = self._sense(i, sensor_ang)
            weight_r = self._sense(i, -sensor_ang)
            rand = ti.random()
            turn = self.turn_speed[idx]

            if weight_f > weight_l and weight_f > weight_r:
                pass
            elif weight_f < weight_l and weight_f < weight_r:
                self.angles[i] += (rand - 0.5) * 2 * turn * dt
            elif weight_r > weight_l:
                self.angles[i] -= rand * turn * dt
            elif weight_l > weight_r:
                self.angles[i] += rand * turn * dt

            direction = ti.Vector([ti.cos(self.angles[i]), ti.sin(self.angles[i])])
            new_pos = self.positions[i] + direction * dt * self.move_speed[idx]

            if (
                new_pos[0] < 0
                or new_pos[0] >= self.s.width
                or new_pos[1] < 0
                or new_pos[1] >= self.s.height
            ):
                rand_angle = ti.random() * 2 * math.pi
                new_pos[0] = ti.min(self.s.width - 1.0, ti.max(0.0, new_pos[0]))
                new_pos[1] = ti.min(self.s.height - 1.0, ti.max(0.0, new_pos[1]))
                self.angles[i] = rand_angle
            else:
                x = ti.cast(new_pos[0], ti.i32)
                y = ti.cast(new_pos[1], ti.i32)
                self.trail_map[x, y] = ti.min(
                    ti.Vector([1.0, 1.0, 1.0, 1.0]),
                    self.trail_map[x, y] + self.species_mask[i] * self.trail_weight[None] * dt,
                )
            self.positions[i] = new_pos

    # FIXED: Use float type hint
    @ti.kernel
    def _diffuse(self, dt: ti.f32):
        for x, y in self.trail_map:
            if x < 0 or x >= self.s.width or y < 0 or y >= self.s.height:
                continue
            acc = ti.Vector([0.0, 0.0, 0.0, 0.0])
            for dx in range(-1, 2):
                for dy in range(-1, 2):
                    sx = ti.min(self.s.width - 1, ti.max(0, x + dx))
                    sy = ti.min(self.s.height - 1, ti.max(0, y + dy))
                    acc += self.trail_map[sx, sy]
            blurred = acc / 9.0
            diffuse_w = ti.min(1.0, self.diffuse_rate[None] * dt)
            blurred = self.trail_map[x, y] * (1 - diffuse_w) + blurred * diffuse_w
            self.diffused_map[x, y] = ti.max(
                ti.Vector([0.0, 0.0, 0.0, 0.0]), blurred - self.decay_rate[None] * dt
            )

    @ti.kernel
    def _copy_diffused(self):
        for x, y in self.trail_map:
            self.trail_map[x, y] = self.diffused_map[x, y]

    @ti.kernel
    def _build_colour(self):
        for x, y in self.colour_map:
            m = self.trail_map[x, y]
            col = ti.Vector([0.0, 0.0, 0.0])
            for i in range(self.num_species):
                mask = ti.Vector([1.0 if i == 0 else 0.0, 1.0 if i == 1 else 0.0, 1.0 if i == 2 else 0.0, 1.0 if i == 3 else 0.0])
                col += self.species_colour[i].xyz * m.dot(mask)
            self.colour_map[x, y] = ti.min(col, ti.Vector([1.0, 1.0, 1.0]))

    def step(self) -> None:
        for _ in range(self.s.steps_per_frame):
            self._update_agents(self.dt)
            self._diffuse(self.dt)
            self._copy_diffused()

    def build_colour_map(self) -> np.ndarray:
        self._build_colour()
        return np.transpose(self.colour_map.to_numpy(), (1, 0, 2))


if __name__ == "__main__":
    # Ensure pygame is available
    try:
        import pygame
    except ImportError:
        print("Please install pygame: pip install pygame")
        exit()

    def run_pygame(
        display_scale: float = 0.8,
        target_fps: int = 60,
        use_taichi: bool = True,
        num_species: int = 3,
        board_scale: float = 1.0,
    ) -> None:
        """Live window; Taichi backend for speed when available."""
        base_species = [
            SpeciesSettings(colour=(1.0, 0.2, 0.2, 1.0), move_speed=70, turn_speed=0.35),
            SpeciesSettings(colour=(0.2, 1.0, 0.6, 1.0), move_speed=65, turn_speed=0.25),
            SpeciesSettings(colour=(0.2, 0.4, 1.0, 1.0), move_speed=75, turn_speed=0.3),
            # SpeciesSettings(colour=(1.0, 1.0, 0.25, 1.0), move_speed=68, turn_speed=0.28),
        ]
        species_subset = base_species[: max(1, min(num_species, len(base_species)))]
        base_w, base_h = 1000, 680  # large viewable board
        width = int(base_w * board_scale)
        height = int(base_h * board_scale)
        agent_scale = board_scale * board_scale
        settings = SlimeSettings(
            steps_per_frame=10,
            width=width,
            height=height,
            num_agents=int(155000 * agent_scale),
            spawn_mode=SpawnMode.INWARD_CIRCLE,
            trail_weight=1.0,
            decay_rate=0.17,  # slower decay so trails persist
            diffuse_rate=7.5,  # slightly gentler diffusion
            species_settings=species_subset,
        )

        if use_taichi and ti is not None:
            backend = TaichiSimulation(settings)
            print("Backend: Taichi (GPU)")
        else:
            backend = Simulation(settings)
            print("Backend: NumPy (CPU)")

        pygame.init()
        window_size = (int(settings.width * display_scale), int(settings.height * display_scale))
        screen = pygame.display.set_mode(window_size)
        pygame.display.set_caption("Slime Simulation")
        clock = pygame.time.Clock()

        running = True
        while running:
            for event in pygame.event.get():
                if event.type == pygame.QUIT or (
                    event.type == pygame.KEYDOWN and event.key == pygame.K_ESCAPE
                ):
                    running = False

            # Simulation Step
            backend.step()

            # Rendering
            frame = (backend.build_colour_map() * 255).astype(np.uint8)
            surface = pygame.surfarray.make_surface(np.transpose(frame, (1, 0, 2)))
            if display_scale != 1.0:
                surface = pygame.transform.smoothscale(surface, window_size)
            screen.blit(surface, (0, 0))
            pygame.display.flip()

            clock.tick(target_fps)
            pygame.display.set_caption(f"Slime Simulation - FPS: {clock.get_fps():.1f}")

        pygame.quit()

    # Automatically select Taichi when installed
    run_pygame(use_taichi=ti is not None)