pybeamshow/beamshow.py

from argparse import ArgumentParser
import time
from typing import Generator, List, Tuple

import pygame as pg
import sys

from effects.effect import Effect, Colors, color_darken, color_fadeout, color_wheel
from effects.presets import Presets


def print_displays() -> None:
    pg.display.init()
    display = pg.display.get_desktop_sizes()
    for i, d in enumerate(display):
        print(f"#{i} - {d[0]}x{d[1]}")
    sys.exit(0)


class Beamshow:
    def __init__(
        self,
        render3d: bool,
        window: bool,
        trails: bool,
        randomize_interval: float,
        fps: bool,
        display: int,
        beat_reactive: bool,
    ) -> None:
        self.render3d = render3d
        self.windowed = window
        self.trails = trails
        self.randomize = bool(randomize_interval)
        self.randomize_interval = randomize_interval or 5
        self.print_fps = fps
        self.display_id = display
        self.effects: List[Effect] = []
        self.clock = pg.time.Clock()
        self.beat_reactive = beat_reactive

        pg.display.init()
        self.window, self.background = self.initialize()

    def initialize(self) -> Tuple[pg.Surface, pg.Surface]:
        displays = pg.display.get_desktop_sizes()
        if not 0 <= self.display_id < len(displays):
            raise ValueError(
                f"Display ID {self.display_id} invalid. Must be between 0 and {len(displays)}!"
            )

        win = pg.display.set_mode(
            size=displays[self.display_id]
            if not self.windowed
            else (displays[self.display_id][0] // 2, displays[self.display_id][1] // 2),
            flags=pg.FULLSCREEN if not self.windowed else 0,
            display=self.display_id,
        )
        if self.trails:
            background = pg.Surface(win.get_size(), flags=pg.SRCALPHA)
            background.fill(pg.Color(0, 0, 0, 5))
        else:
            background = pg.Surface(win.get_size())
            background.fill(Colors.Black)

        self.presets = Presets(bounds=win.get_rect(), beat_reactive=self.beat_reactive)
        if self.randomize:
            self.effects = self.presets.randomize()
        else:
            # self.effects = self.presets.default()
            from effects.spiro import Spiro

            self.effects = [
                Spiro(
                    bounds=win.get_rect(),
                    color=color_wheel(increase=10),
                    sizes=(win.get_rect().height * 0.8, win.get_rect().height * 0.8),
                )
            ]

        return win, background

    def render_loop_normal(
        self,
    ) -> Generator[List[pg.event.Event], Tuple[bool, float], None]:
        is_beat = False
        frames_per_beat = 0.0
        while True:
            unhandled_events: List[pg.event.Event] = []
            for event in pg.event.get():
                unhandled_events.append(event)

            self.window.blit(self.background, (0, 0))
            for e in self.effects:
                e.update(is_beat=is_beat, frames_per_beat=frames_per_beat)
                e.draw(self.window)
            is_beat, frames_per_beat = yield unhandled_events

    def render_loop_3d(
        self,
    ) -> Generator[List[pg.event.Event], Tuple[bool, float], None]:
        stage = pg.Surface(size=self.window.get_size())
        stage.fill(Colors.Black)

        full_size = stage.get_size()
        scaled_sizes = [
            (((full_size[0] * i) // 100), ((full_size[1] * i) // 100))
            for i in range(101)
        ]

        scaled_positions = [
            ((full_size[0] - s[0]) // 2, (full_size[1] - s[1]) // 2)
            for s in scaled_sizes
        ]

        is_beat = False
        frames_per_beat = 0.0
        while True:
            unhandled_events: List[pg.event.Event] = []
            for event in pg.event.get():
                unhandled_events.append(event)

            stage.blit(self.background, (0, 0))
            for e in self.effects:
                e.update(is_beat=is_beat, frames_per_beat=frames_per_beat)
                e.draw(stage)

            self.window.fill(Colors.Black)

            stage.set_colorkey(Colors.Black)
            for i in range(0, 101, 4):
                stage.set_alpha(192 - 192 * ((i / 100) ** 0.5))
                self.window.blit(
                    pg.transform.scale(stage, scaled_sizes[int(100 * (i / 100) ** 2)]),
                    scaled_positions[int(100 * (i / 100) ** 2)],
                )
            is_beat, frames_per_beat = yield unhandled_events

    def main(self):

        if self.render3d:
            loop = self.render_loop_3d()
        else:
            loop = self.render_loop_normal()
        next(loop)

        framecounter = 0
        blackout = False
        single_random = False
        taps = []
        tap_mean = 0
        last_beat = 0
        fps_slidewindow = []
        frames_per_beat = 0
        while True:
            is_beat = False
            fps_mean = (
                sum(fps_slidewindow) / len(fps_slidewindow) if fps_slidewindow else 0
            )
            if tap_mean and last_beat:
                this_beat = time.time()
                is_beat = this_beat >= last_beat + tap_mean
                frames_per_beat = fps_mean * tap_mean
                if is_beat:
                    last_beat = this_beat

            common_events = loop.send((is_beat, frames_per_beat))
            reinitialize = False
            for event in common_events:
                if event.type == pg.QUIT or (
                    event.type == pg.KEYDOWN and event.key == pg.K_ESCAPE
                ):
                    pg.quit()
                    sys.exit()
                elif event.type == pg.KEYDOWN:
                    if event.key == pg.K_F5:
                        single_random = True
                        print("Switching to new random preset")
                    elif event.key == pg.K_F8:
                        self.randomize = not self.randomize
                        state_str = (
                            f"on, {self.randomize_interval} s"
                            if self.randomize
                            else "off"
                        )
                        print(f"Random preset switching: [{state_str}]")
                    elif event.key == pg.K_F9:
                        self.render3d = not self.render3d
                        print(f'Pseudo-3d view [{"on" if self.render3d else "off"}]')
                        reinitialize = True
                    elif event.key == pg.K_F10:
                        self.trails = not self.trails
                        print(f'Trails [{"on" if self.trails else "off"}]')
                        reinitialize = True
                    elif event.key == pg.K_F11:
                        self.windowed = not self.windowed
                        print(f'Windowed mode [{"on" if self.windowed else "off"}]')
                        reinitialize = True
                    elif event.key == pg.K_SPACE:
                        blackout = not blackout
                        print(f'BLACKOUT [{"ON" if blackout else "off"}]')
                    elif event.key == pg.K_HOME:
                        print("resetting beat timing")
                        this_tap = time.time()
                        last_beat = this_tap
                    elif event.key == pg.K_END:
                        print("starting new tap measurement")
                        taps.clear()
                    elif event.key == pg.K_PAGEUP:
                        tap_mean /= 2
                        print(
                            f"Taps: Mean {tap_mean:5.3f} s, {60/tap_mean:5.3f} BPM, {fps_mean * tap_mean:5.3f} FPB [x2]"
                        )
                    elif event.key == pg.K_PAGEDOWN:
                        tap_mean *= 2
                        print(
                            f"Taps: Mean {tap_mean:5.3f} s, {60/tap_mean:5.3f} BPM, {fps_mean * tap_mean:5.3f} FPB [/2]"
                        )
                    elif event.key == pg.K_KP_ENTER or event.key == pg.K_RETURN:
                        this_tap = time.time()
                        if taps and this_tap - taps[-1] > 10:
                            print("starting new tap measurement")
                            taps.clear()
                        taps.append(this_tap)
                        last_beat = this_tap
                        if len(taps) >= 2:
                            tap_mean = sum(
                                map(lambda t1, t2: t2 - t1, taps, taps[1:])
                            ) / (len(taps) - 1)
                            print(
                                f"Taps: Mean {tap_mean:5.3f} s, {60/tap_mean:5.3f} BPM, {fps_mean * tap_mean:5.3f} FPB"
                            )

            if reinitialize:
                self.window, self.background = self.initialize()
                if self.render3d:
                    loop = self.render_loop_3d()
                else:
                    loop = self.render_loop_normal()
                next(loop)

            if blackout:
                self.window.fill(Colors.Black)

            if is_beat:
                pg.draw.rect(self.window, Colors.White, (0, 0, 30, 30))

            pg.display.flip()
            self.clock.tick(60)
            framecounter += 1

            if (
                self.randomize and (framecounter % (self.randomize_interval * 60)) == 0
            ) or single_random:
                single_random = False
                new_preset = self.presets.randomize()
                while new_preset == self.effects:
                    new_preset = self.presets.randomize()
                self.effects.clear()
                self.effects.extend(self.presets.randomize())

            current_fps = self.clock.get_fps()
            fps_slidewindow.append(current_fps)
            while len(fps_slidewindow) > 120:
                fps_slidewindow.pop(0)
            if self.print_fps:
                print(f"FPS: {current_fps:5.2f}, avg {fps_mean:5.2f}")


def app_main() -> None:
    argparser = ArgumentParser(
        description="beamshow - Render a light show for a video projector"
    )
    argparser.add_argument(
        "--3d",
        dest="render3d",
        action="store_true",
        help="Render a 3D preview instead of the normal output",
    )
    argparser.add_argument(
        "--list-displays", action="store_true", help="Show available displays"
    )
    argparser.add_argument(
        "-w", "--window", action="store_true", help="Display in a window"
    )
    argparser.add_argument(
        "-b", "--beat", action="store_true", help="Effects should react to beats"
    )
    argparser.add_argument(
        "--trails", action="store_true", help="Fade patterns out (trail mode)"
    )
    argparser.add_argument(
        "--randomize",
        metavar="N",
        type=int,
        nargs="?",
        default=0,
        const=5,
        help="Select random effect presets after <N> seconds",
    )
    argparser.add_argument("--fps", action="store_true", help="Show FPS in console")
    argparser.add_argument(
        "-d", "--display", type=int, default=0, help="ID of the display to use"
    )

    # print some nice output after the pygame banner to separate our stuff from theirs
    print("")
    print("-" * (len(str(argparser.description)) + 4))
    print(f"  {argparser.description}  ")
    print("-" * (len(str(argparser.description)) + 4))
    print("")

    args = argparser.parse_args()

    if args.list_displays:
        print_displays()

    show = Beamshow(
        render3d=args.render3d,
        window=args.window,
        trails=args.trails,
        randomize_interval=args.randomize,
        fps=args.fps,
        display=args.display,
        beat_reactive=args.beat,
    )
    show.main()


if __name__ == "__main__":
    app_main()