rigid-body/main.py

import pygame as pg
from abc import ABC, abstractmethod
from dataclasses import dataclass
from itertools import combinations
from collections import namedtuple
from random import random

GLOBAL_GRAVITY=250
DIMENSIONS = namedtuple("Dimensions", ["WIDTH", "HEIGHT"])(500,500)

@dataclass
class ColliderContact:
    __slots__ = ["point", "normal", "penetration"]
    point: pg.Vector2
    normal: pg.Vector2
    penetration: float

@dataclass
class RigidBody:
    position: pg.Vector2
    collider: "BaseCollider"
    velocity: pg.Vector2
    mass: float = 1.0
    restitution: float = 0.5

    @property
    def inv_mass(self):
        return 0.0 if self.mass == 0 else 1.0 / self.mass

class BaseCollider(ABC):
     
    def collide(self, other: "BaseCollider", this_position: pg.Vector2, other_position: pg.Vector2) -> ColliderContact | None:
        if isinstance(other, CircleCollider):
            return self.collide_circle(other, this_position, other_position)
        

    @abstractmethod
    def collide_circle(self, other: "CircleCollider", this_position: pg.Vector2, other_position: pg.Vector2) -> ColliderContact | None:
        pass

class CircleCollider(BaseCollider):

    def __init__(self, radius: float):
        self.radius = radius

    def collide_circle(self, other: "CircleCollider", this_position: pg.Vector2, other_position: pg.Vector2) -> ColliderContact | None:
        delta = this_position - other_position
        dist = delta.length()
        radii = self.radius + other.radius
        if dist >= radii or dist == 0:
            return None
        normal = delta.normalize()
        return ColliderContact(
            point=other_position + normal * other.radius,
            normal=normal,
            penetration=radii - dist,
        )  

class PhysicsSystem:
    
    def __init__(self):
        self.bodies: list[RigidBody] = []
        self.gravity = pg.Vector2(0,GLOBAL_GRAVITY)
    
    def add_body(self, body: RigidBody) -> None:
        self.bodies.append(body)

    def update(self, dt: float):
        g = self.gravity * dt
        for body in self.bodies:
            body.velocity += g
            body.position += dt * body.velocity
            self.resolve_bounds(body)

        for a, b in combinations(self.bodies, 2):
            if collision := a.collider.collide(b.collider, a.position, b.position):
                self.resolve_collision(a, b, collision)
        
    def resolve_collision(self, a: RigidBody, b: RigidBody, collision: ColliderContact) -> None:
        SLACK=0.4
        correction = collision.penetration / (a.inv_mass + b.inv_mass) * SLACK * collision.normal
        a.position += correction
        b.position -= correction

        v_rel = a.velocity - b.velocity
        restitution = a.restitution * b.restitution
        impulse = (-(1 + restitution) * (collision.normal.dot(v_rel))) / (a.inv_mass + b.inv_mass)
        a.velocity += collision.normal*impulse*a.inv_mass
        b.velocity -= collision.normal*impulse*b.inv_mass

    def resolve_bounds(self, body: RigidBody) -> None:
        r = body.collider.radius

        if body.position.x - r < 0:
            body.position.x = r
            body.velocity.x = abs(body.velocity.x) * body.restitution
        elif body.position.x + r > DIMENSIONS.WIDTH:
            body.position.x = DIMENSIONS.WIDTH - r
            body.velocity.x = -abs(body.velocity.x) * body.restitution

        if body.position.y - r < 0:
            body.position.y = r
            body.velocity.y = abs(body.velocity.y) * body.restitution
        elif body.position.y + r > DIMENSIONS.HEIGHT:
            body.position.y = DIMENSIONS.HEIGHT - r
            body.velocity.y = -abs(body.velocity.y) * body.restitution

class Circle:

    def __init__(self, radius: float, body: RigidBody, sprite: pg.Surface):
        self.radius = radius
        self.body = body
        self.sprite = sprite

    def draw(self, screen: pg.Surface):
        screen.blit(
            self.sprite,
            (
                self.body.position.x - self.radius,
                self.body.position.y - self.radius
            )
        )

class World:

    def __init__(self, screen: pg.Surface):
        self.physics = PhysicsSystem()
        self.screen = screen
        self.entities = []

    def update(self, dt) -> None:
        self.physics.update(dt)
        
    def draw(self) -> None:
        for entity in self.entities:
            entity.draw(self.screen)

    def spawn_circle(self, radius: float, position: pg.Vector2, velocity: pg.Vector2 | None = None, restitution:float = 0.5, mass: float = 1.0) -> None:
        if velocity is None:
            velocity = pg.Vector2(0,0)
        collider = CircleCollider(radius)
        body = RigidBody(position=position, collider=collider, velocity=velocity, restitution=restitution,mass=mass)
        self.physics.add_body(body)
        sprite = pg.Surface((2*radius, 2*radius),pg.SRCALPHA)
        sprite.fill((255,255,255,0))
        pg.draw.circle(
            sprite,
            color=(0,255,0,255),
            center=(radius, radius),
            radius=radius
        )
        self.entities.append(Circle(radius,body,sprite))

def main():
    running=True
    pg.init()
    clock=pg.time.Clock()
    screen = pg.display.set_mode(size=DIMENSIONS)
    world=World(screen)

    while running:
        dt = clock.tick(144) / 1000
        screen.fill((0,0,0,0))
        world.update(dt)
        world.draw()
        pg.display.flip()
        for event in pg.event.get():
            if event.type == pg.MOUSEBUTTONDOWN:
                world.spawn_circle(
                    10,
                    pg.Vector2(pg.mouse.get_pos()),
                    restitution=random()
                )
            if event.type == pg.QUIT:
                running = False
        
if __name__ == "__main__":
    main()