gravity_controller.py

"""
This is a module that implements a funny simulation called 'Gravity Controller'
using Python and the Pygame and Pymunk libraries.

In this simulation, the user is allowed to use the computer mouse to change the
magnitude and direction of the gravity of the Pymunk Space. In the center of
the screen, there is a static box with some dynamic balls inside. These balls
will move attracted by the gravity. In other words, the balls will move towards
the mouse position at each time step. Move your mouse to guide the balls.

Pygame is a Python library that serves as a versatile framework for developing
2D games and multimedia applications. It provides a straightforward and
efficient way to manage graphics, sound, user input, and basic collision
detection.

PyMunk is a Python library that extends Pygame by adding robust 2D physics
simulation capabilities, allowing developers to create realistic physics
interactions within their 2D games and simulations. It provides functionality
for creating and managing physics objects such as rigid bodies, shapes,
constraints, and handling complex collision detection and response.

PyMunk's integration with Pygame simplifies the process of combining graphics
with dynamic physics, making it a valuable tool for those seeking to develop
games or simulations with compelling and lifelike physics behavior.

KEYWORDS: Pygame, Pymunk, YAML, Simulations, Physics using Python.
"""

import os
import random
from typing import Any, Dict, Tuple

import pygame
import pymunk
import yaml


class Ball:
    """
    Ball class. Object that is attracted towards the direction of
    gravity generated by the computer mouse (at the computer mouse position).
    Dynamic pymunk body, circular pymunk shape. All balls will be contained in
    a BoundingBox.

    The mass, radius, and color will vary depending on the instance.
    """

    def __init__(
            self,
            center: Tuple[int, int],
            mass: int,
            radius: int,
            color: Tuple[int, int, int],
    ):
        """
        Initialize a Ball instance.

        :param center: (x,y) initial coordinates of the ball
        :param mass: mass of the Ball
        :param radius: radius of the Ball (circle)
        :param color: color in RGB format
        """

        self.color = color  # in RGB format

        # Create the Pymunk Body of the Ball
        self.body = pymunk.Body(mass=mass, moment=1)
        self.body.position = center

        # Create the Pymunk Shape (Circle) of the Ball
        self.shape = pymunk.Circle(self.body, radius=radius)

    def draw(self, screen: pygame.surface) -> None:
        """
        Display the Ball on the given 'screen' (pygame surface)

        :param screen: pygame surface where the Ball is displayed
        :return: None
        """

        pygame.draw.circle(
            surface=screen,
            color=self.color,
            center=self.body.position,
            radius=self.shape.radius
        )


class BoundingBox:
    """
    BoundingBox class. Square and static Box that will serve as a container for
    the balls. The BoundingBox will be placed in the center of the screen.
    The balls within the BoundingBox will be moved in the direction of the
    gravity generated by the computer mouse, but they will not be allowed to
    leave the box.
    """

    def __init__(
            self,
            center: Tuple[int, int],
            width: int,
            height: int,
            radius: int,
            line_color: Tuple[int, int, int],
            background_color: Tuple[int, int, int]
    ) -> None:
        """
        Initialize a BoundingBox instance.

        :param center: (x,y) coordinates of the BoundingBox center
        :param width: width of the BoundingBox
        :param height: height of the BoundingBox
        :param radius: line width (contour)
        :line_color: color of the contour, in RGB format
        :background_color: color that fills the box, in RGB format
        """

        self.center = center
        self.width = width
        self.height = height
        self.radius = radius
        self.line_color = line_color
        self.background_color = background_color

        # Define the (x,y) coordinates of the corners of the BoundingBox
        tl = (center[0] - width//2, center[1] - height//2)  # Top Left
        tr = (center[0] + width//2, center[1] - height//2)  # Top Right
        bl = (center[0] - width//2, center[1] + height//2)  # Bottom Left
        br = (center[0] + width//2, center[1] + height//2)  # Bottom Right

        # Create a body for each of the four Segments (lines that form the Box)
        self.body = (
            pymunk.Body(mass=1, moment=1, body_type=pymunk.Body.STATIC),
            pymunk.Body(mass=1, moment=1, body_type=pymunk.Body.STATIC),
            pymunk.Body(mass=1, moment=1, body_type=pymunk.Body.STATIC),
            pymunk.Body(mass=1, moment=1, body_type=pymunk.Body.STATIC)
        )

        # Create four Segments (shapes) that form the BoundingBox
        self.shape = (
            pymunk.Segment(body=self.body[0], a=tl, b=tr, radius=self.radius),
            pymunk.Segment(body=self.body[1], a=tl, b=bl, radius=self.radius),
            pymunk.Segment(body=self.body[2], a=tr, b=br, radius=self.radius),
            pymunk.Segment(body=self.body[3], a=bl, b=br, radius=self.radius)
        )

        # Define the background as a Rectangle (pygame surface)
        self.background_rect = pygame.Rect(tl[0], tl[1], width, height)
        # order of arguments in pygame.Rect: left, top, width, height

    def draw(self, screen: pygame.surface) -> None:
        """
        Displays the Box on the given 'screen' (pygame surface)

        :return: None
        """

        # First, fill the rectangle (background as a pygame Rect)
        pygame.draw.rect(
            surface=screen,
            color=self.background_color,
            rect=self.background_rect
        )
        # Then, display the four contours (lines)
        for segment in self.shape:
            pygame.draw.line(
                surface=screen,
                color=self.line_color,
                start_pos=segment.a,
                end_pos=segment.b,
                width=self.radius
            )


class GravityController:
    """
    GravityController class. This class manages the logic of the simulation.
    It creates the pymunk space, and initializes a set of Balls inside a
    BoundingBox instance. It also updates the gravity based on the current
    position of the mouse, and displays all the elements on a pygame surface.
    """

    def __init__(self) -> None:
        """
        Initialize a GravityController instance.
        """

        # Read the configuration file
        self._config = self._get_config()

        # Initialize the Pygame elements
        pygame.init()
        self._screen = pygame.display.set_mode(
            (self._config['screen_width'], self._config['screen_height'])
        )
        pygame.display.set_caption(self._config['screen_caption'])
        self._clock = pygame.time.Clock()

        # Create a Pymunk Space
        self.space = pymunk.Space()
        # Note that the gravity will depend on the mouse position

        # Create the Box and fill it with Balls
        screen_center = (
            self._config['screen_width'] // 2,
            self._config['screen_height'] // 2
        )
        self.bounding_box = BoundingBox(
            center=screen_center,
            width=self._config['box_width'],
            height=self._config['box_height'],
            radius=self._config['box_line_width'],
            line_color=self._config['box_line_color'],
            background_color=self._config['box_background_color']
        )
        # Add the box (its pymunk Segments) to the pymunk space
        for body, shape in zip(self.bounding_box.body, self.bounding_box.shape):
            self.space.add(body, shape)

        # Create a set of Balls inside the BoundingBox
        self.balls_list = []  # list of the Ball instances
        for _ in range(self._config['n_balls_inside_box']):
            ball = Ball(
                center=self.bounding_box.center,  # inside the BoundingBox
                mass=random.randrange(
                    start=self._config['min_mass'],
                    stop=self._config['max_mass']
                ),
                radius=random.randrange(
                    start=self._config['min_radius'],
                    stop=self._config['max_radius']
                ),
                color=random.choices(population=range(255), k=3)
            )
            self.balls_list.append(ball)

        # Add the Balls to the pymunk space
        for ball in self.balls_list:
            self.space.add(ball.body, ball.shape)

    @staticmethod
    def _get_config() -> Dict[str, Any]:
        """
        Read the configuration file and return it as a python dictionary.
        The configuration file is called 'gravity_controller/config.yml'

        :return: configuration dictionary
        """

        this_file_path = os.path.abspath(__file__)
        this_project_dir_path = '/'.join(this_file_path.split('/')[:-1])
        config_path = this_project_dir_path + '/config.yml'

        with open(config_path, 'r') as yml_file:
            config = yaml.safe_load(yml_file)[0]['config']
        return config

    @staticmethod
    def process_events() -> bool:
        """
        Process the actions of the user:
            - if user wants to end the simulation: quit pygame

        :return: whether the simulation is running
        """

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

    def run_logic(self) -> None:
        """
        Run the logic of the simulation:
            - Update the positions and velocities of the Ball instances
                accordingly in discrete time step using the 'step' method.
            - Updates the gravity based on the current position of the mouse.

        :return: None
        """

        self.space.step(1/60)

        # The idea is to model the gravity so that the balls (Ball instances)
        # follow the mouse, like if they were chasing the mouse but found
        # the BoundingBox limits on their way. After some testing, this formula
        # achieves the desired effect.
        mouse_x, mouse_y = pygame.mouse.get_pos()
        gravity_x = self._config['screen_width'] / 2 - mouse_x
        gravity_y = self._config['screen_height'] / 2 - mouse_y
        self.space.gravity = -3 * int(gravity_x), -3 * int(gravity_y)

    def draw(self) -> None:
        """
        Display the elements of the simulation on the 'screen' attribute.

        :return: None
        """

        self._screen.fill(self._config['background_color'])
        self.bounding_box.draw(screen=self._screen)
        for ball in self.balls_list:
            ball.draw(screen=self._screen)
        pygame.display.update()

    def clock_tick(self) -> None:
        """
        Updates the pygame clock (attribute '_clock')

        :return: None
        """

        self._clock.tick(self._config['pygame_clock_tick'])


if __name__ == '__main__':
    gravity_controller = GravityController()

    running = True
    while running:
        running = gravity_controller.process_events()
        gravity_controller.run_logic()
        gravity_controller.draw()
        gravity_controller.clock_tick()

    pygame.quit()