pybeamshow/effects/rainbowwave.py

from typing import Any, Optional
import pygame as pg
from effects.effect import MovingEffect
from util.color import Colors, rainbow_surface
from util.transform import PositionGenerator
import math


class RainbowWave(MovingEffect):
    def __init__(
        self,
        bounds: pg.Rect,
        wave_count: float = 2,
        wave_height: int = 0,
        thickness: int = 10,
        hue: int = 0,
        color_inc: int = 1,
        start_phase: int = 0,
        scroll_speed: float = 5,
        mover: Optional[PositionGenerator] = None,
        *groups: pg.sprite.Group
    ) -> None:
        image = pg.Surface(bounds.size)
        image.fill(Colors.Black)
        image.set_colorkey(Colors.Black)
        super().__init__(image, bounds, mover, *groups)
        self.subrect = pg.rect.Rect(0, 0, bounds.width, bounds.height)

        self.wave_height = wave_height or bounds.height
        self.bounds = bounds
        self.scroll_speed = scroll_speed

        self.rainbow = pg.Surface((bounds.width, self.wave_height))
        rainbow_surface(self.rainbow, "h", hue, color_inc)

        self.wave_image = pg.Surface(
            (bounds.width * (wave_count + 1) / wave_count, self.wave_height)
        )
        self.wave_image.fill(Colors.Black)
        self.wave_image.set_colorkey(Colors.Black)

        y_halfheight = self.wave_image.get_height() // 2
        points = [
            (
                x,
                y_halfheight
                + (y_halfheight - thickness)
                * math.sin(
                    math.pi
                    * (
                        start_phase / 180
                        + x / self.wave_image.get_width() * 2 * (wave_count + 1)
                    )
                ),
            )
            for x in range(self.wave_image.get_width())
        ]

        pg.draw.lines(self.wave_image, Colors.White, False, points, width=thickness)

        pg.draw.lines(
            self.wave_image,
            Colors.White,
            False,
            [(x + 2, y) for x, y in points],
            width=thickness,
        )

        self.update(is_beat=False)

    def update(self, *args: Any, **kwargs: Any) -> None:

        self.image.fill(Colors.Black)

        # rainbow_surface(self.image)

        pos = self.mover.send(
            ((self.bounds.width, self.wave_height), kwargs["is_beat"])
        )
        # pg.draw.rect(self.image, Colors.Black, (0, 0, self.bounds.width, pos[1]))
        # pg.draw.rect(
        #     self.image,
        #     Colors.Black,
        #     (
        #         0,
        #         pos[1] + self.wave_height,
        #         self.bounds.width,
        #         self.bounds.height - pos[1] + self.wave_height,
        #     ),
        # )

        self.image.blit(self.rainbow, pos)
        self.image.blit(self.wave_image, pos, self.subrect, pg.BLEND_MULT)

        self.subrect.move_ip(self.scroll_speed, 0)
        if self.subrect.right >= self.wave_image.get_rect().right:
            self.subrect.left = 0

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

        # self.ticks += int(self.velocity / 180 * math.pi)
        # self.velocity = random.randint(self.min_velocity, self.max_velocity)