portal/examples/animations.py

#!/usr/bin/env python3
"""Run LED matrix animations locally (SPI/PIO) or on Pico panel (UDP)."""

from __future__ import annotations

import argparse
import os
import sys
import time

from leds import LedMatrix, LedStrip
from leds.animations import ANIMATIONS, DEFAULT_FPS
from leds.array_config import (
    MATRIX_BRIGHTNESS,
    MATRIX_HEIGHT,
    MATRIX_SPI_BUS,
    MATRIX_SPI_DEVICE,
    MATRIX_WIDTH,
)
from leds.config import spi_strip_config
from leds.panel_udp import (
    add_panel_network_args,
    animation_playlist,
    format_panel_targets,
    panel_targets_from_args,
    run_panel_animation_loop,
    run_panel_playlist,
)


def _flash_test(surface, brightness: float) -> None:
    for color in ((255, 0, 0), (0, 255, 0), (0, 0, 255)):
        surface.fill(color)
        surface.show()
        time.sleep(0.35)
    surface.fill((0, 0, 0))
    surface.show()


def _run_local_playlist(surface, playlist: list[str], args: argparse.Namespace) -> None:
    while True:
        for name in playlist:
            draw = ANIMATIONS[name]
            fps = args.fps or DEFAULT_FPS.get(name, 30)
            delay = 1.0 / fps
            frames = max(int(args.duration * fps), 1)
            print(f"  {name} ({fps:.0f} fps)")
            for frame in range(frames):
                draw(surface, frame)
                surface.show()
                time.sleep(delay)


def main() -> int:
    names = list(ANIMATIONS.keys())
    parser = argparse.ArgumentParser(description="LED matrix animations")
    parser.add_argument(
        "animation",
        nargs="?",
        choices=names + ["all"],
        default="all",
    )
    parser.add_argument("--width", type=int, default=MATRIX_WIDTH)
    parser.add_argument("--height", type=int, default=MATRIX_HEIGHT)
    parser.add_argument(
        "--spi-bus",
        type=int,
        default=int(os.environ.get("PORTAL_SPI_BUS", MATRIX_SPI_BUS)),
    )
    parser.add_argument("--spi-device", type=int, default=MATRIX_SPI_DEVICE)
    parser.add_argument("--brightness", type=float, default=MATRIX_BRIGHTNESS)
    parser.add_argument("--duration", type=float, default=12.0)
    parser.add_argument("--fps", type=float, default=None)
    parser.add_argument("--linear", action="store_true")
    parser.add_argument(
        "--serpentine",
        action="store_true",
        help="zigzag odd rows (default: all rows left → right)",
    )
    parser.add_argument("--no-test", action="store_true")
    parser.add_argument(
        "--panel",
        action="store_true",
        help="send to Pico panel(s) over UDP instead of local SPI/PIO",
    )
    add_panel_network_args(parser)
    parser.add_argument("--once", action="store_true", help="panel: play playlist once")
    args = parser.parse_args()

    playlist = animation_playlist(args.animation)

    if args.panel:
        targets = panel_targets_from_args(args)
        print(f"Panel animations -> {format_panel_targets(targets, args.port)}")
        print(f"Playlist: {', '.join(playlist)}")
        try:
            if len(playlist) == 1 and args.animation != "all":
                run_panel_animation_loop(
                    targets,
                    args.port,
                    playlist[0],
                    fps=args.fps,
                    brightness=args.brightness,
                    panel_id=args.panel_id,
                )
            else:
                run_panel_playlist(
                    targets,
                    args.port,
                    playlist,
                    duration=args.duration,
                    fps_override=args.fps,
                    brightness=args.brightness,
                    panel_id=args.panel_id,
                    loop=not args.once,
                )
        except KeyboardInterrupt:
            print("\nStopped.")
        except OSError as exc:
            print(f"Network error: {exc}", file=sys.stderr)
            return 1
        return 0

    count = args.width * args.height
    serpentine = "rows" if args.serpentine else "none"
    device = f"/dev/spidev{args.spi_bus}.{args.spi_device}"
    mode = f"linear {count}" if args.linear else f"{args.width}×{args.height}"
    print(f"Matrix {mode} ({count} LEDs) → {device}")

    strip_cfg = spi_strip_config(
        count,
        spi_bus=args.spi_bus,
        spi_device=args.spi_device,
        brightness=args.brightness,
    )

    try:
        if args.linear:
            with LedStrip(
                count,
                backend="spi",
                spi_bus=args.spi_bus,
                spi_device=args.spi_device,
                brightness=args.brightness,
            ) as surface:
                if not args.no_test:
                    _flash_test(surface, args.brightness)
                _run_local_playlist(surface, playlist, args)
        else:
            with LedMatrix(
                args.width,
                args.height,
                serpentine=serpentine,
                config0=strip_cfg,
            ) as surface:
                if not args.no_test:
                    _flash_test(surface, args.brightness)
                _run_local_playlist(surface, playlist, args)
    except KeyboardInterrupt:
        print("\nStopped.")
        return 0
    except Exception as exc:
        print(f"Error: {exc}", file=sys.stderr)
        print("Try: --spi-bus 1  or  --panel --host <pico-ip>", file=sys.stderr)
        return 1
    return 0


if __name__ == "__main__":
    raise SystemExit(main())