pybeamshow/effects/drops.py

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


class Drop(pg.sprite.Sprite):
    def __init__(
        self,
        bounds: pg.rect.Rect,
        color: ColorGenerator,
        acceleration=1.0,
    ) -> None:

        self.bounds = bounds
        self.image = pg.Surface((self.bounds.width, self.bounds.width))
        self.rect = pg.rect.Rect(
            self.bounds.topleft, (self.bounds.width, self.bounds.width)
        )
        super().__init__()

        self.image.fill(Colors.Black)
        self.image.set_colorkey(Colors.Black)

        self.color = color
        pg.draw.ellipse(
            self.image,
            next(self.color),
            ((0, 0), self.rect.size),
        )

        self.fall = transform_falling(
            bounds=self.bounds,
            acceleration=acceleration,
            initial_pos=(self.bounds.left, self.bounds.top),
            initial_velocity=0,
        )
        next(self.fall)
        self.finished = False
        self.update(is_beat=False)

    def update(self, *args, **kwargs):
        if isinstance(self.color, Generator):
            pg.draw.ellipse(
                self.image,
                next(self.color),
                ((0, 0), self.rect.size),
            )
        if not self.finished:
            self.rect.topleft = self.fall.send(
                ((self.bounds.width, self.bounds.width), kwargs["is_beat"])
            )
        if self.rect.bottom >= self.bounds.bottom - self.bounds.width:
            self.finished = True

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


class Drops(Effect):
    def __init__(
        self,
        bounds: pg.Rect,
        colors: Tuple[
            ColorGenerator,
            ColorGenerator,
        ],
        sizes=(10, 100),
        drop_rate=0.2,
        drop_acceleration=1,
        *groups: pg.sprite.Group
    ) -> None:
        self.min_size = sizes[0]
        self.max_size = sizes[1]
        self.drop_size = random.randint(self.min_size, self.max_size)
        self.drop_rate = drop_rate
        self.drop_acceleration = drop_acceleration
        self.bounds = bounds
        self.colors = colors

        image = pg.Surface(self.bounds.size)
        image.fill(Colors.Black)
        image.set_colorkey(Colors.Black)

        super().__init__(image, self.bounds, *groups)

        self.drops: List[Drop] = []
        self.num_drops = int(self.bounds.width / (self.drop_size * 2))

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

        if len(self.drops) < self.num_drops:
            spawn_new = random.uniform(0, 1) < self.drop_rate

            if spawn_new:
                initial_color = (
                    self.colors[0]
                    if isinstance(self.colors[0], pg.Color)
                    else next(self.colors[0])
                )
                end_color = (
                    self.colors[1]
                    if isinstance(self.colors[1], pg.Color)
                    else next(self.colors[1])
                )
                new_drop = Drop(
                    pg.rect.Rect(
                        random.randrange(0, self.bounds.width - self.drop_size),
                        0,
                        self.drop_size,
                        self.bounds.height,
                    ),
                    color=color_fade(initial_color, end_color, random.randint(10, 120)),
                    acceleration=self.drop_acceleration,
                )
                self.drops.append(new_drop)

        for drop in self.drops:
            if not drop.finished:
                drop.update(is_beat=kwargs["is_beat"])
                drop.draw(self.image)

        if self.drops:
            first_drop = self.drops[0]
            while self.drops and first_drop.finished:
                first_drop = self.drops.pop(0)
                del first_drop
                if self.drops:
                    first_drop = self.drops[0]