mipt_labs/lab5/gun.py

from __future__ import annotations
import math
from copy import deepcopy
import time
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
pygame.init()
font = pygame.font.Font('IntroRust.otf', 25)
font2 = pygame.font.Font('IntroRust.otf', 16)


@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
    source: int = 0

    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:
    gun_id: int = field(default_factory=lambda: rnd(1, 1000000))
    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)
        new_ball.source = self.gun_id
        game.balls.append(new_ball)
        self.f2_on = False
        self.f2_power = 10

    def targetting(self, event):
        """Прицеливание. Зависит от положения мыши."""
        if event and event.pos[0] != self.x0:
            self.an = math.atan2((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

    def hittest(self, obj):
        if self.x0 - 50 - obj.r < obj.position[0] < self.x0 + 50 + obj.r \
                and self.y0 - obj.r < obj.position[1] < self.y0 + 35 + obj.r \
                and hasattr(obj, 'source') and obj.source != self.gun_id:
            return True
        return False


@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
    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, game: Game, points=1):
        """Попадание шарика в цель."""
        game.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
    points: int = 0
    screen: pygame.Surface = field(default_factory=lambda: pygame.display.set_mode((WIDTH, HEIGHT)))

    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 gun in self.guns:
                if gun.hittest(ball):
                    self.guns.remove(gun)
            for target in self.targets:
                if ball.hittest(target) and target.live:
                    target.live = 0
                    ball.live = -1
                    target.hit(self)
                    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)
        score_text = font.render(str(self.points), 1, BLACK)
        screen.blit(score_text, (10, 10))
        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):
        while not self.finished:
            self.draw(self.screen)
            self.clock.tick(self.fps)
            for event in pygame.event.get():
                self.process_event(event)
            self.move()
        pygame.quit()


rules_text = '''
Управление левой пушкой:
a - влево, d - вправо, qe - наведение, s - стрелять
Управление правой:
jl - движение, uo - наведение, k - стрелять

Некоторые снаряды разрывные - разбиваются на 2.
Цели бывают обычные (те, которые полностью красные) - они 
движутся, просто отражаясь от стенок, - 
а также те, которые движутся случайно, и те, которые вращаются

Кликните, чтобы начать
'''
game = Game()
y = 100
game.screen.fill(0x2e2646)
for line in rules_text.split('\n'):
    surface = font2.render(line, True, WHITE)
    game.screen.blit(surface, (WIDTH / 2 - surface.get_width() / 2, y))
    y += 5 + surface.get_height()
rules_surface = font2.render(rules_text, True, WHITE)
pygame.display.update()
while not any(ev.type == pygame.MOUSEBUTTONUP or ev.type == pygame.QUIT for ev in pygame.event.get()):
    pass
game.main_loop()