simulation.py

import random

from pyglm import glm

from body import Body
from spatial_grid import SpatialGrid


class Simulation:
    """
    Simulation class for managing a collection of bodies and their interactions.
    This class manages the setup and update of a physics simulation containing
    multiple bodies. It handles body initialization, position tracking,
    collision detection, and resolution.
    """

    def __init__(self, N: int, aspect_ratio: float = 1.0) -> None:
        """
        Initialize simulation with N bodies.
        Args:
            N: Number of bodies.
            aspect_ratio: Display aspect ratio. Defaults to 1.0.
        """

        self._N = N

        # Internal lists to hold bodies and their properties
        self._bodies = []
        self._positions = []
        self._sizes = []
        self._colors = []

        # Simulation properties from GUI
        self._gravity_vector = glm.vec2(0.0, 0.0)
        self._color_based_on_speed = True
        self._aspect_ratio = aspect_ratio
        self._color_data = [0.01, 0.4, [0, 0, 1], [1, 0, 0]]

        # Parameters for spatial grid
        self._max_radius, self._min_radius = Body.calculate_size_range(N)
        self._grid = SpatialGrid(2 * self._max_radius)

    @property
    def positions(self) -> list[glm.vec2]:
        return self._positions

    @property
    def sizes(self) -> list[float]:
        return self._sizes

    @property
    def colors(self) -> list[glm.vec3]:
        return self._colors

    @property
    def gravity_vector(self) -> glm.vec2:
        return self._gravity_vector

    @gravity_vector.setter
    def gravity_vector(self, new: glm.vec2) -> glm.vec2:
        self._gravity_vector = new

    @property
    def color_based_on_speed(self) -> bool:
        return self._color_based_on_speed

    @color_based_on_speed.setter
    def color_based_on_speed(self, new: bool) -> None:
        self._color_based_on_speed = new

    @property
    def aspect_ratio(self) -> float:
        return self._aspect_ratio

    @aspect_ratio.setter
    def aspect_ratio(self, new: float) -> float:
        self._aspect_ratio = new

    @property
    def color_data(self) -> list:
        return self._color_data

    @color_data.setter
    def color_data(self, data: list) -> None:
        self._color_data = data

    def setup_simulation(self) -> None:
        """Initialize the simulation with N bodies with random properties"""
        for _ in range(self._N):
            pos = glm.vec2(
                random.uniform(-0.9, 0.9) * self._aspect_ratio,
                random.uniform(-0.9, 0.9),
            )
            vel = glm.vec2(random.uniform(-0.2, 0.2), random.uniform(-0.2, 0.2))
            size = random.uniform(self._min_radius, self._max_radius)
            self._sizes.append(size)

            mass = random.uniform(self._max_radius, self._max_radius * 2)

            color = glm.vec3(
                random.uniform(0.3, 1), random.uniform(0.3, 1), random.uniform(0.3, 1)
            )

            body = Body(pos, vel, size, mass, color, self._aspect_ratio)
            self._bodies.append(body)

    def update_bodies(self, dT: float, body: Body) -> None:
        """Update individual body properties such as position, color, and acceleration"""

        if self.color_based_on_speed:
            body.set_color_based_on_speed(self._color_data)
        else:
            body.reset_color()

        body.update(dT)
        body.acc = self._gravity_vector
        body.aspect_ratio = self._aspect_ratio

    def update(self, dT: float) -> None:
        """Update the simulation state by handling collisions and updating body positions"""

        for body in self._bodies:
            self._grid.insert(body)
            self.update_bodies(dT, body)

        checked_pairs = set()
        for body in self._bodies:
            nearby = self._grid.get_nearby(body)
            for other in nearby:
                if body is other:
                    continue

                pair = tuple(sorted([id(body), id(other)]))
                if pair in checked_pairs:
                    continue
                checked_pairs.add(pair)

                if Body.is_colliding(body, other):
                    Body.resolve_collision(body, other)

        self._grid.clear()
        self._positions = [body.pos for body in self._bodies]
        self._colors = [body.visibleColor for body in self._bodies]

    def reset(self) -> None:
        """Reset the simulation to its initial state"""
        self._bodies.clear()
        self._positions.clear()
        self._sizes.clear()
        self._colors.clear()
        self.setup_simulation()