pybeamshow/effects/doublespot.py

from typing import Any, List, Tuple
import pygame as pg
from effects.effect import Effect
from util import XYCoord
from util.color import Colors, ColorGenerator
import random
from typing import Generator


def fade_statemachine(
    fade_in: int, hold: int, fade_out: int
) -> Generator[Tuple[int, bool, bool], None, None]:
    for t in range(fade_in):
        yield 255 * t // fade_in, False, False
    for t in range(hold):
        yield 255, False, False
    for t in range(fade_out):
        yield 255 * (fade_out - t) // fade_out, True, False
    yield 0, True, True


class Spot(pg.sprite.Sprite):
    def __init__(
        self,
        color: pg.Color,
        position: XYCoord,
        radius: int,
        fade_in: int,
        hold: int,
        fade_out: int,
    ) -> None:
        super().__init__()
        self.rect = pg.Rect(
            position[0] - radius, position[1] - radius, radius * 2, radius * 2
        )
        self.image = pg.Surface(self.rect.size)
        self.image.set_colorkey(Colors.Black)
        self.image.fill(Colors.Black)
        pg.draw.ellipse(
            self.image,
            color,
            ((0, 0), self.rect.size),
        )
        self.start_next = False
        self.finished = False
        self.state = fade_statemachine(fade_in, hold, fade_out)

    def update(self) -> None:
        if not self.finished:
            alpha, self.start_next, self.finished = next(self.state)
            self.image.set_alpha(alpha)

    def draw(self, dest: pg.Surface) -> None:
        dest.blit(self.image, self.rect)


class DoubleSpot(Effect):
    def __init__(
        self,
        bounds: pg.Rect,
        color: ColorGenerator,
        radius: int = 200,
        hold: int = 60 * 1,
        fade_out: bool = False,
        fade_in: bool = False,
        beat_adapt: bool = False,
        *groups: pg.sprite.Group
    ) -> None:
        self.color = color
        self.radius: int = radius
        self.fade_in_time: int = hold // 4 if fade_in else 0
        self.fade_out_time: int = hold // 4 if fade_out else 0
        self.hold_time: int = hold - (self.fade_in_time + self.fade_out_time)
        self.beat_adapt: bool = beat_adapt
        self.spots: List[Spot] = []

        self.randrange = (
            (bounds.width - 2 * self.radius) // (self.radius * 2),
            (bounds.height - 2 * self.radius) // (self.radius * 2),
        )

        image = pg.Surface(size=bounds.size)
        image.set_colorkey(Colors.Black)
        super().__init__(image, bounds, *groups)

        self.update(is_beat=False, frames_per_beat=0.0)

    def add_spot(self, spot_color):
        position = (
            random.randint(0, self.randrange[0]) * self.radius * 2 + self.radius,
            random.randint(0, self.randrange[1]) * self.radius * 2 + self.radius,
        )
        while any(
            (
                (s.rect.centerx == position[0] and s.rect.centery == position[1])
                for s in self.spots
            )
        ):
            position = (
                random.randint(0, self.randrange[0]) * self.radius * 2 + self.radius,
                random.randint(0, self.randrange[1]) * self.radius * 2 + self.radius,
            )

        self.spots.append(
            Spot(
                color=spot_color,
                position=position,
                radius=self.radius,
                fade_in=self.fade_in_time,
                hold=self.hold_time,
                fade_out=self.fade_out_time,
            )
        )

    def update(self, *args: Any, **kwargs: Any) -> None:
        self.image.fill(Colors.Black)

        if self.beat_adapt and kwargs["frames_per_beat"]:
            hold = kwargs["frames_per_beat"] - 2
            self.fade_in_time = int(hold // 4 if self.fade_in_time else 0)
            self.fade_out_time = int(hold // 4 if self.fade_out_time else 0)
            self.hold_time = int(hold - self.fade_in_time)

        if not self.beat_adapt:
            if len(self.spots) == 0 or (
                self.fade_in_time != 0
                and self.spots[0].start_next
                and len(self.spots) == 2
            ):
                spot_color = (
                    next(self.color),
                )
                self.add_spot(spot_color)
                self.add_spot(spot_color)
        elif kwargs["is_beat"]:
            spot_color = (
                next(self.color),
            )
            self.add_spot(spot_color)
            self.add_spot(spot_color)

        for spot in self.spots:
            if not spot.finished:
                spot.update()
                spot.draw(self.image)

        if self.spots:
            first_spot = self.spots[0]
            while self.spots and first_spot.finished:
                first_spot = self.spots.pop(0)
                del first_spot
                if self.spots:
                    first_spot = self.spots[0]