mipt_labs/lab5/gun.py

from __future__ import annotations
import math
from copy import deepcopy
from random import choice, uniform, randrange as rnd
from dataclasses import dataclass, field

import pygame
import typing
import numpy as np

FPS = 30

RED = 0xFF0000
BLUE = 0x0000FF
YELLOW = 0xFFC91F
GREEN = 0x00FF00
MAGENTA = 0xFF03B8
CYAN = 0x00FFCC
BLACK = (0, 0, 0)
WHITE = 0xFFFFFF
GREY = 0x7D7D7D
GAME_COLORS = [RED, BLUE, YELLOW, GREEN, MAGENTA, CYAN]

WIDTH = 800
HEIGHT = 600


@dataclass
class Ball:
    position: np.array
    r: float = 10
    v: np.array = field(default_factory=lambda: np.array([0., 0.]))
    color: int = field(default_factory=lambda: choice(GAME_COLORS))
    time_to_split: typing.Optional[int] = field(default_factory=lambda: choice([None, rnd(5, 25)]))
    live: int = 30

    def move(self):
        """Переместить мяч по прошествии единицы времени.

        Метод описывает перемещение мяча за один кадр перерисовки. То есть, обновляет значения
        self.x и self.y с учетом скоростей self.vx и self.vy, силы гравитации, действующей на мяч,
        и стен по краям окна (размер окна 800х600).
        """
        if self.time_to_split is not None:
            self.time_to_split -= 1
        if self.position[1] <= 500:
            self.v[1] -= 1.2
            self.position += self.v * np.array([1, -1])
            self.v[0] *= 0.99
        else:
            if np.linalg.norm(self.v) > 10 ** 0.5:
                self.v *= np.array([0.5, -0.5])
                self.position[1] = 499
            self.live -= 1
        if self.position[0] > 780:
            self.v[0] *= -0.5
            self.position[0] = 779

    def draw(self, screen):
        pygame.draw.circle(
            screen,
            self.color,
            self.position,
            self.r
        )

    def hittest(self, obj):
        """Функция проверяет, сталкивается ли данный объект с целью, описываемой в объекте obj.

        Args:
            obj: Объект, с которым проверяется столкновение.
        Returns:
            Возвращает True в случае столкновения мяча и цели. В противном случае возвращает False.
        """
        if np.linalg.norm(self.position - obj.position) < self.r + obj.r:
            return True
        else:
            return False

    def get_split(self) -> typing.Tuple[Ball, Ball]:
        ball1, ball2 = deepcopy(self), deepcopy(self)
        ball1.r /= 2 ** 0.5
        ball2.r /= 2 ** 0.5
        ball1.time_to_split = choice([None, rnd(5, 25)])
        ball2.time_to_split = choice([None, rnd(5, 25)])
        v_angle = np.arctan2(*self.v)
        delta_angle = 0.5
        ball1.v = np.linalg.norm(self.v) * np.array([np.sin(v_angle - delta_angle), np.cos(v_angle - delta_angle)])
        ball2.v = np.linalg.norm(self.v) * np.array([np.sin(v_angle + delta_angle), np.cos(v_angle + delta_angle)])
        return ball1, ball2


class Gun:
    def __init__(self):
        self.bullet = 0
        self.f2_power = 10
        self.f2_on = 0
        self.an = 1
        self.color = GREY
        self.x0 = 40
        self.y0 = 450

    def fire2_start(self, _event):
        self.f2_on = 1

    def fire2_end(self, event, game: Game):
        """Выстрел мячом.

        Происходит при отпускании кнопки мыши.
        Начальные значения компонент скорости мяча vx и vy зависят от положения мыши.
        """
        self.bullet += 1
        new_ball = Ball(position=np.array([self.x0 + self.deltas()[0], self.y0 + self.deltas()[1]]))
        new_ball.r += 5
        self.an = math.atan2((event.pos[1] - new_ball.position[1]), (event.pos[0] - new_ball.position[0]))
        new_ball.v[0] = self.f2_power * math.cos(self.an)
        new_ball.v[1] = - self.f2_power * math.sin(self.an)
        game.balls.append(new_ball)
        self.f2_on = 0
        self.f2_power = 10

    def targetting(self, event):
        """Прицеливание. Зависит от положения мыши."""
        if event and event.pos[0] != 20:
            self.an = math.atan((event.pos[1] - 450) / (event.pos[0] - 20))
        if self.f2_on:
            self.color = RED
        else:
            self.color = GREY

    def deltas(self) -> typing.Tuple[float, float]:
        length = 50 + self.f2_power
        return length * math.cos(self.an), length * math.sin(self.an)

    def draw(self, screen):
        r, length = 10, 50 + self.f2_power
        dx, dy = r * math.sin(self.an), -r * math.cos(self.an)
        delta_x, delta_y = self.deltas()
        points = [
            (self.x0 + dx, self.y0 + dy), (self.x0 - dx, self.y0 - dy),
            (self.x0 + delta_x - dx, self.y0 + delta_y - dy), (self.x0 + delta_x + dx, self.y0 + delta_y + dy)]
        pygame.draw.polygon(screen, self.color, points)

    def power_up(self):
        if self.f2_on:
            if self.f2_power < 100:
                self.f2_power += 1
            self.color = RED
        else:
            self.color = GREY


@dataclass
class Target:
    x: float = field(default_factory=lambda: rnd(600, 780))
    y: float = field(default_factory=lambda: rnd(300, 550))
    r: float = field(default_factory=lambda: rnd(9, 50))
    vx: float = field(default_factory=lambda: rnd(-4, 6))
    vy: float = field(default_factory=lambda: rnd(-4, 6))
    color: int = RED
    points: int = 0
    live: int = 1
    ax: float = 0
    ay: float = 0
    randomness_ampl: float = field(default_factory=lambda: uniform(0, 5) * choice([0, 0, 1]))
    oscillation_freq: float = field(default_factory=lambda: 0.05 * choice([-1, 1]))
    oscillation_ampl: float = field(default_factory=lambda: rnd(0, 6) * choice([0, 0, 1]))
    oscillation_phase: float = field(default_factory=lambda: uniform(0, 2 * np.pi))

    def hit(self, points=1):
        """Попадание шарика в цель."""
        self.points += points

    def draw(self, screen):
        pygame.draw.circle(
            screen,
            self.color,
            (self.x, self.y),
            self.r
        )
        pygame.draw.circle(
            screen,
            WHITE,
            (self.x, self.y),
            self.r * 0.7
        )
        pygame.draw.circle(
            screen,
            self.color if not self.oscillation_ampl and not self.randomness_ampl else BLUE,
            (self.x, self.y),
            self.r * 0.45
        )

    def move(self):
        """Переместить мяч по прошествии единицы времени.

        Метод описывает перемещение мяча за один кадр перерисовки. То есть, обновляет значения
        self.x и self.y с учетом скоростей self.vx и self.vy, силы гравитации, действующей на мяч,
        и стен по краям окна (размер окна 800х600).
        """
        if self.y < 0:
            self.vy = -abs(self.vy) * 0.5 + self.ay
        if self.x < 0:
            self.vx = abs(self.vx) * 0.75 + self.ax
        if self.y <= 500:
            self.y -= self.vy + self.oscillation_ampl * np.sin(self.oscillation_phase)
            self.x += self.vx + self.oscillation_ampl * np.cos(self.oscillation_phase)
            self.vx *= 0.98
        else:
            # if self.vx ** 2 + self.vy ** 2 > 10:
            #     self.vy = -self.vy / 2
            #     self.vx = self.vx / 2
            self.vy = 0.9 * abs(self.vy)
            self.y = 499
        if self.x > 780:
            self.vx = -abs(self.vx) / 2
            self.x = 779
        self.oscillation_phase += self.oscillation_freq
        self.ax += self.randomness_ampl * uniform(-1, 1)
        self.ay += self.randomness_ampl * uniform(-1, 1)

    @property
    def position(self):
        return np.array([self.x, self.y])


@dataclass
class Game:
    balls: typing.List[Ball] = field(default_factory=list)
    targets: typing.List[Target] = field(default_factory=lambda: [Target(), Target()])
    gun: Gun = field(default_factory=Gun)
    finished: bool = False
    clock: pygame.time.Clock = field(default_factory=pygame.time.Clock)
    fps: int = FPS

    def move(self):
        for target in self.targets:
            target.move()
        for ball in self.balls:
            ball.move()
            if ball.live < 0:
                self.balls.pop([i for i in range(len(self.balls)) if self.balls[i].position is ball.position][0])
            elif ball.time_to_split is not None and ball.time_to_split < 0:
                ball1, ball2 = ball.get_split()
                self.balls.pop([i for i in range(len(self.balls)) if self.balls[i].position is ball.position][0])
                self.balls.extend([ball1, ball2])
            for target in self.targets:
                if ball.hittest(target) and target.live:
                    target.live = 0
                    ball.live = -1
                    target.hit()
                    self.targets.remove(target)
                    self.targets.append(Target())
        self.gun.power_up()

    def draw(self, screen):
        screen.fill(WHITE)
        self.gun.draw(screen)
        for target in self.targets:
            target.draw(screen)
        for ball in self.balls:
            ball.draw(screen)
        pygame.display.update()

    def process_event(self, event):
        if event.type == pygame.QUIT:
            self.finished = True
        elif event.type == pygame.MOUSEBUTTONDOWN:
            self.gun.fire2_start(event)
        elif event.type == pygame.MOUSEBUTTONUP:
            self.gun.fire2_end(event, self)
        elif event.type == pygame.MOUSEMOTION:
            self.gun.targetting(event)

    def main_loop(self):
        screen = pygame.display.set_mode((WIDTH, HEIGHT))
        while not self.finished:
            self.draw(screen)
            self.clock.tick(self.fps)
            for event in pygame.event.get():
                self.process_event(event)
            self.move()
        pygame.quit()


pygame.init()
game = Game()
game.main_loop()