mipt_labs/lab4/main.py

from __future__ import annotations
import pygame
import sys
from dataclasses import dataclass, field
import time
import numpy as np
import typing
from random import uniform, choice
if typing.TYPE_CHECKING:
    from nptyping import NDArray


pygame.init()
# width and height of the screen
W, H = 900, 800
# ball's lifetime (seconds)
DEFAULT_TTL = 3
# ball's radius (default)
DEFAULT_RADIUS = 20
# time delta, ms
dt = 10
# load the font
font = pygame.font.Font('orbitron-medium.otf', 15)
N_BALLS = 18

# List of ball types: [(radius, speed, color, points)]
BALL_TYPES = [(25, 0.035, (185, 20, 50), 10), (18, 0.055, (220, 0, 0), 30), (20, 0.05, (60, 0, 160), 20)]


@dataclass
class Game:
    """
    Class for everything connected to the game

    :param score: user's score
    :param balls: the list of balls
    :param target_ball_number: the target number of balls on screen at a given time
    """
    score: int = 0
    balls: typing.List[Ball] = field(default_factory=lambda: [Ball.new_random() for _ in range(N_BALLS)])
    target_ball_number: int = N_BALLS

    def draw_everything(self, screen: pygame.Surface):
        """
        Draw everything (for example, balls) on screen
        :param screen: the screen
        """
        for ball in self.balls:
            ball.draw_the_ball(screen)

    def process_click(self, x: float, y: float):
        """
        Process mouse click, check if the user clicked on a ball and process it
        :param x: click's x coordinate
        :param y: click's y coordinate
        """
        for i in range(len(self.balls)):
            if self.balls[i].is_point_inside(x, y):
                self.score += self.balls[i].points
                self.balls.pop(i)
                break

    def tick(self, dt: float):
        """
        Progress for some time

        dt: time period, ms
        """
        if len(self.balls) < self.target_ball_number:
            n_balls = round((self.target_ball_number - len(self.balls))) # * (1 - 0.9 ** dt))
            self.balls.extend([Ball.new_random() for _ in range(n_balls)])
        for ball in self.balls:
            ball.tick(dt)


@dataclass
class Ball:
    """
    The ball the user has to catch

    :param points: the number of points the user gets for catching this ball
    :param position: ball's position on the screen
    :param velocity: ball's velocity, px/ms
    :param time_of_death: the time when this ball will be destroyed. Its death is pre-determined, which is probably
    very sad and depressing
    :param radius: the radius of the ball on screen in pixels
    :param color: ball's color
    """
    points: int = 10
    position: NDArray[(2,), float] = field(default_factory=lambda: np.array([uniform(0, W), uniform(0, H)]))
    velocity: NDArray[(2,), float] = field(default_factory=lambda: np.array([0, 0]))
    time_of_death: float = field(default_factory=lambda: time.time() + DEFAULT_TTL)
    radius: float = DEFAULT_RADIUS
    color: typing.Tuple[int, int, int] = (255, 0, 0)

    @classmethod
    def new_random(cls) -> Ball:
        radius, speed, color, points = choice(BALL_TYPES)
        angle = uniform(0, np.pi)
        return cls(points, radius=radius, color=color, velocity=speed * np.array([np.sin(angle), np.cos(angle)]))

    def draw_the_ball(self, screen: pygame.Surface):
        """
        Just draw the ball on the screen
        :param screen: the screen, obviously
        """
        pygame.draw.circle(screen, self.color, self.position.tolist(), self.radius)
        text_surface = font.render(str(self.points), False, (0, 0, 0))
        screen.blit(text_surface, (
            round(self.position[0] - text_surface.get_rect().width / 2),
            round(self.position[1] - text_surface.get_rect().height / 2)
        ))

    def is_point_inside(self, x: float, y: float) -> bool:
        """
        Check if a given point is inside this circle (to handle click events)
        :param x: point's x coordinate
        :param y: point's y coordinate
        """
        return (self.position[0] - x) ** 2 + (self.position[1] - y) ** 2 <= self.radius ** 2

    def tick(self, dt: float):
        """
        Progress for time dt (move the ball according to speed)
        dt: time in ms
        """
        self.position += self.velocity * dt
        if self.position[0] - self.radius < 0:
            self.position[0] = self.radius
            self.velocity[0] *= -1
        if self.position[0] >= W - self.radius:
            self.position[0] = W - self.radius
            self.velocity[0] *= -1
        if self.position[1] < self.radius:
            self.position[1] = self.radius
            self.velocity[1] *= -1
        if self.position[1] >= H - self.radius:
            self.position[1] = H - self.radius
            self.velocity[1] *= -1


# create the display surface object  
screen = pygame.display.set_mode((W, H))

game = Game()
  
# set the window name 
pygame.display.set_caption("Capture the ball!")
    
# the main loop
while True:
    # time delay of dt
    pygame.time.delay(dt)
      
    # iterate over the list of Event objects  
    # that was returned by pygame.event.get() method.  
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            pygame.quit()
            sys.exit(0)
        if event.type == pygame.MOUSEBUTTONUP:
            game.process_click(*pygame.mouse.get_pos())

    # clear the screen
    screen.fill((40, 46, 70))
    game.tick(dt)
    game.draw_everything(screen)
    
    pygame.display.update()