Pool Ball Movement Physics

gistfile1.txt

import pygame
import math

# Initialize Pygame
pygame.init()

# Set up display
WIDTH, HEIGHT = 800, 400
screen = pygame.display.set_mode((WIDTH, HEIGHT))
pygame.display.set_caption("Pool Ball Simulator")

# Colors
WHITE = (255, 255, 255)
GREEN = (0, 128, 0)
RED = (255, 0, 0)
BLUE = (0, 0, 255)

# Physics Constants
FRICTION = 0.99
BALL_RADIUS = 10
BALL_MASS = 1  # Assuming equal mass for simplicity

# Ball Class
class Ball:
    def __init__(self, x, y, vx, vy, color):
        self.x = x
        self.y = y
        self.vx = vx
        self.vy = vy
        self.radius = BALL_RADIUS
        self.color = color

    def move(self):
        # Update position based on velocity
        self.x += self.vx
        self.y += self.vy

        # Apply friction
        self.vx *= FRICTION
        self.vy *= FRICTION

        # Handle collision with the table's edges
        if self.x - self.radius <= 0 or self.x + self.radius >= WIDTH:
            self.vx = -self.vx
        if self.y - self.radius <= 0 or self.y + self.radius >= HEIGHT:
            self.vy = -self.vy

    def draw(self, screen):
        pygame.draw.circle(screen, self.color, (int(self.x), int(self.y)), self.radius)

    def collide(self, other):
        # Calculate the distance between two balls
        dx = other.x - self.x
        dy = other.y - self.y
        distance = math.sqrt(dx**2 + dy**2)

        # Check if the distance is less than the sum of the radii (collision)
        if distance < self.radius + other.radius:
            # Calculate angle of collision
            angle = math.atan2(dy, dx)

            # Total velocity before collision
            u1 = math.sqrt(self.vx**2 + self.vy**2)
            u2 = math.sqrt(other.vx**2 + other.vy**2)

            # Direction of velocities before collision
            dir1 = math.atan2(self.vy, self.vx)
            dir2 = math.atan2(other.vy, other.vx)

            # New velocities after collision
            self.vx = u2 * math.cos(dir2 - angle)
            self.vy = u2 * math.sin(dir2 - angle)
            other.vx = u1 * math.cos(dir1 - angle)
            other.vy = u1 * math.sin(dir1 - angle)

            # Reposition balls to prevent overlap
            overlap = 0.5 * (self.radius + other.radius - distance + 1)
            self.x -= math.cos(angle) * overlap
            self.y -= math.sin(angle) * overlap
            other.x += math.cos(angle) * overlap
            other.y += math.sin(angle) * overlap

# Initialize balls
balls = [
    Ball(WIDTH // 3, HEIGHT // 2, 5, 2, RED),
    Ball(2 * WIDTH // 3, HEIGHT // 2, -4, -3, BLUE),
    Ball(WIDTH // 2, HEIGHT // 3, 2, -5, WHITE),
]

# Main game loop
running = True
clock = pygame.time.Clock()

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

    # Move and check collisions between all balls
    for i, ball in enumerate(balls):
        ball.move()
        for other in balls[i + 1:]:
            ball.collide(other)

    # Draw everything
    screen.fill(GREEN)
    for ball in balls:
        ball.draw(screen)
    pygame.display.flip()

    # Cap the frame rate
    clock.tick(60)

pygame.quit()