import pygame
import math
import random

# Initialize
pygame.init()
WIDTH, HEIGHT = 1000, 1000
screen = pygame.display.set_mode((WIDTH, HEIGHT))
pygame.display.set_caption("Ball Bounce in a Circle")
clock = pygame.time.Clock()
font = pygame.font.SysFont(None, 36)

# Constants
CENTER = (WIDTH // 2, HEIGHT // 2)
CIRCLE_RADIUS = 450
START_RADIUS = 20
MIN_RADIUS = 5
BALL_SPEED = 2
SPAWN_COOLDOWN = 30

# Colors
BLACK = (0, 0, 0)
WHITE = (255, 255, 255)
GRAY = (150, 150, 150)

class Ball:
    def __init__(self, x=CENTER[0], y=CENTER[1], radius=START_RADIUS):
        angle = random.uniform(0, 2 * math.pi)
        self.x = x
        self.y = y
        self.dx = math.cos(angle) * BALL_SPEED
        self.dy = math.sin(angle) * BALL_SPEED
        self.radius = radius
        self.spawn_cooldown = 0

    def update(self):
        self.x += self.dx
        self.y += self.dy
        if self.spawn_cooldown > 0:
            self.spawn_cooldown -= 1

        dist = math.hypot(self.x - CENTER[0], self.y - CENTER[1])
        if dist + self.radius >= CIRCLE_RADIUS:
            nx = (self.x - CENTER[0]) / dist
            ny = (self.y - CENTER[1]) / dist
            dot = self.dx * nx + self.dy * ny
            self.dx -= 2 * dot * nx
            self.dy -= 2 * dot * ny

            if self.spawn_cooldown <= 0 and self.radius > MIN_RADIUS:
                self.radius *= 0.9
                self.spawn_cooldown = SPAWN_COOLDOWN
                return True  # Signal to spawn new balls
        return False

    def draw(self, surface):
        pygame.draw.circle(surface, WHITE, (int(self.x), int(self.y)), int(self.radius))

    def collide_with(self, other):
        dx = other.x - self.x
        dy = other.y - self.y
        distance = math.hypot(dx, dy)
        if distance < self.radius + other.radius:
            angle = math.atan2(dy, dx)

            # Elastic collision velocities exchange
            v1 = (self.dx, self.dy)
            v2 = (other.dx, other.dy)
            self.dx, self.dy = v2
            other.dx, other.dy = v1

            # Separate balls
            overlap = self.radius + other.radius - distance
            shift_x = math.cos(angle) * (overlap / 2)
            shift_y = math.sin(angle) * (overlap / 2)
            self.x -= shift_x
            self.y -= shift_y
            other.x += shift_x
            other.y += shift_y

balls = [Ball()]

running = True
while running:
    screen.fill(BLACK)
    pygame.draw.circle(screen, GRAY, CENTER, CIRCLE_RADIUS, 2)

    new_balls = []
    for i, ball in enumerate(balls):
        if ball.update():
            new_radius = ball.radius * 0.9
            if new_radius > MIN_RADIUS:
                # Spawn 10 smaller balls at center
                for _ in range(10):
                    new_balls.append(Ball(radius=new_radius))

        for j in range(i + 1, len(balls)):
            ball.collide_with(balls[j])

        ball.draw(screen)

    balls.extend(new_balls)

    # Ball counter
    text = font.render(f"Balls: {len(balls)}", True, WHITE)
    screen.blit(text, (20, 20))

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

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

pygame.quit()