complex balls.py

import pygame
import random
import math

pygame.init()
screen = pygame.display.set_mode((800,800))
clock = pygame.time.Clock()
pygame.event.set_allowed([])

white = (255,255,255)
black = (0,0,0)

n = 20
vel = 5
radius = 20
randcolor = False

balls = []
for i in range(n):
    z = (random.randrange(radius,800-radius) + random.randrange(radius,800-radius) * 1j)
    t = random.random() * math.tau
    v = vel * (math.cos(t) + math.sin(t) * 1j)
    if randcolor:
        color = (random.randrange(128,256),random.randrange(128,256),random.randrange(128,256))
    else: 
        color = white
    balls.append([z,v,color])

balls = tuple(balls)

done = False
while not done:
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            done = True

    for ball in balls:
        ball[0] += ball[1]

        if ball[0].real < radius:
            ball[1] = abs(ball[1].real) + ball[1].imag * 1j
        elif ball[0].real > 800-radius:
            ball[1] = -abs(ball[1].real) + ball[1].imag * 1j
        if ball[0].imag < radius:
            ball[1] = ball[1].real + abs(ball[1].imag) * 1j
        elif ball[0].imag > 800-radius:
            ball[1] = ball[1].real - abs(ball[1].imag) * 1j
    
    for ball1 in balls:
        for ball2 in balls:
            if ball1 != ball2:
                z1, v1, _ = ball1
                z2, _, _ = ball2
                if abs(z1-z2) < 2 * radius:
                    v1 = - (z1-z2)**2 / v1 * vel**2 / abs(z1-z2)**2
                    ball1[1] = v1
    
    screen.fill(black)
    
    for ball in balls:
        pygame.draw.circle(screen, ball[2], (ball[0].real,ball[0].imag), radius)

    clock.tick(60)
    pygame.display.flip()