pybeamshow/effects/rotatingpoly.py

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


def calc_radii(size, f):
    r_c = size / (5 * f + 2)
    r_a = size / 2 - r_c
    return r_c, r_a


class RotatingPoly(MovingEffect):
    def __init__(
        self,
        bounds: pg.Rect,
        color: ColorGenerator,
        beat_color: bool = False,
        size: int = 100,
        rot_speed: float = 5,
        outer: int = 5,
        spot_factor: float = 0.5,
        mover: Optional[PositionGenerator] = None,
        *groups: pg.sprite.Group
    ) -> None:

        self.rot_speed = rot_speed
        self.rotation = 0.0
        self.color = color
        self.beat_color = beat_color
        self.o_count = outer
        self.spot_factor = spot_factor

        self.o_f = 1 / math.sin(math.pi / self.o_count)

        self.spot_radius, self.o_radius = calc_radii(size, self.o_f)

        image = pg.Surface((size, size))
        image.fill(Colors.Black)
        image.set_colorkey(Colors.Black)
        super().__init__(
            image,
            pg.Rect(
                bounds.centerx - size / 2,
                bounds.centery - size / 2,
                size,
                size,
            ),
            mover,
            *groups
        )
        self.bounds = bounds
        self.o_color = next(self.color)
        self.update(is_beat=False)

    def update(self, *args: Any, **kwargs: Any) -> None:
        self.rect.center = self.mover.send((self.rect.size, kwargs["is_beat"]))

        self.image.fill(Colors.Black)

        if not self.beat_color or kwargs["is_beat"]:
            if isinstance(self.color, Generator):
                self.o_color = next(self.color)

        node_positions = []
        for phi in range(0, 360, 360 // self.o_count):
            rads = math.pi * (phi + self.rotation) / 180
            pos_x = int(self.rect.width // 2 + self.o_radius * math.cos(rads))
            pos_y = int(self.rect.height // 2 + self.o_radius * math.sin(rads))
            node_positions.append((pos_x, pos_y))

        pg.draw.polygon(
            self.image, self.o_color, node_positions, width=int(self.spot_radius * 1.5)
        )

        for pos_x, pos_y in node_positions:
            pg.draw.circle(
                self.image,
                Colors.White,
                (pos_x, pos_y),
                self.spot_radius * self.spot_factor,
            )

        self.rotation += self.rot_speed