original source

import abc
from typing import Tuple

import numpy as np
import pygame
from pygame import Color

pygame.init()

screen_width = 320
screen_height = 320

background_color = (255, 255, 255)  # white

screen_size = (screen_width, screen_height)

clock = pygame.time.Clock()

screen = pygame.display.set_mode(screen_size)
pygame.display.set_caption("Collision at boundary")


class Vector2D:
    @staticmethod
    def from_tuple(vector):
        return Vector2D(vector[0], vector[1])

    def __init__(self, x=0.0, y=0.0):
        self.x = x
        self.y = y

    def direction(self):
        return np.arctan(self.y / self.x)

    def to_tuple(self):
        return self.x, self.y

    def norm(self):
        return (self.x ** 2 + self.y ** 2) ** 0.5

    def normalize(self):
        self.x /= self.norm()
        self.y /= self.norm()

    def normalized(self):
        return Vector2D(self.x / self.norm(), self.y / self.norm())

    def __add__(self, other):
        return Vector2D(self.x + other.x, self.y + other.y)

    def __sub__(self, other):
        return Vector2D(self.x - other.x, self.y - other.y)

    def __mul__(self, other):
        if isinstance(other, float) or isinstance(other, int):
            return Vector2D(self.x * other, self.y * other)
        elif isinstance(other, Vector2D):
            return self.x * other.x + self.y * other.y

    def __rmul__(self, other):
        if isinstance(other, float) or isinstance(other, int):
            return Vector2D(self.x * other, self.y * other)

    def __truediv__(self, other):
        if isinstance(other, float) or isinstance(other, int):
            return Vector2D(self.x / other, self.y / other)

    def __str__(self):
        return "Vector2D: " + str(self.to_tuple())


class Shape(metaclass=abc.ABCMeta):
    def __init__(
            self,
            cur_screen: pygame.surface.Surface,
            coordinates: Tuple[int, int],
            initial_v: Tuple[float, float],
            mass: int,
            color: Color
    ):
        self.coors = Vector2D.from_tuple(coordinates)
        self.velocity = Vector2D.from_tuple(initial_v)
        self.mass = mass
        self.screen = cur_screen
        self.color = color

        self.hit_flag = False

    @abc.abstractmethod
    def draw(self):
        pass

    def move(self):
        self.coors += self.velocity


class Circle(Shape):
    def __init__(
            self,
            cur_screen: pygame.surface.Surface,
            coordinates: Tuple[int, int],
            initial_v: Tuple[float, float],
            mass: int,
            color: Color,
            radius: int,
    ):
        super().__init__(cur_screen, coordinates, initial_v, mass, color)
        self.radius = radius

    def draw(self):
        pygame.draw.circle(self.screen, self.color, self.coors.to_tuple(), self.radius)

    def detect_wall_collision(self):
        boundary = self.screen.get_size()
        if self.coors.x - self.radius <= 0 or self.coors.x + self.radius >= boundary[0]:
            self.velocity = Vector2D(-self.velocity.x, self.velocity.y)
        if self.coors.y - self.radius <= 0 or self.coors.y + self.radius >= boundary[1]:
            self.velocity = Vector2D(self.velocity.x, -self.velocity.y)


def detect_collision(body1: Circle, body2: Circle):
    distance_v = body1.coors - body2.coors
    rad_sum = body1.radius + body2.radius
    if distance_v.norm() <= rad_sum:
        return True
    return False


def collision(body1: Shape, body2: Shape):
    m1 = body1.mass
    m2 = body2.mass
    v1_i = body1.velocity
    v2_i = body2.velocity

    d = body1.coors - body2.coors

    t = ((v2_i - v1_i) * d) * 2 / ((d * d) * (1 / m1) + (d * d) * (1 / m2))
    I: Vector2D | float = d * t
    v1_f = v1_i + I * (1 / m1)
    v2_f = v2_i - I * (1 / m2)

    body1.velocity = v1_f
    body2.velocity = v2_f


def update_pos(circles: list):
    for i in range(len(circles)):
        c1 = circles[i]
        for j in range(i + 1, len(circles)):
            c2 = circles[j]
            if detect_collision(c1, c2):
                collision(c1, c2)

        c1.detect_wall_collision()
        c1.move()
        c1.draw()


circle1 = Circle(screen, (160, 80), (0.5, 1.0), 10, Color("blue"), 30)
circle2 = Circle(screen, (80, 240), (-0.6, 0.3), 10, Color("green"), 30)
circle3 = Circle(screen, (240, 240), (-0.8, -0.8), 10, Color("red"), 30)

circles = [circle1, circle2, circle3]

running = True


while running:
    clock.tick(120)
    screen.fill(background_color)

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

    update_pos(circles)

    pygame.display.flip()