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()