pybeamshow/beamshow.py

from argparse import ArgumentParser
from typing import Iterable, List, Tuple
from pygame.locals import *
import math
import pygame as pg
import random
import sys
import time

from effects.bouncingspot import BouncingSpot
from effects.doublespot import DoubleSpot
from effects.effect import Effect, color_randomize, color_wheel, Colors
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)


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

    win = pg.display.set_mode(
        size=displays[display_id]
        if not windowed
        else (displays[display_id][0] // 2, displays[display_id][1] // 2),
        flags=pg.FULLSCREEN if not windowed else 0,
        display=display_id,
    )
    if 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)
    return win, background


def render_loop_normal(
    window: pg.Surface,
    background: pg.surface,
    effects: Iterable[Effect],
    clock: pg.time.Clock,
) -> None:
    blackout = False

    while True:
        for event in pg.event.get():
            if event.type == K_SPACE:
                blackout = not blackout
            if event.type == QUIT:
                pg.quit()
                sys.exit()

        window.blit(background, (0, 0))
        if not blackout:
            for e in effects:
                e.update()
                e.draw(window)
        yield


def render_loop_3d(
    window: pg.Surface,
    background: pg.surface,
    effects: Iterable[Effect],
    clock: pg.time.Clock,
) -> None:
    stage = pg.Surface(size=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
    ]

    while True:
        for event in pg.event.get():
            if event.type == QUIT:
                pg.quit()
                sys.exit()

        stage.blit(background, (0, 0))
        for e in effects:
            e.update()
            e.draw(stage)

        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))
            window.blit(
                pg.transform.scale(stage, scaled_sizes[int(100 * (i / 100) ** 2)]),
                scaled_positions[int(100 * (i / 100) ** 2)],
            )
        yield


def 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(
        "--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(argparser.description) + 4))
    print(f"  {argparser.description}  ")
    print("-" * (len(argparser.description) + 4))
    print("")

    args = argparser.parse_args()

    if args.list_displays:
        print_displays()

    window, background = initialize(
        display_id=args.display, windowed=args.window, trails=args.trails
    )

    presets = Presets(bounds=window.get_rect())
    if args.randomize:
        effects = presets.randomize()
    else:
        effects = [
            BouncingSpot(
                bounds=window.get_rect(),
                color=Colors.White,
                # color=color_wheel(),
                sizes=(window.get_height() / 8, window.get_height() / 8),
                velocity=(1, 1),
                x_factor=(1, 1),
                y_factor=(2.2, 2.2),
            ),
        ]

    clock = pg.time.Clock()

    if args.render3d:
        loop = render_loop_3d(window, background, effects, clock)
    else:
        loop = render_loop_normal(window, background, effects, clock)

    framecounter = 0
    while True:
        next(loop)
        pg.display.flip()
        clock.tick(60)
        framecounter += 1

        if args.randomize:
            if (framecounter % (args.randomize * 60)) == 0:
                effects.clear()
                effects.extend(presets.randomize())

        if args.fps:
            print(clock.get_fps())


if __name__ == "__main__":
    main()