How to put my game loop into a function?

[temlotresid6]

I am not sure if this is too much code to post, but this is the most code I've had in a single program so now it's overwhelming trying to solve this. I didn't think of adding a game_loop function until after I created my classes and other functions and now I get errors no matter what I try. I am currently trying to press and hold down a key (such as f) and then it will run the game loop, otherwise the game does not start. If I get that to work then I'll be able to replace that key if statement with my separate main menu program and integrate it with my game here.

When I put the entire code within the "while Running:" loop into the "def game_loop():" function, there were a bunch of errors to do with all of my variables such as them not existing or being setup properly, so I made most/all of them global variables and referenced them outside of the function but that would cause more problems.

import pygame
import random
import math

pygame.init()

# Screen dimensions
WIDTH, HEIGHT = 800, 600
screen = pygame.display.set_mode((WIDTH, HEIGHT))
pygame.display.set_caption('Collision!')

# Colours
BLACK = (0, 0, 0)
WHITE = (255, 255, 255)
RED = (255, 0, 0)

# Have either 5 or 10 total AI objects spawned into the game
spawned_at_least_5 = False
spawned_at_least_10 = False

# Powerup variables - before the main game loop
powerup = None
powerup_spawned = False
powerup_destroy_time = 0

# Player brace mechanic
brace_charges = 0
charge_timer = 0  # Timer to track elapsed time for adding charges
charge_interval = 30  # Time in seconds to add a charge

# Score related variables, set all to 0 prior to the game starting
score = 0
last_score_increase = 0
last_score_double = 0

# Timer variables
def display_time(time_ms):
    """Return time in milliseconds as string mm:ss."""
    minutes, seconds = divmod(int(time_ms / 1000), 60)
    return f"{minutes: 2}:{seconds:02}"
    
timer_event = 10000  # This is 10 seconds in milliseconds
last_spawn_time = pygame.time.get_ticks()
last_time = pygame.time.get_ticks()
start_timer = pygame.time.get_ticks()
clock = pygame.time.Clock()
elapsed_time = pygame.time.get_ticks()
start_time = pygame.time.get_ticks()

class Player:
    def __init__(self):
        self.x = WIDTH // 2
        self.y = HEIGHT // 2
        self.speed = 2
        self.width = 30
        self.height = 30
        self.update_rect()

    def update_rect(self):
        self.rect = pygame.Rect(self.x, self.y, self.width, self.height)
        
    def move(self, dx, dy):
        self.x += dx * self.speed
        self.y += dy * self.speed
        self.x = max(0, min(self.x, WIDTH - self.width))
        self.y = max(0, min(self.y, HEIGHT - self.height))
        self.update_rect()

    def grow(self):
        self.width += self.width * 0.1
        self.height += self.height * 0.1
        self.width = min(self.width, WIDTH)
        self.height = min(self.height, HEIGHT)
        self.update_rect()

    def reduce(self):
        self.width -= self.width * 0.15
        self.height -= self.height * 0.15
        self.update_rect()

    def draw(self):
        if pygame.key.get_pressed()[pygame.K_c] and brace_charges > 0:
            self.color = (0, 0, 255)  # Blue for active brace state
        else:
            self.color = RED
        pygame.draw.rect(screen, self.color, (self.x, self.y, self.width, self.height))
        
class AIObject:
    def __init__(self):
        self.x, self.y = random.randint(50, WIDTH-50), random.randint(50, HEIGHT-50)
        self.target_x, self.target_y = self.x, self.y
        self.speed = 2.5
        self.movement_phase = 0
        self.timer = 0

    def set_new_target(self):
        self.target_x = random.randint(50, WIDTH-50)
        self.target_y = random.randint(50, HEIGHT-50)

    def move(self):
        dx = self.target_x - self.x
        dy = self.target_y - self.y
        distance = math.sqrt(dx**2 + dy**2)
        if distance > self.speed:
            self.x += (dx / distance) * self.speed
            self.y += (dy / distance) * self.speed
        else:
            self.x, self.y = self.target_x, self.target_y
            self.timer += 1
            if self.timer >= random.randint(200, 300):
                if self.movement_phase == 0:
                    self.set_new_target()
                    self.movement_phase = 1
                elif self.movement_phase == 1:
                    self.set_new_target()
                    self.movement_phase = 2
                elif self.movement_phase == 2:
                    self.movement_phase = 0
                    self.timer = 0
                    return "delete"
                self.timer = 0

    def check_overlap_with_player(self, player):
        potential_rect = pygame.Rect(self.x - 15, self.y - 15, 30, 30)
        return player.rect.colliderect(potential_rect)
    
    @staticmethod
    def spawn_multiple():
        global spawned_at_least_5, spawned_at_least_10
        if len(ai_objects) >= 5 and not spawned_at_least_5:
            for _ in range(3):
                ai_objects.append(AIObject())
            spawned_at_least_5 = True
        if len(ai_objects) >= 10 and not spawned_at_least_10:
            for _ in range(3):
                ai_objects.append(AIObject())
            spawned_at_least_10 = True
            
    def draw(self):
        pygame.draw.circle(screen, WHITE, (int(self.x), int(self.y)), 15)

class Powerup:
    def __init__(self):
        self.x, self.y = random.randint(50, WIDTH-50), random.randint(50, HEIGHT-50)
        self.width = 20
        self.height = 20
        self.color = (200, 100, 0)
        self.update_rect()

    def update_rect(self):
        self.rect = pygame.Rect(self.x, self.y, self.width, self.height)
        
    def draw(self, surface):
        pygame.draw.rect(surface, self.color, (self.x, self.y, self.width, self.height))

# Game setup
ai_objects = []
player = Player()
game_over = False
running = True

def game_loop():
    global score, last_score_increase, last_score_double, brace_charges
    global powerup, powerup_spawned, powerup_destroy_time, elapsed_time
    global last_spawn_time, last_time, start_timer, start_time, game_over

    # Reset game state variables
    ai_objects.clear()
    player.__init__()
    score = 0
    last_score_increase = 0
    last_score_double = 0
    brace_charges = 0
    powerup = None
    powerup_spawned = False
    powerup_destroy_time = 0
    elapsed_time = 0
    last_spawn_time = pygame.time.get_ticks()
    last_time = pygame.time.get_ticks()
    start_timer = pygame.time.get_ticks()
    start_time = pygame.time.get_ticks()

    while True:
        current_time = pygame.time.get_ticks()
        charge_timer = current_time - start_timer
        
        # Update the timer
        if not game_over:
            elapsed_time = current_time - start_time
            if elapsed_time - last_score_increase >= 5000:
                score += 1
                last_score_increase += 5000

            if elapsed_time - last_score_double >= 30000:
                score *= 2
                last_score_double += 30000

            if charge_timer >= 30000:
                brace_charges += 1
                start_timer = pygame.time.get_ticks()

        if game_over:
            screen.fill(BLACK)
            font = pygame.font.Font(None, 36)

            game_over_text = font.render("Game Over! Press SPACE to restart or ESC to quit.", True, WHITE)
            game_over_rect = game_over_text.get_rect(center=(WIDTH // 2, HEIGHT // 2 - 100))

            time_text = font.render(f"Time: {display_time(elapsed_time)}", True, WHITE)
            screen.blit(time_text, (WIDTH // 2 - time_text.get_width() // 2, HEIGHT // 2 - time_text.get_height() // 2))

            score_text = font.render(f"Score: {score}", True, WHITE)
            screen.blit(score_text, (WIDTH // 2 - time_text.get_width() // 2, HEIGHT // 2 - score_text.get_height() // 2 + 25))

            screen.blit(game_over_text, game_over_rect)
            pygame.display.flip()
            
            for event in pygame.event.get():
                if event.type == pygame.QUIT or (event.type == pygame.KEYDOWN and event.key == pygame.K_ESCAPE):
                    return  # Exit the game loop
                if event.type == pygame.KEYDOWN and event.key == pygame.K_SPACE:
                    game_over = False
                    continue
            continue  # Skip the rest of the loop

        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                return
            
        # Player movement
        keys = pygame.key.get_pressed()
        if keys[pygame.K_UP]:
            player.move(0, -1)
        if keys[pygame.K_DOWN]:
            player.move(0, 1)
        if keys[pygame.K_LEFT]:
            player.move(-1, 0)
        if keys[pygame.K_RIGHT]:
            player.move(1, 0)
        
        if current_time - last_spawn_time > random.randint(2000, 4000):
            new_obj = AIObject()
            attempts = 0
            max_attempts = 100
            while new_obj.check_overlap_with_player(player) and attempts < max_attempts:
                new_obj.x, new_obj.y = random.randint(50, WIDTH-50), random.randint(50, HEIGHT-50)
                attempts += 1
            if attempts < max_attempts:
                ai_objects.append(new_obj)
            last_spawn_time = current_time

        if current_time - last_time >= timer_event:
            player.grow()
            last_time = current_time

        AIObject.spawn_multiple()

        # Powerup logic
        if powerup is None and not powerup_spawned and current_time - powerup_destroy_time >= 15000:
            powerup = Powerup()
            powerup_spawned = True

        if powerup:
            player.update_rect()
            if player.rect.colliderect(powerup.rect):
                player.reduce()
                powerup = None
                powerup_spawned = False
                powerup_destroy_time = current_time

        to_remove = []
        
        # Update AI objects and check for collisions
        for obj in ai_objects[:]:
            action = obj.move()
            if action == "delete":
                ai_objects.remove(obj)

        brace_key = pygame.key.get_pressed()
        if brace_key[pygame.K_c] and brace_charges > 0:
            for obj in ai_objects[:]:
                if player.rect.colliderect(pygame.Rect(obj.x - 15, obj.y - 15, 30, 30)):
                    to_remove.append(obj)
                    brace_charges -= 1
                    score += 5
                    if brace_charges < 0:
                        brace_charges = 0
        else:
            player.update_rect()
            for obj in ai_objects:
                if player.rect.colliderect(pygame.Rect(obj.x - 15, obj.y - 15, 30, 30)):
                    game_over = True
                    to_remove.append(obj)
                    break

        # Remove all objects that are to be deleted
        for obj in to_remove:
            if obj in ai_objects:
                ai_objects.remove(obj)

        if game_over:
            continue  # Skip the drawing and updating display part if game over

        screen.fill(BLACK)

        for obj in ai_objects[:]:
            action = obj.move()
            if action != "delete":
                obj.draw()
        
        player.draw()

        if powerup:
            powerup.draw(screen)

        keyp = pygame.key.get_pressed()
        if keyp[pygame.K_b]:
            font = pygame.font.Font(None, 36)
            score_text = font.render(f"Score: {score}", True, WHITE)
            screen.blit(score_text, (10, 10))
            time_text = font.render(f"Time elapsed: {display_time(elapsed_time)}", True, WHITE)
            screen.blit(time_text, (565, 10))
            brace_text = font.render(f"Charges: {brace_charges}", True, WHITE)
            screen.blit(brace_text, (10, 35))
            
        pygame.display.flip()
        clock.tick(50)  # 50 FPS for smooth movement

# Main loop
while True:
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            pygame.quit()
            exit()
        if event.type == pygame.KEYDOWN:
            if event.key == pygame.K_f:
                game_loop()

pygame.quit()

Edit: Now there's no errors, but when pressing "space" to restart the game, the game does not reset but it's just a continuation of the previous game. The time and score also don't reset, and the AI circle objects start to flicker.

[menator01]

I moved the global score above the score = 0 and it worked. I didn't do any more than that.
You should try to write your code without using global if possible.
It would be better to pass it as an argument.

You did good. I like it.

[temlotresid6]

I moved the global score above the score = 0 and it worked. I didn't do any more than that.
You should try to write your code without using global if possible.
It would be better to pass it as an argument.

You did good. I like it.

Do you mean that it worked once you put the entire game loop into a function and then moved global score? The game as it is actually works great and there's no faults with it I think, but I'm just not sure on how to place it inside a function so that if I have a main menu class, I can call the game_loop function when I press "Play", if that makes sense. Thanks for the reply

[deanhystad]

Post the code that doesn’t work. I don’t want to guess what errors you make when moving the game loop to a function.

[temlotresid6]

Post the code that doesn’t work. I don’t want to guess what errors you make when moving the game loop to a function.

I updated my main post with a recent change I made, there's no error, but nothing resets upon restarting the game, even though from what I can see, it should reset. The main game loop during the first run time seems to work fine, but pressing space to restart causes issues which I edited into the post at the bottom as well. I basically used "global" for almost everything, but not sure if that's a good idea.

Edit: I do get this error when pressing escape to exit:

---------------------------------------------------------------------------
error                                     Traceback (most recent call last)
Cell In[1], line 331
    329 # Main loop
    330 while True:
--> 331     for event in pygame.event.get():
    332         if event.type == pygame.QUIT:
    333             pygame.quit()

error: video system not initialized

And even in the first runtime before ever restarting, the AI objects do flicker but I can't figure out why yet.

[menator01]

I think it's because you are using pygame.quit(). In the while loop I usuall set running = True then in the event set running = False

I'm trying to write my own version of your game. For the timers I'm using pygame's time.set_timer. Don't have to write timers then. Just using events. Am having problems right now as how to get the balls to move. They jump to the next position instead of moving.

---

This is what I have so far.
Updated code. Now have movement
todo:

• collision detection
  
• keep from at player when first starting
  
• and more
  

import pygame
import random
import math

pygame.init()

WIDTH, HEIGHT = (800, 600)
screen = pygame.display.set_mode((WIDTH, HEIGHT))
pygame.display.set_caption('Collision')

CLOCK = pygame.time.Clock()

FPS = 60

BLACK = (0, 0, 0)
WHITE = (255, 255, 255)
RED = (255, 0, 0)
ORANGE = (255, 180, 0)
BLUE = (0, 180, 255)
GREEN = (0, 255, 0)
CYAN = (0, 255, 255)
YELLOW = (255, 255, 0)

running = True

# Create some sprites groups
# ball_sprite and player_sprite are for collision detection
# allsprites is for updating the window sprites
player_sprite = pygame.sprite.Group()
ball_sprite = pygame.sprite.Group()
allsprites = pygame.sprite.Group()

class Player(pygame.sprite.Sprite):
    def __init__(self):
        pygame.sprite.Sprite.__init__(self)
        self.width = self.height = 30
        self.image = pygame.Surface((self.width, self.height))
        self.image.fill(RED)
        self.rect = self.image.get_rect()
        self.rect.center = ((WIDTH/2, HEIGHT/2))
        self.speedx = self.speedy = 2

        # Create a user event
        self.grow_timer = pygame.USEREVENT + 1

        # Set the event to fire the grow method every 10 seconds
        pygame.time.set_timer(self.grow_timer, 10000)

        player_sprite.add(self)
        allsprites.add(self)

    def grow(self):
        ''' Method for growing the player '''
        self.width += self.rect.width * 0.1
        self.height += self.rect.height * 0.1

        # Don't want the size to big 80
        if self.width > 80:
            self.width = 80
            self.height = 80

        # Redraw the player
        self.image = pygame.Surface((self.width, self.height))
        self.image.fill(RED)

    def shrink(self):
        ''' Method for shrinking player '''
        self.width -= self.rect.width * 0.15
        self.height -= self.rect.height * 0.15

        # Don't want the size to be smaller than 20
        if self.width <= 20:
            self.width = 20
            self.height = 20

        self.image = pygame.Surface((self.width, self.height))
        self.image.fill(RED)



    def update(self):
        self.speedx = self.speedy = 0

        # Get keys pressed
        key = pygame.key.get_pressed()

        if key[pygame.K_LEFT]:
            self.speedx = -2

        if key[pygame.K_RIGHT]:
            self.speedx = 2

        if key[pygame.K_UP]:
            self.speedy = -2

        if key[pygame.K_DOWN]:
            self.speedy = 2

        # Set bounds
        if self.rect.left <= 0:
            self.rect.left = 0

        if self.rect.right >= WIDTH:
            self.rect.right = WIDTH

        if self.rect.top <= 0:
            self.rect.top = 0

        if self.rect.bottom >= HEIGHT:
            self.rect.bottom = HEIGHT

        self.rect.centerx += self.speedx
        self.rect.centery += self.speedy


class Ball(pygame.sprite.Sprite):
    def __init__(self):
        pygame.sprite.Sprite.__init__(self)
        self.image = pygame.Surface((30,30))
        self.image.fill(BLACK)
        self.image.set_colorkey(BLACK)
        self.rect = self.image.get_rect()
        self.rect.x = random.randint(30, WIDTH-30)
        self.rect.y = random.randint(30, HEIGHT-30)
        self.speedx = 2.5
        self.speedy = 2.5
        self.dx = self.rect.x
        self.dy = self.rect.y


        colors = [ORANGE, WHITE, GREEN, YELLOW, CYAN]
        pygame.draw.circle(self.image, random.choice(colors), (self.rect.width/2, self.rect.height/2), 15)

        self.timer = pygame.USEREVENT + random.randint(1, 1000)
        pygame.time.set_timer(self.timer, random.randrange(1000, 2000))

        ball_sprite.add(self)
        allsprites.add(self)

    
    def update(self):
        
        if self.rect.x < self.dx:
            self.speedx = 2.5

        if self.rect.x > self.dx:
            self.speedx = -2.5

        if self.rect.y < self.dy:
            self.speedy = 2.5

        if self.rect.y > self.dy:
            self.speedy = -2.5

        if self.rect.x != self.dx and self.rect.y != self.dy:
            self.rect.x += self.speedx
            self.rect.y += self.speedy

    def move(self):
        self.dx = random.randint(30, WIDTH-30)
        self.dy = random.randint(30, HEIGHT-30)

# Create the player     
player = Player()

# Create empty list for balls
balls = []

# Set for counter
i = 1

# Create a user event
starter = pygame.USEREVENT + 2000

# Start timer for user event
pygame.time.set_timer(starter, 2000)

while running:

    screen.fill(BLACK)           

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

        # Fire user event for growth
        # if event.type == player.grow_timer:
        #     player.grow()

        # Check for user event
        if event.type == starter:
            
            # Check how many balls in balls list - if less than 5 create a ball
            # and create a user event and set a timer
            if len(balls) < 5:
                balls.append(Ball())
                start = pygame.USEREVENT + 1000+1
                pygame.time.set_timer(start, 1000)
            
        # For moving around the balls in the list
        for index, ball in enumerate(balls):
            if event.type == start:
                balls[index].move()

    # Update all sprites
    allsprites.update()
    allsprites.draw(screen)

    pygame.display.update()
    CLOCK.tick(FPS)

pygame.quit()

[temlotresid6]

I think it's because you are using pygame.quit(). In the while loop I usuall set running = True then in the event set running = False

I'm trying to write my own version of your game. For the timers I'm using pygame's time.set_timer. Don't have to write timers then. Just using events. Am having problems right now as how to get the balls to move. They jump to the next position instead of moving.

---

This is what I have so far.

import pygame
import random
import math

pygame.init()

WIDTH, HEIGHT = (800, 600)
screen = pygame.display.set_mode((WIDTH, HEIGHT))
pygame.display.set_caption('Collision')

CLOCK = pygame.time.Clock()

FPS = 60



BLACK = (0, 0, 0)
WHITE = (255, 255, 255)
RED = (255, 0, 0)
ORANGE = (255, 180, 0)
BLUE = (0, 180, 255)
GREEN = (0, 255, 0)
CYAN = (0, 255, 255)
YELLOW = (255, 255, 0)

running = True

player_sprite = pygame.sprite.Group()
enemy_sprite = pygame.sprite.Group()
allsprites = pygame.sprite.Group()

class Player(pygame.sprite.Sprite):
    def __init__(self):
        pygame.sprite.Sprite.__init__(self)
        self.width = self.height = 30
        self.image = pygame.Surface((self.width, self.height))
        self.image.fill(RED)
        self.rect = self.image.get_rect()
        self.rect.center = ((WIDTH/2, HEIGHT/2))
        self.speedx = self.speedy = 2

        # Create a user event
        self.grow_timer = pygame.USEREVENT + 1

        # Set the event to fire the grow method every 10 seconds
        pygame.time.set_timer(self.grow_timer, 10000)

        player_sprite.add(self)
        allsprites.add(self)

    def grow(self):
        ''' Method for growing the player '''
        self.width += self.rect.width * 0.1
        self.height += self.rect.height * 0.1

        # Don't want the size to big 80
        if self.width > 80:
            self.width = 80
            self.height = 80

        # Redraw the player
        self.image = pygame.Surface((self.width, self.height))
        self.image.fill(RED)

    def shrink(self):
        ''' Method for shrinking player '''
        self.width -= self.rect.width * 0.15
        self.height -= self.rect.height * 0.15

        # Don't want the size to be smaller than 20
        if self.width <= 20:
            self.width = 20
            self.height = 20

        self.image = pygame.Surface((self.width, self.height))
        self.image.fill(RED)



    def update(self):
        self.speedx = self.speedy = 0

        # Get keys pressed
        key = pygame.key.get_pressed()

        if key[pygame.K_LEFT]:
            self.speedx = -2

        if key[pygame.K_RIGHT]:
            self.speedx = 2

        if key[pygame.K_UP]:
            self.speedy = -2

        if key[pygame.K_DOWN]:
            self.speedy = 2

        # Set bounds
        if self.rect.left <= 0:
            self.rect.left = 0

        if self.rect.right >= WIDTH:
            self.rect.right = WIDTH

        if self.rect.top <= 0:
            self.rect.top = 0

        if self.rect.bottom >= HEIGHT:
            self.rect.bottom = HEIGHT

        self.rect.centerx += self.speedx
        self.rect.centery += self.speedy


class Ball(pygame.sprite.Sprite):
    def __init__(self):
        pygame.sprite.Sprite.__init__(self)
        self.image = pygame.Surface((30,30))
        self.image.fill(BLACK)
        self.image.set_colorkey(BLACK)
        self.rect = self.image.get_rect()
        self.rect.x = random.randint(30, WIDTH-30)
        self.rect.y = random.randint(30, HEIGHT-30)
        self.speed = 2.5

        colors = [ORANGE, WHITE, GREEN, YELLOW, CYAN]
        pygame.draw.circle(self.image, random.choice(colors), (self.rect.width/2, self.rect.height/2), 15)

        self.timer = pygame.USEREVENT + random.randint(1, 1000)
        pygame.time.set_timer(self.timer, random.randrange(1000, 5000))

        enemy_sprite.add(self)
        allsprites.add(self)


    def move(self):
        dx = random.randint(30, WIDTH-30)
        dy = random.randint(30, HEIGHT-30)

        distance = math.sqrt(dx**2+dy**2)
        
        self.rect.x = dx
        self.rect.y = dy


player = Player()
mob = Ball()

while running:
    
    screen.fill(BLACK) 

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

        # if event.type == player.grow_timer:
        #     player.grow()

        if event.type == mob.timer:
            mob.move()

        


    allsprites.update()
    allsprites.draw(screen)

    pygame.display.update()
    CLOCK.tick(FPS)

pygame.quit()

In my def(move) function inside the AI (ball) class, I have code that I made which makes the movement of the AI smooth rather than that snap/teleport effect. I also had that issue originally but managed to solve it.

My AI ball movement code is a bit different though since I have additional movement features for it where, I spawn balls at a random interval -> move the ball after another random interval (move 1) -> after another random interval, move the ball again (move 2) -> after a brief delay, remove the ball after it has completed two moves. Then this just repeats for every ball object that spawns into the game. You can just remove the whole movement_phase == 1 and == 2, etc, parts of that move function and I think that's a good method for getting smooth movement rather than instant.

[menator01]

Updated code

Have everything working, did not add a scoring system
Did add a start page. In your code the balls will move and then stop.
Mine they just keep moving. After 2 moves they disappear.
Checks how many sprites are in the ball group if it's not the amount indicated
one is created.

---

Noticed that, if the game has been running for a bit, player collides with a ball and goes to start page, hit spacebar again, it seems to take a couple of seconds for the game to start again. Even noticed that sometimes it's hard to exit game from start page. The shmup game I wrote a while back doesn't seem to do that. Of coarse I don't have all those pygame.time.set_timer either.

---

Added a scoring system - every one seconds gives 0.20 points, every shrinker gives 0.35 points
Re-wrote the code a little, everything appears to be working right.
Updated code in case the music folder doesn't exist
Added a function to create a rgb range of colors for the balls. Range is from 50 to 255. Didn't want the color to be blackish. That
would be hard to impossible to see on a black background.

---

My final version

import pygame
from pygame import mixer
from random import randint, choice, random, shuffle, randrange
import math
from pathlib import Path
import os

pygame.init()
mixer.init()

path = Path(__file__).parent

# Set width and height of game window
WIDTH, HEIGHT = (1280, 840)
screen = pygame.display.set_mode((WIDTH, HEIGHT))
pygame.display.set_caption('Collision Balls')
 
CLOCK = pygame.time.Clock()

# Frames per second
FPS = 60

# Define some colors
BLACK = (0, 0, 0)
WHITE = (255, 255, 255)
RED = (255, 0, 0)
LIGHTBLUE = (180, 255, 255)
BLUE = (0, 180, 255)
GREEN = (0, 255, 0)
CYAN = (0, 255, 255)
YELLOW = (255, 255, 0)
ORANGE = (255,150,0)
PURPLE = (200,100,250)

# Set game variables
running = True
gamestate = 'page'
ball_number = 8


# Timers
checker = pygame.USEREVENT+3
pygame.time.set_timer(checker, 1000)

show_shrinker = pygame.USEREVENT+4
pygame.time.set_timer(show_shrinker, 5000)

invincible_timer = pygame.USEREVENT+5
pygame.time.set_timer(invincible_timer, 1000)

# Sprite groups
player_sprite = pygame.sprite.Group()
ball_sprite = pygame.sprite.Group()
shrinker_sprite = pygame.sprite.Group()
allsprites = pygame.sprite.Group()


# Define a start page - also is the game over page
def page(score=0):
    screen.fill(BLACK)
    font = pygame.font.SysFont(None, 180)
    text = 'Collision Balls'
    title_width, title_height = font.size(text)
    title = font.render(text, True, ORANGE)

    font = pygame.font.SysFont(None, 80)
    text = f'Final Score: {score}'
    score_width, score_height = font.size(text)
    scores = font.render(text, True, PURPLE)

    font = pygame.font.SysFont(None, 80)
    text = 'Press spacebar to start'
    instruct_width, instruct_height = font.size(text)
    instructions = font.render(text, True, GREEN)

    screen.blit(title, ((WIDTH/2-title_width/2), HEIGHT*0.05))
    screen.blit(instructions, ((WIDTH/2-instruct_width/2), HEIGHT*0.5))
    if score > 0:
        screen.blit(scores, ((WIDTH/2-score_width/2), HEIGHT*0.3))

    pygame.display.update()
    player_sprite.empty()
    ball_sprite.empty()
    shrinker_sprite.empty()
    allsprites.empty()


def create_color():
    color = ()
    for i in range(3):
        n = randrange(50,255)
        color = color + (n,)
    return color

class Music:
    def __init__(self):
        self.tracks = None
        if os.path.exists(f'{path}/media/Royalty Free'):
            self.tracks = os.listdir(f'{path}/media/Royalty Free')

    def play(self):
        if self.tracks:
            tracks = []
            if not tracks:
                tracks = self.tracks.copy()
                shuffle(tracks)
            track = tracks.pop(0)
            mixer.music.load(f'{path}/media/Royalty Free/{track}')
            mixer.music.set_volume(0.5)
            mixer.music.play()

    def stop(self):
        mixer.music.stop()

    
class Player(pygame.sprite.Sprite):
    def __init__(self):
        pygame.sprite.Sprite.__init__(self)
        self.width = self.height = 30
        self.image = pygame.Surface((self.width, self.height))
        self.image.fill(RED)
        self.rect = self.image.get_rect()
        self.rect.center = ((WIDTH/2, HEIGHT/2))
        self.speedx = self.speedy = 2

        self.counter = 10

        self.score = 0
        self.invincible = False
        self.current = pygame.time.get_ticks()
 
        # Create a user event
        self.grow_timer = pygame.USEREVENT + 1
 
        # Set the event to fire the grow method every 10 seconds
        pygame.time.set_timer(self.grow_timer, 10000)
 
        player_sprite.add(self)
        allsprites.add(self)
 
    def grow(self):
        ''' Method for growing the player '''
        self.width += self.rect.width * 0.1
        self.height += self.rect.height * 0.1
 
        # Don't want the size to big 80
        if self.width > 80:
            self.width = 80
            self.height = 80
 
        # Redraw the player
        self.image = pygame.Surface((self.width, self.height))
        self.image.fill(RED)
 
    def shrink(self):
        ''' Method for shrinking player '''
        self.width -= self.rect.width * 0.15
        self.height -= self.rect.height * 0.15
 
        # Don't want the size to be smaller than 20
        if self.width <= 20:
            self.width = 20
            self.height = 20
 
        self.image = pygame.Surface((self.width, self.height))
        self.image.fill(RED)
 
    def change(self):
        self.invincible = False

    def update(self):
        # Set up a scoring system. For every one secong give .20 score
        # Ever five seconds get one point
        if pygame.time.get_ticks() - self.current >= 1000:
            self.score += 0.20
            self.current = pygame.time.get_ticks()

        self.image.fill(PURPLE) if self.invincible else self.image.fill(RED)

        self.speedx = self.speedy = 0
 
        # Get keys pressed
        key = pygame.key.get_pressed()
 
        if key[pygame.K_LEFT] or key[pygame.K_a]:
            self.speedx = -2
 
        if key[pygame.K_RIGHT] or key[pygame.K_d]:
            self.speedx = 2
 
        if key[pygame.K_UP] or key[pygame.K_w]:
            self.speedy = -2
 
        if key[pygame.K_DOWN] or key[pygame.K_SPACE]:
            self.speedy = 2
 
        # Set bounds
        if self.rect.left <= 0:
            self.rect.left = 0
 
        if self.rect.right >= WIDTH:
            self.rect.right = WIDTH
 
        if self.rect.top <= 0:
            self.rect.top = 0
 
        if self.rect.bottom >= HEIGHT:
            self.rect.bottom = HEIGHT
 
        self.rect.centerx += self.speedx
        self.rect.centery += self.speedy
    

class Ball(pygame.sprite.Sprite):
    ''' Class creates a ball object '''
    def __init__(self):
        pygame.sprite.Sprite.__init__(self)
        self.image = pygame.Surface((30,30))
        self.image.fill(BLACK)
        self.image.set_colorkey(BLACK)
        self.rect = self.image.get_rect()
        self.rect.x = randint(self.rect.width, WIDTH-self.rect.width)
        self.rect.y = randint(self.rect.height, HEIGHT-self.rect.height)
        self.speed = 2.5
        self.target_x, self.target_y = self.rect.x, self.rect.y

        self.moves = 0

        self.colors = create_color()
        pygame.draw.circle(self.image, self.colors, (self.rect.width/2, self.rect.height/2), 15)
 
        self.timer = pygame.USEREVENT + randint(10, 1000)
        pygame.time.set_timer(self.timer, 20)
 
        ball_sprite.add(self)
        allsprites.add(self)

    def target(self):
        ''' Method for setting a new target '''
        self.target_x = randint(self.rect.width, WIDTH-self.rect.width)
        self.target_y = randint(self.rect.height, HEIGHT-self.rect.height)
       
        # Increase counter
        self.moves += 1

    
    def move(self):
        ''' Method for moving ball '''
        dx = self.target_x - self.rect.x
        dy = self.target_y - self.rect.y
        distance = math.sqrt(dx**2 + dy**2)

        if distance > self.speed:
            self.rect.x += (dx/distance) * self.speed
            self.rect.y += (dy/distance) * self.speed
        else:
            self.rect.x, self.rect.y = self.target_x, self.target_y
            self.target()   


class Shrinker(pygame.sprite.Sprite):
    def __init__(self):
        pygame.sprite.Sprite.__init__(self)
        self.image = pygame.Surface((20,20))
        self.image.fill(ORANGE)

        self.rect = self.image.get_rect()

        self.rect.x = randint(self.rect.width, WIDTH-self.rect.width)
        self.rect.y = randint(self.rect.height, HEIGHT-self.rect.height)

        self.change_player = False

        self.born = pygame.time.get_ticks()

        shrinker_sprite.add(self)
        allsprites.add(self)

    def update(self):

        # Kills shrinker after five seconds
        if pygame.time.get_ticks() - self.born > 5000:
            self.born = pygame.time.get_ticks()
            self.kill()

    def change(self):
        ''' Method for random creation of either a regular shrinker or powered shrinker '''
        if random() > 0.7:
            self.image.fill(PURPLE)
            self.change_player = True
        else:
            self.image.fill(ORANGE)


player = Player()

music = Music()

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

        # Countdown timer for invincible time
        elif event.type == invincible_timer and player.invincible:
            player.counter -= 1
            if player.counter <= 0:
                player.change()
                player.counter = 10
                player.invincible = False

        # Create new shrinker
        elif event.type == show_shrinker:
            shrinker = Shrinker()
            shrinker.change()

        # Event for player growth
        elif event.type == player.grow_timer:
            player.grow()

        # Check how many sprites are in ball_sprite
        # If less than ball_number create ball
        elif event.type == checker:
            if len(ball_sprite) < ball_number:
                Ball()

        for sprite in ball_sprite:
            if event.type == sprite.timer:
                sprite.move()
            elif sprite.moves == 3:
                sprite.kill()

    if gamestate == 'play':
        shuffle(music.tracks)
        if not mixer.music.get_busy():
            music.play()
        screen.fill(BLACK)
        keys = pygame.key.get_pressed()

        font = pygame.font.SysFont(None, 35)

        text = f'Score: {round(player.score)}'
        surface = font.render(text, True, CYAN)
        screen.blit(surface, (10,10))

        if player.invincible:
            text = f'Invincible Timer: {player.counter:02}'
            surface = font.render(text, True, ORANGE)
            screen.blit(surface, (WIDTH-250, 10))

        # Detect if player is shieled or not
        # If player is shieled get points for ball collision and not die
        # Shield will be placed random
        shield = False if player.invincible else True

        # Detect sprite collisions
        collision = pygame.sprite.groupcollide(ball_sprite, player_sprite, True, shield)
    
        # Collision with player and ball. Player not invincible
        if collision and not player.invincible:
            score = round(player.score)
            gamestate = 'page'

        # Player is invincible
        elif collision and player.invincible:
            player.score += 0.50

        # Check for collision with shrinker
        shrinker_collision = pygame.sprite.groupcollide(shrinker_sprite, player_sprite, True, False)
        if shrinker_collision:
            player.shrink()
            if shrinker.change_player:
                player.invincible = True
                player.counter = 10

            player.score += 0.35

        allsprites.update()
        allsprites.draw(screen)

        pygame.display.update()
        CLOCK.tick(FPS)

    elif gamestate == 'page':
        music.stop()
        # Call page and pass player score
        page(score=round(player.score))
        keys = pygame.key.get_pressed()
        if keys[pygame.K_SPACE]:
            player = Player()
            gamestate = 'play'
        
pygame.quit()