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


@dataclass
class Gun:
    bullet: bool = 0
    f2_power: int = 10
    f2_on: bool = 0
    an: float = 1
    color: int = GREY
    x0: int = 40
    y0: int = 350
    bg_color: int = BLUE
    keymap: typing.List[int] = field(
        default_factory=lambda: [pygame.K_a, pygame.K_s, pygame.K_d, pygame.K_q, pygame.K_e]
    )

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

    def fire2_end(self, 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
        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 = False
        self.f2_power = 10

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

    def handle_event(self, event, game: Game):
        if event.type == pygame.MOUSEMOTION:
            self.targetting(event)
        if event.type == pygame.KEYDOWN and event.key == self.keymap[1]:
            self.fire2_start(event)
        if event.type == pygame.KEYUP and event.key == self.keymap[1]:
            self.fire2_end(game)
        elif event.type == pygame.MOUSEBUTTONDOWN:
            self.fire2_start(event)

    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)
        pygame.draw.polygon(screen, self.bg_color,
                            [(self.x0 - 50, self.y0), (self.x0 + 50, self.y0), (self.x0 + 50, self.y0 + 35),
                             (self.x0 - 50, self.y0 + 35)])
        pygame.draw.polygon(screen, self.bg_color,
                            [(self.x0 - 15, self.y0 - 10), (self.x0 + 15, self.y0 - 10), (self.x0 + 15, self.y0 + 3),
                             (self.x0 - 15, self.y0 + 3)])

    def move(self, _game):
        if pygame.key.get_pressed()[self.keymap[2]]:
            self.x0 += 5
        if pygame.key.get_pressed()[self.keymap[0]]:
            self.x0 -= 5
        if pygame.key.get_pressed()[self.keymap[3]]:
            self.an -= 0.05
        if pygame.key.get_pressed()[self.keymap[4]]:
            self.an += 0.05
        self.power_up()

    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()])
    guns: typing.List[Gun] = field(default_factory=lambda: [
        Gun(), Gun(x0=300, keymap=[pygame.K_j, pygame.K_k, pygame.K_l, pygame.K_u, pygame.K_o], bg_color=GREEN)
    ])
    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())
        for gun in self.guns:
            gun.move(self)

    def draw(self, screen):
        screen.fill(WHITE)
        for gun in self.guns:
            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.MOUSEBUTTONUP:
            for gun in self.guns:
                gun.targetting(event)
                gun.fire2_end(self)
        for gun in self.guns:
            gun.handle_event(event, self)

    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()