pybeamshow/util/transform.py

import math
import random
from typing import Generator, Optional, Tuple
import pygame as pg

from util import XYCoord

PositionGenerator = Generator[XYCoord, Tuple[XYCoord, Optional[bool]], None]


def transform_static(position: XYCoord) -> PositionGenerator:
    while True:
        _ = yield (position)


def transform_bounce(
    bounds: pg.Rect,
    velocity: Tuple[float, float],
    x_factor: Tuple[float, float],
    y_factor: Tuple[float, float],
    on_beat_random_phase: int = 0,
    on_beat_phase: int = 0,
) -> PositionGenerator:
    min_velocity = velocity[0]
    max_velocity = velocity[1]
    current_velocity = random.uniform(min_velocity, max_velocity)
    phase = random.uniform(0, 360)
    current_x_factor = random.uniform(x_factor[0], x_factor[1])
    current_y_factor = random.uniform(y_factor[0], y_factor[1])

    pos_x = bounds.centerx
    pos_y = bounds.centery
    while True:
        (size_x, size_y), is_beat = yield (pos_x, pos_y)
        pos_x = int(
            math.cos(current_x_factor * phase) * (bounds.width - size_x) // 2
            + bounds.centerx
        )
        pos_y = int(
            math.sin(current_y_factor * phase) * (bounds.height - size_y) // 2
            + bounds.centery
        )

        inc = current_velocity
        if on_beat_random_phase and is_beat:
            inc += random.randrange(0, on_beat_random_phase)
        if on_beat_phase and is_beat:
            inc += on_beat_phase
        phase += inc / 180 * math.pi


def transform_oscillate(
    bounds: pg.Rect,
    period: int,
    initial_pos: XYCoord = (0, 0),
    auto_period: int = 0,
) -> PositionGenerator:
    pos_x = float(initial_pos[0])
    pos_y = float(initial_pos[1])
    direction = "+"

    last_period = 0

    while True:
        (size_x, size_y), is_beat = yield (
            int(pos_x + bounds.left),
            int(pos_y + bounds.top),
        )

        range_x = bounds.width - size_x
        range_y = bounds.height - size_y

        if auto_period:
            if is_beat:
                period = last_period * auto_period
                last_period = 0
            else:
                last_period += 1

        inc_x = range_x / period
        inc_y = range_y / period

        if direction == "+":
            pos_x = pg.math.clamp(pos_x + inc_x, 0, range_x)
            pos_y = pg.math.clamp(pos_y + inc_y, 0, range_y)
        else:
            pos_x = pg.math.clamp(pos_x - inc_x, 0, range_x)
            pos_y = pg.math.clamp(pos_y - inc_y, 0, range_y)

        if (inc_x and (pos_x > range_x - inc_x)) or (
            inc_y and (pos_y > range_y - inc_y)
        ):
            direction = "-"
        elif (inc_x and (pos_x < inc_x)) or (inc_y and (pos_y < inc_y)):
            direction = "+"


def transform_falling(
    bounds: pg.Rect,
    acceleration: float,
    initial_pos: XYCoord,
    initial_velocity: float = 1.0,
    on_beat_reset: bool = False,
) -> PositionGenerator:
    pos_x = float(initial_pos[0])
    pos_y = float(initial_pos[1])

    velocity = initial_velocity
    pos_x = bounds.left
    pos_y = bounds.top

    while True:
        (_, size_y), is_beat = yield (int(pos_x), int(pos_y))
        range_y = bounds.height - size_y

        pos_y += velocity
        velocity += acceleration

        if (not on_beat_reset and (pos_y > range_y)) or (on_beat_reset and is_beat):
            pos_y = initial_pos[1]
            velocity = initial_velocity