275 lines
9.1 KiB
Python
275 lines
9.1 KiB
Python
import pygame
|
|
import math
|
|
import random
|
|
|
|
# Initialize Pygame
|
|
pygame.init()
|
|
|
|
# Screen setup
|
|
WIDTH, HEIGHT = 800, 600
|
|
screen = pygame.display.set_mode((WIDTH, HEIGHT))
|
|
pygame.display.set_caption("balls")
|
|
|
|
# Colors
|
|
BLACK = (0, 0, 0)
|
|
WHITE = (255, 255, 255)
|
|
RED = (220, 50, 50)
|
|
GREEN = (50, 220, 50)
|
|
CYAN = (50, 220, 220)
|
|
YELLOW = (220, 220, 50)
|
|
|
|
# Clock
|
|
clock = pygame.time.Clock()
|
|
FPS = 60
|
|
|
|
# Global physics
|
|
gravity = 0.5
|
|
elasticity = 0.85
|
|
air_resistance = 0.99 # 1.0 = no resistance, lower = more resistance
|
|
friction = 0.98 # Ground friction when rolling
|
|
|
|
# Ball class
|
|
class Ball:
|
|
def __init__(self, x, y):
|
|
self.radius = 20
|
|
self.x = x
|
|
self.y = y
|
|
self.mass = self.radius / 10.0 # Mass proportional to radius
|
|
# Give random initial velocity in both x and y
|
|
self.vx = random.uniform(-5, 5)
|
|
self.vy = random.uniform(-5, 0)
|
|
self.tracer_length = 50
|
|
self.tracers = []
|
|
self.dragging = False
|
|
self.drag_start_pos = None
|
|
self.drag_last_pos = (x, y)
|
|
self.drag_history = []
|
|
|
|
def update(self):
|
|
if not self.dragging:
|
|
# Apply gravity
|
|
self.vy += gravity
|
|
|
|
# Apply air resistance (quadratic drag)
|
|
speed = math.hypot(self.vx, self.vy)
|
|
if speed > 0:
|
|
drag_force = (1 - air_resistance) * speed
|
|
drag_x = -drag_force * (self.vx / speed)
|
|
drag_y = -drag_force * (self.vy / speed)
|
|
self.vx += drag_x
|
|
self.vy += drag_y
|
|
|
|
# Update position
|
|
self.x += self.vx
|
|
self.y += self.vy
|
|
|
|
# Bounce off floor with friction
|
|
if self.y + self.radius > HEIGHT:
|
|
self.y = HEIGHT - self.radius
|
|
self.vy = -self.vy * elasticity
|
|
self.vx *= friction # Apply friction when bouncing on ground
|
|
|
|
# Stop if moving very slowly
|
|
if abs(self.vy) < 0.1:
|
|
self.vy = 0
|
|
if abs(self.vx) < 0.1:
|
|
self.vx = 0
|
|
|
|
# Bounce off ceiling
|
|
if self.y - self.radius < 0:
|
|
self.y = self.radius
|
|
self.vy = -self.vy * elasticity
|
|
|
|
# Bounce off right wall
|
|
if self.x + self.radius > WIDTH:
|
|
self.x = WIDTH - self.radius
|
|
self.vx = -self.vx * elasticity
|
|
|
|
# Bounce off left wall
|
|
if self.x - self.radius < 0:
|
|
self.x = self.radius
|
|
self.vx = -self.vx * elasticity
|
|
|
|
# Update tracers
|
|
self.tracers.append((int(self.x), int(self.y)))
|
|
if len(self.tracers) > self.tracer_length:
|
|
self.tracers.pop(0)
|
|
|
|
def start_drag(self, mx, my):
|
|
self.dragging = True
|
|
self.drag_start_pos = (mx, my)
|
|
self.drag_last_pos = (mx, my)
|
|
self.drag_history = [(mx, my)]
|
|
self.vx = 0
|
|
self.vy = 0
|
|
|
|
def update_drag(self, mx, my):
|
|
if self.dragging:
|
|
self.x, self.y = mx, my
|
|
self.drag_history.append((mx, my))
|
|
if len(self.drag_history) > 5: # Keep last 5 positions for velocity calculation
|
|
self.drag_history.pop(0)
|
|
self.drag_last_pos = (mx, my)
|
|
|
|
def end_drag(self):
|
|
if self.dragging and len(self.drag_history) >= 2:
|
|
# Calculate velocity from drag motion (throw mechanics)
|
|
dx = self.drag_history[-1][0] - self.drag_history[0][0]
|
|
dy = self.drag_history[-1][1] - self.drag_history[0][1]
|
|
frames = len(self.drag_history)
|
|
|
|
# Set velocity based on drag motion
|
|
self.vx = dx / frames * 2
|
|
self.vy = dy / frames * 2
|
|
|
|
self.dragging = False
|
|
self.drag_start_pos = None
|
|
self.drag_history = []
|
|
|
|
def draw(self, surface):
|
|
# Draw tracers with fade effect
|
|
for i, pos in enumerate(self.tracers):
|
|
alpha = int(255 * (i / len(self.tracers)))
|
|
size = int(3 + 2 * (i / len(self.tracers)))
|
|
color = (50, 220, 220)
|
|
pygame.draw.circle(surface, color, pos, size)
|
|
|
|
# Draw ball
|
|
pygame.draw.circle(surface, RED, (int(self.x), int(self.y)), self.radius)
|
|
|
|
# Draw velocity vector (scaled for visibility)
|
|
if not self.dragging:
|
|
vel_scale = 3
|
|
pygame.draw.line(surface, GREEN,
|
|
(int(self.x), int(self.y)),
|
|
(int(self.x + self.vx * vel_scale),
|
|
int(self.y + self.vy * vel_scale)), 3)
|
|
|
|
# Draw drag trajectory preview
|
|
if self.dragging and self.drag_start_pos:
|
|
pygame.draw.line(surface, YELLOW,
|
|
self.drag_start_pos,
|
|
(int(self.x), int(self.y)), 2)
|
|
# Draw arrow at the end
|
|
pygame.draw.circle(surface, YELLOW, (int(self.x), int(self.y)), 5)
|
|
|
|
# Ball-to-ball collision detection and response
|
|
def handle_ball_collisions(balls):
|
|
for i, ball1 in enumerate(balls):
|
|
for ball2 in balls[i+1:]:
|
|
dx = ball2.x - ball1.x
|
|
dy = ball2.y - ball1.y
|
|
distance = math.hypot(dx, dy)
|
|
|
|
# Check if balls are colliding
|
|
if distance < ball1.radius + ball2.radius and distance > 0:
|
|
# Normalize collision vector
|
|
nx = dx / distance
|
|
ny = dy / distance
|
|
|
|
# Separate balls
|
|
overlap = (ball1.radius + ball2.radius) - distance
|
|
ball1.x -= nx * overlap * 0.5
|
|
ball1.y -= ny * overlap * 0.5
|
|
ball2.x += nx * overlap * 0.5
|
|
ball2.y += ny * overlap * 0.5
|
|
|
|
# Calculate relative velocity
|
|
dvx = ball2.vx - ball1.vx
|
|
dvy = ball2.vy - ball1.vy
|
|
|
|
# Calculate relative velocity in collision normal direction
|
|
dot_product = dvx * nx + dvy * ny
|
|
|
|
# Only resolve if balls are moving toward each other
|
|
if dot_product < 0:
|
|
# Calculate impulse (conservation of momentum with equal mass)
|
|
impulse = 2 * dot_product / (ball1.mass + ball2.mass)
|
|
|
|
# Apply impulse to both balls
|
|
ball1.vx += impulse * ball2.mass * nx * elasticity
|
|
ball1.vy += impulse * ball2.mass * ny * elasticity
|
|
ball2.vx -= impulse * ball1.mass * nx * elasticity
|
|
ball2.vy -= impulse * ball1.mass * ny * elasticity
|
|
|
|
# List of balls
|
|
balls = [Ball(WIDTH//2, HEIGHT//2)]
|
|
|
|
running = True
|
|
while running:
|
|
mx, my = pygame.mouse.get_pos()
|
|
|
|
for event in pygame.event.get():
|
|
if event.type == pygame.QUIT:
|
|
running = False
|
|
|
|
# Drag start
|
|
elif event.type == pygame.MOUSEBUTTONDOWN:
|
|
for ball in balls:
|
|
if math.hypot(mx - ball.x, my - ball.y) < ball.radius + 10:
|
|
ball.start_drag(mx, my)
|
|
break
|
|
|
|
# Drag end
|
|
elif event.type == pygame.MOUSEBUTTONUP:
|
|
for ball in balls:
|
|
if ball.dragging:
|
|
ball.end_drag()
|
|
|
|
# Spawn new ball
|
|
elif event.type == pygame.KEYDOWN:
|
|
if event.key == pygame.K_SPACE:
|
|
balls.append(Ball(random.randint(50, WIDTH-50),
|
|
random.randint(50, HEIGHT-50)))
|
|
elif event.key == pygame.K_c:
|
|
balls.clear()
|
|
balls.append(Ball(WIDTH//2, HEIGHT//2))
|
|
|
|
# Keyboard control for physics
|
|
keys = pygame.key.get_pressed()
|
|
if keys[pygame.K_UP]:
|
|
gravity = min(2.0, gravity + 0.01)
|
|
if keys[pygame.K_DOWN]:
|
|
gravity = max(0, gravity - 0.01)
|
|
if keys[pygame.K_RIGHT]:
|
|
elasticity = min(1.0, elasticity + 0.01)
|
|
if keys[pygame.K_LEFT]:
|
|
elasticity = max(0, elasticity - 0.01)
|
|
if keys[pygame.K_w]:
|
|
air_resistance = min(1.0, air_resistance + 0.001)
|
|
if keys[pygame.K_s]:
|
|
air_resistance = max(0.9, air_resistance - 0.001)
|
|
|
|
# Update dragging balls
|
|
for ball in balls:
|
|
if ball.dragging:
|
|
ball.update_drag(mx, my)
|
|
|
|
# Update balls
|
|
for ball in balls:
|
|
ball.update()
|
|
|
|
# Handle collisions between balls
|
|
handle_ball_collisions(balls)
|
|
|
|
# Drawing
|
|
screen.fill(BLACK)
|
|
for ball in balls:
|
|
ball.draw(screen)
|
|
|
|
# Physics info
|
|
font = pygame.font.SysFont(None, 22)
|
|
# info = [
|
|
# f"Gravity: {gravity:.2f} Elasticity: {elasticity:.2f} (←→)",
|
|
# f"Air Resistance: {air_resistance:.3f} (WS) Balls: {len(balls)}",
|
|
# f"SPACE: Add ball C: Clear Drag to throw!"
|
|
# ]
|
|
|
|
# for i, text in enumerate(info):
|
|
# surface = font.render(text, True, WHITE)
|
|
# screen.blit(surface, (10, 10 + i * 25))
|
|
|
|
pygame.display.flip()
|
|
clock.tick(FPS)
|
|
|
|
pygame.quit() |