pygame2-claude-3.7.py

import pygame
import random
import math

# Initialize Pygame
pygame.init()

# Screen dimensions
WIDTH, HEIGHT = 800, 600
screen = pygame.display.set_mode((WIDTH, HEIGHT))
pygame.display.set_caption("2D Physics Engine")

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

# Physics constants
GRAVITY = 0.5
ELASTICITY = 0.75  # Coefficient of restitution (0 = inelastic, 1 = elastic)

class Ball:
    def __init__(self, x, y, radius, mass, color):
        self.x = x
        self.y = y
        self.radius = radius
        self.mass = mass
        self.color = color
        self.vx = random.uniform(-5, 5)  # Initial horizontal velocity
        self.vy = random.uniform(-5, 5)  # Initial vertical velocity

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

    def update(self):
        # Apply gravity
        self.vy += GRAVITY

        # Update position
        self.x += self.vx
        self.y += self.vy

        # Boundary collisions (bounce off the walls and floor)
        if self.x + self.radius > WIDTH or self.x - self.radius < 0:
            self.vx *= -ELASTICITY
            self.x = min(WIDTH - self.radius, max(self.radius, self.x))  # Keep ball within bounds
        if self.y + self.radius > HEIGHT or self.y - self.radius < 0:
            self.vy *= -ELASTICITY
            self.y = min(HEIGHT - self.radius, max(self.radius, self.y))  # Keep ball within bounds

def collide(ball1, ball2):
    # Calculate the distance between the centers of the balls
    dx = ball2.x - ball1.x
    dy = ball2.y - ball1.y
    distance = math.sqrt(dx**2 + dy**2)

    # Check for collision
    if distance < ball1.radius + ball2.radius:
        # Calculate the angle of collision
        angle = math.atan2(dy, dx)

        # Calculate the velocities along the line of collision
        v1 = ball1.vx * math.cos(angle) + ball1.vy * math.sin(angle)
        v2 = ball2.vx * math.cos(angle) + ball2.vy * math.sin(angle)

        # Calculate the new velocities after the collision (1D collision formula)
        v1_new = (v1 * (ball1.mass - ball2.mass) + 2 * ball2.mass * v2) / (ball1.mass + ball2.mass)
        v2_new = (v2 * (ball2.mass - ball1.mass) + 2 * ball1.mass * v1) / (ball1.mass + ball2.mass)

        # Update the velocities of the balls
        ball1.vx = v1_new * math.cos(angle) - ball1.vy * math.sin(angle)
        ball1.vy = v1_new * math.sin(angle) + ball1.vy * math.cos(angle)
        ball2.vx = v2_new * math.cos(angle) - ball2.vy * math.sin(angle)
        ball2.vy = v2_new * math.sin(angle) + ball2.vy * math.cos(angle)

# Create multiple balls
balls = [
    Ball(random.randint(50, 750), random.randint(50, 200), 20, 1, RED),
    Ball(random.randint(50, 750), random.randint(50, 200), 30, 2, GREEN),
    Ball(random.randint(50, 750), random.randint(50, 200), 25, 1.5, BLUE)
]

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

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

    # Update ball positions
    for ball in balls:
        ball.update()

    # Check for collisions between balls
    for i in range(len(balls)):
        for j in range(i + 1, len(balls)):
            collide(balls[i], balls[j])

    # Clear the screen
    screen.fill(BLACK)

    # Draw the balls
    for ball in balls:
        ball.draw(screen)

    # Update the display
    pygame.display.flip()

    # Limit frame rate
    clock.tick(60)

# Quit Pygame
pygame.quit()