pybeamshow/effects/doublespot.py

from typing import Any, List, Tuple
import pygame as pg
from effects.effect import Effect, Colors
import random
import math
from typing import Union, Generator


def fade_statemachine(fade_in: int, hold: int, fade_out: int) -> Tuple[int, bool, bool]:
    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: Tuple[int, int],
        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: Union[pg.Color, Generator[pg.Color, None, None]],
        radius: int = 200,
        hold: int = 60 * 1,
        fade_out: bool = False,
        fade_in: 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.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=image, rect=bounds, *groups)

        self.update()

    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 len(self.spots) == 0 or (
            self.fade_in_time != 0 and self.spots[0].start_next and len(self.spots) == 2
        ):
            spot_color = (
                self.color if isinstance(self.color, pg.Color) else 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]