import pygame
import time
import sys
import random

WIDTH, HEIGHT = 1000, 600
WHITE = (255, 255, 255)
BLUE = (0, 100, 200)
BLACK = (0, 0, 0)
PINK = (255, 0, 120)

canvas = pygame.display.set_mode((WIDTH,HEIGHT))
pygame.display.set_caption('Balls')
clock = pygame.time.Clock()
pygame.init()

def quitGame():
    pygame.quit()
    sys.exit(0)

def randomizeBall():
        x = random.randint(60, WIDTH - 60)
        y = random.randint(60, HEIGHT - 60)
        w = random.randint(MIN_BALL_WIDTH, MAX_BALL_WIDTH)
        xspeed = yspeed =  random.choice(SPEED)
        fun = random.choice(FUN)

        return x, y, w, xspeed, yspeed, fun

class Ball():
    def __init__(self, x, y, w, xspeed, yspeed, fun):
        self.x = x
        self.y = y
        self.w = w
        self.xspeed = xspeed
        self.yspeed = yspeed
        self.fun = fun
        self.rec = pygame.Rect(self.x, self.y, self.w, self.w)

    def update(self, borders):
        self.rec.x += self.xspeed
        self.rec.y += self.yspeed

        for border, value in borders.items():
            if self.rec.colliderect(value.rec):
                if border == 'left' or border == 'right':
                    self.xspeed *= -1
                    self.xspeed += random.randint(-self.fun, self.fun)
                if border == 'up' or border == 'down':
                    self.yspeed *= -1
                    self.xspeed += random.randint(-self.fun, self.fun)

    def show(self):
        pygame.draw.rect(canvas, WHITE, self.rec, 0)

    def highlight(self, ball):
        pygame.draw.rect(canvas, PINK, ball, 0)

class Border():
    def __init__(self, x, y, w, h):
        self.x = x
        self.y = y
        self.w = w
        self.h = h
        self.rec = pygame.Rect(self.x, self.y, self.w, self.h)

    def show(self):
        pygame.draw.rect(canvas, WHITE, self.rec, 0)

### GAME SETTINGS
OFFSET = 50
BORDER_WIDTH = 10

BALL_COLOR = WHITE
BACKGROUND_COLOR = BLACK
HIGHLIGHT_COLOR = PINK

SPEED = list(range(-2, 2))
SPEED.remove(0)
FUN =  list(range(0, 2))

MIN_BALL_WIDTH = 30
MAX_BALL_WIDTH = 50

BALL_DENSITY = random.randint(5, 20)
###

ballPool = []
borders = {'left': Border(OFFSET, OFFSET, BORDER_WIDTH, HEIGHT - 2*OFFSET), 'right': Border(WIDTH - OFFSET, OFFSET, BORDER_WIDTH, HEIGHT - 2*OFFSET),
            'up': Border(OFFSET, OFFSET, WIDTH - 2*OFFSET + BORDER_WIDTH, BORDER_WIDTH), 'down': Border(OFFSET, HEIGHT - OFFSET, WIDTH - 2*OFFSET + BORDER_WIDTH, BORDER_WIDTH)}

for i in range(BALL_DENSITY):
    x, y, w, xspeed, yspeed, fun = randomizeBall()
    ballPool.append(Ball(x, y, w, xspeed, yspeed, fun))

while True:
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            quitGame()
    canvas.fill(BLACK)

    for border, value in borders.items():
        value.show()

    ballRect = []
    for ball in ballPool:
        ball.update(borders)
        ball.show()
        ballRect.append(ball.rec)

    for rectBall in ballRect:
        ballRect.remove(rectBall)
        collision = rectBall.collidelist(ballRect)
        if collision != -1:
            ball.highlight(rectBall)
            ball.highlight(ballRect[collision])

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