pybeamshow/effects/starfield.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 Starfield(Effect):
    def __init__(
        self,
        bounds: pg.Rect,
        color: ColorGenerator,
        radius: int = 20,
        star_factor: float = 0.2,
        hold: int = 60 * 3,
        fade_out: bool = False,
        fade_in: bool = False,
        beat_adapt: bool = False,
        *groups: pg.sprite.Group
    ) -> None:
        self.color = color
        self.radius: int = radius * 2
        self.fade_in_time: int = hold // 10 if fade_in else 0
        self.fade_out_time: int = int(hold * 0.8 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.stars: List[Spot] = []

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

        self.num_stars = int(self.randrange[0] * self.randrange[1] * star_factor)

        # self.time_to_next = (self.hold_time + self.fade_in_time) // 3
        self.time_to_next = hold // 5

        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_star(self, spot_color):
        loc_y = random.randint(0, self.randrange[1])
        pos_y = loc_y * self.radius * 2 + self.radius
        pos_x = (
            random.randint(0, self.randrange[0]) * self.radius * 2
            + self.radius
            + (self.radius if (loc_y & 1 != 0) else 0)
        )

        while any(
            ((s.rect.centerx == pos_x and s.rect.centery == pos_y) for s in self.stars)
        ):
            loc_y = random.randint(0, self.randrange[1])
            pos_y = loc_y * self.radius * 2 + self.radius
            pos_x = (
                random.randint(0, self.randrange[0]) * self.radius * 2
                + self.radius
                + (self.radius if (loc_y & 1 != 0) else 0)
            )

        self.stars.append(
            Spot(
                color=spot_color,
                position=(pos_x, pos_y),
                radius=self.radius // 2,
                fade_in=random.randint(0, self.fade_in_time * 2),
                hold=random.randint(0, self.hold_time * 2),
                fade_out=random.randint(0, self.fade_out_time * 2),
            )
        )

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

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

        if (
            self.time_to_next == 0 and (not self.beat_adapt or kwargs["is_beat"])
        ) and len(self.stars) < int(self.num_stars * 0.8):
            missing_stars = self.num_stars - len(self.stars)
            for _ in range(random.randint(missing_stars // 3, missing_stars // 2)):
                star_color = (
                    next(self.color),
                )
                self.add_star(star_color)

            self.time_to_next = hold // 5
        else:
            self.time_to_next -= 1
            if self.time_to_next < 0:
                self.time_to_next = 0

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

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