# Ian Binstock
# 6/11/2026
# Edit of the bounceballsplit_v2.py example using sprites instead

import random, pygame, sys
from pygame.locals import *

# Colors taken from dvd.py, modified slightly for more vibrant colors
def changeColor():
    color = pygame.Color(0, 0, 0)

    nice = False
    while not nice:
        r = random.randint(50, 255)
        g = random.randint(50, 255)
        b = random.randint(50, 255)

        if abs(r-g) > 175 or abs(r-b) > 175 or abs(g-b) > 175:
            sum = r + g + b
            # <400 = RGB, >525 = CYMK
            if sum < 400 or sum > 525:
                nice = True
                color = pygame.Color(r, g, b)

    return color

# We'll assign colors to each ball sequentially from this list with wraparound
COLORLIST = []
for i in range(10):
    COLORLIST.append(changeColor())

# Useful Constants
SCREEN_W = 1000
SCREEN_H = 900
BALL_R = 10
NUM_BALLS = 200

# Returns a random integer in between low and high not equal to 0
def myrand(low,high):
    res = 0
    while res == 0:
        res = random.randint(low, high)
    return res

class BallSprite(pygame.sprite.Sprite):
    
    def __init__(self, x, y, dx, dy, radius, color):
        super().__init__()
        self.radius = radius
        self.color = color
        self.dx = dx
        self.dy = dy
        
        # 'image' surface required for sprites
        self.image = pygame.Surface((radius * 2, radius * 2), pygame.SRCALPHA)
        self.draw_ball()
        # 'rect' required for sprite positioning
        self.rect = self.image.get_rect(center=(x, y))

    def draw_ball(self):
        self.image.fill((0, 0, 0, 0))  # Clear image before drawing
        pygame.draw.circle(self.image, self.color, (self.radius, self.radius), self.radius)

    def update(self):
        self.rect.x += self.dx
        self.rect.y += self.dy

        # Bounce off left border
        if self.rect.left <= 0:
            self.rect.left = 0
            self.dx *= -1
        # Bounce off right border
        elif self.rect.right >= SCREEN_W:
            self.rect.right = SCREEN_W
            self.dx *= -1

        # Bounce off top border
        if self.rect.top <= 0:
            self.rect.top = 0
            self.dy *= -1
        # Bounce off bottom border
        elif self.rect.bottom >= SCREEN_H:
            self.rect.bottom = SCREEN_H
            self.dy *= -1

def main():

    # Basic Set Up
    pygame.init()
    DISPLAYSURF = pygame.display.set_mode((SCREEN_W, SCREEN_H))
    pygame.display.set_caption("Object Oriented Bouncing w/ Sprites")

    clock = pygame.time.Clock()

    # Make NUM_BALLS random BallSprites
    ball_group = pygame.sprite.Group()
    for i in range(NUM_BALLS):

        # Somewhere on the screen
        x = random.randint(25, SCREEN_W-BALL_R - 25)
        y = random.randint(25, SCREEN_H-BALL_R - 25)

        # Random non-zero movement in both directions
        dx = myrand(-2,2)
        dy = myrand(-2,2)

        # Add ball to group
        ball = BallSprite(x,y,dx,dy,BALL_R,COLORLIST[i%len(COLORLIST)])
        ball_group.add(ball)
        
    # Game loop
    while True:
        for event in pygame.event.get():
            if event.type == QUIT:
                pygame.quit()
                sys.exit()

        # Black bg
        DISPLAYSURF.fill(pygame.Color("black"))

        # Update every ball
        ball_group.update()
            
        # Update what we put on the canvas.
        ball_group.draw(DISPLAYSURF)

        # Collision Handling
        balls = ball_group.sprites()
        for i in range(len(balls)):
            for j in range(i+1, len(balls)):
                if pygame.sprite.collide_circle(balls[i], balls[j]):

                    # Switch dx, dy.
                    tempDX = balls[i].dx
                    tempDY = balls[i].dy
                    balls[i].dx = balls[j].dx
                    balls[i].dy = balls[j].dy
                    balls[j].dx = tempDX
                    balls[j].dy = tempDY

                    # Swap colors too
                    tempColor = balls[i].color
                    balls[i].color = balls[j].color
                    balls[j].color = tempColor

                    # Redraw to update colors
                    balls[i].draw_ball()
                    balls[j].draw_ball()

        pygame.display.update()
        # Wait a bit!
        clock.tick(100)

# Run it
main()