ship.py

"""Spaceships, shooting through space."""

from __future__ import annotations

import math
import random
from enum import Enum
from re import S
from typing import TYPE_CHECKING

import pygame
from pygame import Color
from pygame.math import Vector2 as Vec2

from physics import Disk
from projectiles import Bullet, Rocket, Missile

from constants import (
    BULLET_SPEED,
    DAMAGE_INDICATOR_TIME,
    ENEMY_BULLET_COOLDOWN,
    GUNBARREL_LENGTH,
    GUNBARREL_WIDTH,
    ENEMY_SHOOT_RANGE,
    ENEMY_THRUST_MULTIPLIER,
    ENEMY_ACTION_TIMER,
    ENEMY_HEALTH,
    ENEMY_ACTION_WEIGHTS,
    ENEMY_ROCKET_COOLDOWN,
    GUN_COOLDOWN,
    ENEMY_VISUAL_RANGE,
    WORLD_SIZE,
    SMOL,
    ENEMY_MISSILE_COOLDOWN,
)

if TYPE_CHECKING:
    from camera import Camera


class Ship(Disk):
    """A basic spaceship."""

    def __init__(
        self,
        pos: Vec2,
        vel: Vec2,
        density: float,
        size: float,
        color: Color,
        bullet_color: Color,
    ) -> None:
        """Create a new spaceship.

        Args:
        ----
            pos (Vec2): Initial position
            vel (Vec2): Initial velocity
            density (float): Density (of disk-body)
            size (float): Radius of disk-body
            color (Color): Material color
            bullet_color (Color): Bullet_color

        """
        super().__init__(pos, vel, density, size, color)
        self.size: float = size
        self.angle: float = 0
        self.health: float = 100.0
        self.projectiles: list[Bullet] = []
        self.gun_cooldown: float = 0
        # This just just determines cooldown once at the very start.
        self.has_trophy: bool = False
        self.bullet_color = Color(bullet_color)
        self.ammo: int = 3700
        self.thrust: float = 250 * self.mass
        self.rotation_thrust: float = 230
        self.thruster_rot_left: bool = False
        self.thruster_rot_right: bool = False
        self.thruster_backward: bool = False
        self.thruster_forward: bool = False
        self.max_fuel: float = 100.0
        self.fuel: float = self.max_fuel
        self.fuel_consumption_rate: float = 0.5
        self.fuel_rot_consumption_rate: float = 0.5
        self.damage_indicator_timer: float = 0

    def get_faced_direction(self) -> Vec2:
        """Get `self`'s faced direction from its `angle`.

        Returns
        -------
            Vec2: Faced direction, normalized

        """
        # For unknown reasons, `Vec2.from_polar((self.angle, 1))` won't work.
        direction = Vec2()
        direction.from_polar((1, self.angle))
        return direction

    def shoot(self) -> None:
        """Try to shoot a bullet."""
        if self.gun_cooldown <= 0 and self.ammo > 0:
            forward = self.get_faced_direction()
            bullet_pos = self.pos + forward * self.radius * GUNBARREL_LENGTH
            bullet_vel = self.vel + forward * BULLET_SPEED
            self.projectiles.append(Bullet(bullet_pos, bullet_vel, self.bullet_color))
            self.gun_cooldown = GUN_COOLDOWN
            self.ammo -= 1

    def suffer_damage(self, damage: float) -> None:
        """Deal damage to the ship and activate its damage-indicator.

        Does nothing if damage is <= 0.

        Args:
        ----
            damage (float): Amount of damage to deal.

        """
        if damage > 0:
            self.health -= damage
            self.damage_indicator_timer = DAMAGE_INDICATOR_TIME

    def step(self, dt: float) -> None:
        """Physics, control, and bullet-stepping for `self`.

        Args:
        ----
            dt (float): Passed time

        """
        if self.fuel > 0:
            if self.thruster_rot_left:
                self.fuel = max(0, self.fuel - dt * self.fuel_rot_consumption_rate)
                self.angle += self.rotation_thrust * dt
            if self.thruster_rot_right:
                self.fuel = max(0, self.fuel - dt * self.fuel_rot_consumption_rate)
                self.angle -= self.rotation_thrust * dt

            forward = self.get_faced_direction()
            if self.thruster_forward:
                self.fuel = max(0, self.fuel - dt * self.fuel_consumption_rate)
                self.apply_force(forward * self.thrust, dt)
            if self.thruster_backward:
                self.fuel = max(0, self.fuel - dt * self.fuel_consumption_rate)
                self.apply_force(-forward * self.thrust, dt)

        self.damage_indicator_timer = max(0, self.damage_indicator_timer - dt)

        super().step(dt)

        for projectile in self.projectiles:
            projectile.step(dt)

        self.gun_cooldown = max(0, self.gun_cooldown - dt)

    def draw(self, camera: Camera) -> None:
        """Draw `self` on `camera.

        Args:
        ----
            camera (Camera): Camera to draw on

        """
        forward = self.get_faced_direction()
        right = Vec2(-forward.y, forward.x)
        left = -right
        backward = -forward

        base_color: Color = self.color.lerp(Color("red"), self.damage_indicator_timer)
        darker_color: Color = base_color.lerp(Color("black"), 0.5)

        # Helper function for drawing polygons relative to the ship-position
        def drawy(color: Color, points: list[Vec2]) -> None:
            camera.draw_polygon(color, [self.pos + self.radius * p for p in points])

        # thruster_backward (active)
        if self.thruster_backward:
            drawy(Color("orange"), [forward * 2, left * 1.25, right * 1.25])

        # "For his neutral special, he wields a gun"
        camera.draw_line(
            darker_color,
            self.pos,
            self.pos + forward * self.radius * GUNBARREL_LENGTH,
            GUNBARREL_WIDTH * self.radius,
        )

        # thruster_rot_left (material)
        drawy(
            darker_color,
            [
                0.7 * left + 0.7 * forward,
                0.5 * left + 0.5 * backward,
                2.0 * left + 1.0 * backward,
            ],
        )
        # thruster_rot_left (active)
        if self.thruster_rot_left:
            drawy(
                Color("orange"),
                [
                    1.5 * left + 1.25 * backward,
                    0.5 * left + 0.5 * backward,
                    2.0 * left + 1.0 * backward,
                ],
            )

        # thruster_rot_right (material)
        drawy(
            darker_color,
            [
                0.7 * right + 0.7 * forward,
                0.5 * right + 0.5 * backward,
                2.0 * right + 1.0 * backward,
            ],
        )
        # thruster_rot_right (active)
        if self.thruster_rot_right:
            drawy(
                Color("orange"),
                [
                    1.5 * right + 1.25 * backward,
                    0.5 * right + 0.5 * backward,
                    2.0 * right + 1.0 * backward,
                ],
            )

        # thruster_forward (flame)
        if self.thruster_forward:
            drawy(
                Color("orange"),
                [
                    0.7 * left + 0.7 * backward,
                    0.5 * left + 1.5 * backward,
                    1.25 * backward,
                    0.5 * right + 1.5 * backward,
                    0.7 * right + 0.7 * backward,
                ],
            )
        # thruster_forward (material)
        drawy(
            darker_color,
            [
                0.7 * left + 0.7 * backward,
                0.5 * left + 1.25 * backward,
                1.0 * backward,
                0.5 * right + 1.25 * backward,
                0.7 * right + 0.7 * backward,
            ],
        )

        # Ugly hack
        backup_self_color = Color(self.color)
        self.color = base_color
        super().draw(camera)  # Draw circular body ("hitbox")
        self.color = backup_self_color

        for projectile in self.projectiles:
            projectile.draw(camera)


class ShipInput:
    """Specification for which keys trigger what spaceship-action."""

    type PygameKey = int

    def __init__(
        self,
        thruster_rot_left: PygameKey,
        thruster_rot_right: PygameKey,
        thruster_forward: PygameKey,
        thruster_backward: PygameKey,
        shoot: PygameKey,
    ) -> None:
        """Create a new map from keys to spaceship-actions.

        Args:
        ----
            thruster_rot_left (pygame_key): Left rotation thruster's key
            thruster_rot_right (pygame_key): Right rotation thruster's key
            thruster_forward (pygame_key): Forward thruster's key
            thruster_backward (pygame_key): Backward thruster's key
            shoot (pygame_key): Pew pew key

        """
        self.thruster_rot_left = thruster_rot_left
        self.thruster_rot_right = thruster_rot_right
        self.thruster_forward = thruster_forward
        self.thruster_backward = thruster_backward
        self.shoot = shoot


class PlayerShip(Ship):
    """A player-controlled spaceship."""

    def __init__(
        self,
        pos: Vec2,
        vel: Vec2,
        density: float,
        size: float,
        color: Color,
        bullet_color: Color,
        spaceship_input: ShipInput,
        image_path: str,
    ) -> None:
        """Create a new player-spaceship.

        Args:
        ----
            pos (Vec2): Initial position
            vel (Vec2): Initial velocity
            density (float): Density (of disk-body)
            size (float): Radius of disk-body
            color (Color): Material color
            bullet_color (Color): Bullet_color
            spaceship_input (SpaceshipInput): Map from keys to actions
            image_path (str): Path to image

        """
        super().__init__(pos, vel, density, size, color, bullet_color)
        self.spaceship_input = spaceship_input
        self.image = pygame.image.load(image_path)

    def handle_input(self, keys: pygame.key.ScancodeWrapper) -> None:
        """Handle input for `self` using ScancodeWrapper `keys`.

        `keys` is typically retreived using `pygame.key.get_pressed()`

        Args:
        ----
            keys (pygame.key.ScancodeWrapper): Pressed keys

        """
        self.thruster_rot_left = keys[self.spaceship_input.thruster_rot_left]
        self.thruster_rot_right = keys[self.spaceship_input.thruster_rot_right]
        self.thruster_forward = keys[self.spaceship_input.thruster_forward]
        self.thruster_backward = keys[self.spaceship_input.thruster_backward]
        if keys[self.spaceship_input.shoot]:
            self.shoot()

    def draw(self, camera: Camera) -> None:
        """Draw `self` on `camera.

        Args:
        ----
            camera (Camera): Camera to draw on

        """
        forward = self.get_faced_direction()

        right = Vec2(-forward.y, forward.x)
        left = -right
        backward = -forward

        base_color: Color = self.color.lerp(Color("red"), self.damage_indicator_timer)
        darker_color: Color = base_color.lerp(Color("black"), 0.5)

        # Helper function for drawing polygons relative to the ship-position
        def drawy(color: Color, points: list[Vec2]) -> None:
            camera.draw_polygon(color, [self.pos + self.radius * p for p in points])

        # thruster_backward (active)
        if self.thruster_backward:
            drawy(Color("orange"), [forward * 2, left * 1.25, right * 1.25])

        # "For his neutral special, he wields a gun"
        camera.draw_line(
            darker_color,
            self.pos,
            self.pos + forward * self.radius * GUNBARREL_LENGTH,
            GUNBARREL_WIDTH * self.radius,
        )

        # thruster_rot_left (material)
        drawy(
            darker_color,
            [
                0.7 * left + 0.7 * forward,
                0.5 * left + 0.5 * backward,
                2.0 * left + 1.0 * backward,
            ],
        )
        # thruster_rot_left (active)
        if self.thruster_rot_left:
            drawy(
                Color("orange"),
                [
                    1.5 * left + 1.25 * backward,
                    0.5 * left + 0.5 * backward,
                    2.0 * left + 1.0 * backward,
                ],
            )

        # thruster_rot_right (material)
        drawy(
            darker_color,
            [
                0.7 * right + 0.7 * forward,
                0.5 * right + 0.5 * backward,
                2.0 * right + 1.0 * backward,
            ],
        )
        # thruster_rot_right (active)
        if self.thruster_rot_right:
            drawy(
                Color("orange"),
                [
                    1.5 * right + 1.25 * backward,
                    0.5 * right + 0.5 * backward,
                    2.0 * right + 1.0 * backward,
                ],
            )

        # thruster_forward (flame)
        if self.thruster_forward:
            drawy(
                Color("orange"),
                [
                    0.7 * left + 0.7 * backward,
                    0.5 * left + 1.5 * backward,
                    1.25 * backward,
                    0.5 * right + 1.5 * backward,
                    0.7 * right + 0.7 * backward,
                ],
            )
        # thruster_forward (material)
        drawy(
            darker_color,
            [
                0.7 * left + 0.7 * backward,
                0.5 * left + 1.25 * backward,
                1.0 * backward,
                0.5 * right + 1.25 * backward,
                0.7 * right + 0.7 * backward,
            ],
        )

        rotated_image = pygame.transform.rotate(self.image, -self.angle - 90)

        # Get the width and height of the rotated image
        image_rect = rotated_image.get_rect()
        center_offset = Vec2(image_rect.width / 2, image_rect.height / 2)

        # Adjust the position to center the image
        adjusted_pos = self.pos - center_offset

        # Ugly hack
        backup_self_color = Color(self.color)
        self.color = base_color
        super().draw(camera)  # Draw circular body ("hitbox")
        self.color = backup_self_color

        for projectile in self.projectiles:
            projectile.draw(camera)

        camera.draw_image(rotated_image, adjusted_pos)


class BulletEnemy(Ship):
    """An enemy ship, targeting a specific other ship."""

    Action = Enum(
        "Action",
        [
            "accelerate_to_player",
            "accelerate_randomly",
            "decelerate",
        ],
    )

    def __init__(
        self,
        pos: Vec2,
        vel: Vec2,
        target_ship: Ship,
        color: Color = Color("lime"),
        bullet_color: Color = Color("hotpink"),
    ) -> None:
        """Create a new enemy ship.

        Args:
        ----
            pos (Vec2): Initial position
            vel (Vec2): Initial velocity
            target_ship (Ship): Ship to target
            color (Color, optional): Material color. Defaults to Color("purple").
            bullet_color (Color): Color of shot projectiles

        """
        super().__init__(pos, vel, 1, 8, color, bullet_color)
        self.thrust *= ENEMY_THRUST_MULTIPLIER
        self.action_timer = 0
        self.health = ENEMY_HEALTH
        self.current_action: BulletEnemy.Action = (
            BulletEnemy.Action.accelerate_randomly
        )
        self.target_ship = target_ship
        self.projectiles: list[Bullet] = []
        self.random_point = Vec2(0, 0)

    def step(self, dt: float) -> None:
        """Apply physics and "AI" to `self`.

        Args:
        ----
            dt (float): Passed time

        """
        self.action_timer -= dt
        delta_target_ship = self.target_ship.pos - self.pos

        if self.action_timer <= 0:
            self.random_point = Vec2(
                random.uniform(0, WORLD_SIZE.x), random.uniform(0, WORLD_SIZE.y)
            )
            # print(self.current_action)
            # print(self.random_point)
            if delta_target_ship == Vec2(0, 0):
                delta_target_ship = Vec2(SMOL, SMOL)

            if delta_target_ship.magnitude_squared() < ENEMY_VISUAL_RANGE**2:
                self.current_action = BulletEnemy.Action.accelerate_to_player
            else:
                [self.current_action] = random.choices(
                    population=[
                        BulletEnemy.Action.accelerate_randomly,
                        BulletEnemy.Action.decelerate,
                    ],
                    weights=ENEMY_ACTION_WEIGHTS,
                )

            self.action_timer = ENEMY_ACTION_TIMER


        match self.current_action:
            case BulletEnemy.Action.accelerate_to_player:
                desired_velocity = (
                    delta_target_ship * self.thrust / delta_target_ship.magnitude()
                )
                perfect_multiplier = max(
                    self.target_ship.vel.magnitude() * 1.5, desired_velocity.magnitude()
                )
                perfect_velocity = desired_velocity.normalize() * perfect_multiplier
                required_acceleration = perfect_velocity - self.vel
                force_direction = required_acceleration
            case BulletEnemy.Action.decelerate:
                force_direction = -self.vel
            case BulletEnemy.Action.accelerate_randomly:
                delta_random_point = self.random_point - self.pos
                force_direction = delta_random_point

        if force_direction.magnitude() != 0:
            force = force_direction * self.thrust / force_direction.magnitude()
            self.apply_force(force, dt)

        super().step(dt)

        # Shooting logic
        if (
            delta_target_ship.magnitude_squared() < ENEMY_SHOOT_RANGE**2
        ):
            self.shoot()

    def shoot(self) -> None:
        """Shoot a bullet."""
        if (
            self.gun_cooldown <= 0
            and self.ammo > 0
            and self.current_action == BulletEnemy.Action.accelerate_to_player
        ):
            forward = self.get_faced_direction()
            bullet_pos = self.pos + forward * self.radius * GUNBARREL_LENGTH
            bullet_vel = self.vel + forward * BULLET_SPEED
            self.projectiles.append(Bullet(bullet_pos, bullet_vel, self.bullet_color))
            self.gun_cooldown = ENEMY_BULLET_COOLDOWN
            self.ammo -= 1


class RocketEnemy(BulletEnemy):
    """An enemy ship shooting rockets, targeting a specific other ship."""

    def __init__(
        self,
        pos: Vec2,
        vel: Vec2,
        target_ship: Ship,
        color: Color = Color("plum"),
    ) -> None:
        """Create a new Rocket-Ship.

        Args:
        ----
            pos (Vec2): Initial position
            vel (Vec2): Initial velocity
            target_ship (Ship): Ship to target
            color (Color, optional): Material color. Defaults to Color("red").

        """
        super().__init__(pos, vel, target_ship, color)

    def shoot(self) -> None:
        """Shoot a Rocket."""
        if (
            self.gun_cooldown <= 0
            and self.ammo > 0
            and self.current_action == BulletEnemy.Action.accelerate_to_player
        ):
            forward = self.get_faced_direction()
            bullet_pos = self.pos + forward * self.radius * GUNBARREL_LENGTH
            bullet_vel = self.vel
            self.projectiles.append(
                Rocket(bullet_pos, bullet_vel, self.color, self.target_ship),
            )
            self.gun_cooldown = ENEMY_ROCKET_COOLDOWN
            self.ammo -= 1


class MissileEnemy(BulletEnemy):
    """An enemy ship shooting powerful, smart, homing missiles, targeting a specific other ship."""

    def __init__(
        self,
        pos: Vec2,
        vel: Vec2,
        target_ship: Ship,
        color: Color = Color("blue"),
    ) -> None:
        """Create a new Missile-Ship.

        Args:
        ----
            pos (Vec2): Initial position
            vel (Vec2): Initial velocity
            target_ship (Ship): Ship to target
            color (Color, optional): Material color. Defaults to Color("red").

        """
        super().__init__(pos, vel, target_ship, color)

    def shoot(self) -> None:
        """Shoot a smart Missile."""
        if (
            self.gun_cooldown <= 0
            and self.ammo > 0
            and self.current_action == BulletEnemy.Action.accelerate_to_player
        ):
            forward = self.get_faced_direction()
            bullet_pos = self.pos + forward * self.radius * GUNBARREL_LENGTH
            bullet_vel = self.vel
            self.projectiles.append(
                Missile(
                    bullet_pos,
                    bullet_vel,
                    self.color,
                    self.target_ship,
                    "assets/missile.png",
                ),
            )
            self.gun_cooldown = ENEMY_MISSILE_COOLDOWN
            self.ammo -= 1