pybeamshow/effects/effect.py

from dataclasses import dataclass
import math
import random
from typing import Generator, Literal, Tuple
import pygame as pg


def copy_color(source: pg.Color) -> pg.Color:
    return pg.Color(source.r, source.g, source.b, source.a)


@dataclass(frozen=True, slots=True)
class Colors:
    Black = pg.Color(0, 0, 0)
    White = pg.Color(255, 255, 255)
    Red = pg.Color(255, 0, 0)
    Green = pg.Color(0, 255, 0)
    Blue = pg.Color(0, 0, 255)
    Cyan = pg.Color(0, 255, 255)
    Yellow = pg.Color(255, 255, 0)
    Magenta = pg.Color(255, 0, 255)


def color_wheel(hue=0, increase=1) -> Generator[pg.Color, None, None]:
    color = copy_color(Colors.Red)
    h, s, l, a = color.hsla
    h = hue

    while True:
        color.hsla = h, s, l, a
        yield color
        h = (h + increase) % 360


def color_randomize() -> Generator[pg.Color, None, None]:
    color = copy_color(Colors.Red)
    h, s, l, a = color.hsla
    color.hsla = random.randrange(0, 360 // 5) * 5, s, l, a

    while True:
        yield color
        color.hsla = random.randrange(0, 360 // 5) * 5, s, l, a


def color_fadeout(
    initial_color: pg.Color, fade_speed: float
) -> Generator[pg.Color, None, None]:
    color = copy_color(initial_color)
    h, s, l, a = color.hsla
    l_float = float(l)

    while True:
        yield color
        color.hsla = h, s, int(l_float), a
        l_float -= fade_speed
        if l_float < 0:
            l_float = 0


def color_fade(
    initial_color: pg.Color, end_color: pg.Color, duration: int
) -> Generator[pg.Color, None, None]:
    color = copy_color(initial_color)
    h, s, l, a = color.hsla
    h2, s2, l2, a2 = end_color.hsla
    h_inc = (h2 - h) / duration
    s_inc = (s2 - s) / duration
    l_inc = (l2 - l) / duration
    a_inc = (a2 - a) / duration
    h_f = float(h)
    s_f = float(s)
    l_f = float(l)
    a_f = float(a)

    while True:
        yield color
        color.hsla = int(h_f), int(s_f), int(l_f), int(a_f)
        if h_f < h2:
            h_f += h_inc
        if s_f < s2:
            s_f += s_inc
        if l_f < l2:
            l_f += l_inc
        if a_f < a2:
            a_f += a_inc


def color_darken(color: pg.Color, factor: float) -> pg.Color:
    h, s, l, a = color.hsla
    new_color = pg.Color(0, 0, 0, 255)
    new_color.hsla = h, s, l * factor, a
    return new_color


def rainbow_surface(
    image: pg.Surface,
    orientation: Literal["h", "v"] = "h",
    hue: int = 0,
    increase: int = 1,
):
    wheel = color_wheel(hue, increase)
    if orientation == "h":
        h = image.get_height()
        for x in range(image.get_width()):
            pg.draw.line(image, next(wheel), (x, 0), (x, h))
    elif orientation == "v":
        w = image.get_width()
        for y in range(image.get_height()):
            pg.draw.line(image, next(wheel), (0, y), (w, y))


def transform_bounce(
    bounds: pg.Rect,
    velocity: Tuple[float, float],
    x_factor: Tuple[float, float],
    y_factor: Tuple[float, float],
) -> Generator[Tuple[int, int], Tuple[int, int], None]:
    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])

    size_x, size_y = yield (bounds.centerx, bounds.centery)
    while True:
        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
        )

        phase += current_velocity / 180 * math.pi
        # current_velocity = random.uniform(min_velocity, max_velocity)
        size_x, size_y = yield (pos_x, pos_y)


def transform_oscillate(
    bounds: pg.Rect,
    period: int,
    initial_pos: Tuple[int, int] = (-1, -1),
) -> Generator[Tuple[int, int], Tuple[int, int], None]:
    pos_x = float(initial_pos[0] if initial_pos[0] > 0 else bounds.left)
    pos_y = float(initial_pos[1] if initial_pos[1] > 0 else bounds.top)
    direction = "+"

    size_x, size_y = yield (bounds.left, bounds.top)

    while True:
        range_x = bounds.width - size_x
        range_y = bounds.height - size_y

        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 = "+"

        size_x, size_y = yield (int(pos_x), int(pos_y))


def transform_falling(
    bounds: pg.Rect,
    acceleration: float,
    initial_pos: Tuple[int, int],
    initial_velocity: float = 1.0,
) -> Generator[Tuple[int, int], Tuple[int, int], None]:
    pos_x = float(initial_pos[0])
    pos_y = float(initial_pos[1])

    size_x, size_y = yield (bounds.left, bounds.top)

    velocity = initial_velocity

    while True:
        range_y = bounds.height - size_y

        pos_y += velocity
        velocity += acceleration

        if pos_y > range_y:
            pos_y = initial_pos[1]
            velocity = initial_velocity

        size_x, size_y = yield (int(pos_x), int(pos_y))


class Effect(pg.sprite.Sprite):
    def __init__(
        self, image: pg.Surface, rect: pg.Rect, *groups: pg.sprite.Group
    ) -> None:
        super().__init__(*groups)
        self.rect = rect
        self.image = image

    def draw(self, surface: pg.Surface):
        surface.blit(self.image, self.rect)