portal/leds/animations.py

"""Animations for LedMatrix and LedStrip."""

from __future__ import annotations

import math
import random
from typing import Callable, Union

from leds.colors import dim, heat_color, hsv_to_rgb, wheel
from leds.matrix import LedMatrix
from leds.strip import LedStrip

Surface = Union[LedMatrix, LedStrip]
DrawFn = Callable[[Surface, int], None]


def _is_matrix(surface: Surface) -> bool:
    return hasattr(surface, "width") and hasattr(surface, "height")


def _size(surface: Surface) -> tuple[int, int]:
    if _is_matrix(surface):
        return surface.width, surface.height
    n = len(surface)
    return n, 1


def _set_index(surface: Surface, x: int, y: int, color: tuple[int, int, int]) -> None:
    if isinstance(surface, LedStrip):
        w, _ = _size(surface)
        surface[y * w + x] = color
    else:
        surface[x, y] = color


def _each_pixel(surface: Surface):
    if _is_matrix(surface):
        w, h = surface.width, surface.height
        for y in range(h):
            for x in range(w):
                yield x, y
    else:
        for i in range(len(surface)):
            yield i, 0


def rainbow_scroll(surface: Surface, frame: int) -> None:
    if isinstance(surface, LedStrip):
        n = len(surface)
        for i in range(n):
            surface[i] = wheel(frame * 2 + i * 256 // max(n, 1))
        return

    w, h = surface.width, surface.height
    strip = getattr(surface, "strip", None)
    if strip is not None and hasattr(surface, "index_at"):
        for y in range(h):
            for x in range(w):
                strip[surface.index_at(x, y)] = wheel(x * 3 + y * 8 + frame * 2)
    else:
        for y in range(h):
            for x in range(w):
                surface[x, y] = wheel(x * 3 + y * 8 + frame * 2)


def plasma(surface: Surface, frame: int) -> None:
    w, h = _size(surface)
    t = frame * 0.08
    for y in range(h):
        for x in range(w):
            v = (
                math.sin(x * 0.15 + t)
                + math.sin(y * 0.2 - t * 1.3)
                + math.sin((x + y) * 0.1 + t * 0.7)
            ) / 3.0
            _set_index(surface, x, y, hsv_to_rgb((v + 1) / 2, 1.0, 1.0))


def sine_wave(surface: Surface, frame: int) -> None:
    w, h = _size(surface)
    for y in range(h):
        for x in range(w):
            wave = math.sin(x * 0.25 + frame * 0.15) * math.cos(y * 0.35 - frame * 0.1)
            _set_index(surface, x, y, dim((0, 80, 200), (wave + 1) / 2))


def scanner(surface: Surface, frame: int) -> None:
    w, h = _size(surface)
    pos = frame % max(w * 2 - 2, 1)
    if pos >= w:
        pos = w * 2 - 2 - pos
    if _is_matrix(surface):
        surface.fill((0, 0, 0))
    else:
        surface.fill((0, 0, 0))
    for y in range(h):
        for dx in range(-3, 4):
            x = pos + dx
            if 0 <= x < w:
                _set_index(surface, x, y, dim((255, 0, 0), 1.0 - abs(dx) / 4.0))


def comet(surface: Surface, frame: int) -> None:
    w, h = _size(surface)
    head = frame % (w + 12)
    if _is_matrix(surface):
        surface.fill((0, 0, 0))
    else:
        surface.fill((0, 0, 0))
    for y in range(h):
        for i in range(12):
            x = head - i
            if 0 <= x < w:
                _set_index(surface, x, y, dim(wheel(frame * 3 + i * 10), (12 - i) / 12))


def pulse(surface: Surface, frame: int) -> None:
    level = (math.sin(frame * 0.12) + 1) / 2
    color = dim((120, 0, 180), 0.15 + level * 0.85)
    if _is_matrix(surface):
        surface.fill(color)
    else:
        surface.fill(color)


def sparkle(surface: Surface, frame: int) -> None:
    for x, y in _each_pixel(surface):
        if _is_matrix(surface):
            r, g, b = surface[x, y]
        else:
            r, g, b = surface[x]
        c = (r >> 2, g >> 2, b >> 2)
        _set_index(surface, x, y, c)
    w, h = _size(surface)
    for _ in range(6):
        _set_index(surface, random.randint(0, w - 1), random.randint(0, h - 1), (220, 220, 255))


def fire(surface: Surface, frame: int) -> None:
    w, h = _size(surface)
    size = w * h
    if not hasattr(fire, "_heat") or len(fire._heat) != size:
        fire._heat = [0.0] * size

    heat: list[float] = fire._heat
    for x in range(w):
        heat[(h - 1) * w + x] = random.random() * 0.5 + 0.5

    for y in range(h - 2, -1, -1):
        for x in range(w):
            below = (y + 1) * w + x
            left = below - 1 if x > 0 else below
            right = below + 1 if x < w - 1 else below
            decay = random.uniform(0.0, 0.18)
            heat[y * w + x] = max(0.0, (heat[left] + heat[below] + heat[right]) / 3 - decay)

    for y in range(h):
        for x in range(w):
            _set_index(surface, x, y, heat_color(heat[y * w + x]))


def rain(surface: Surface, frame: int) -> None:
    w, h = _size(surface)
    if not hasattr(rain, "_drops"):
        rain._drops = [
            {"x": random.randint(0, w - 1), "y": random.randint(-h, 0), "speed": random.randint(1, 3)}
            for _ in range(max(w // 2, 1))
        ]

    if _is_matrix(surface):
        surface.fill((0, 0, 0))
    else:
        surface.fill((0, 0, 0))

    for drop in rain._drops:
        drop["y"] += drop["speed"]
        if drop["y"] >= h:
            drop["y"] = random.randint(-3, -1)
            drop["x"] = random.randint(0, w - 1)
            drop["speed"] = random.randint(1, 3)
        x, y = drop["x"], int(drop["y"])
        if 0 <= y < h:
            _set_index(surface, x, y, (0, 220, 80))
            if y > 0:
                _set_index(surface, x, y - 1, dim((0, 220, 80), 0.4))


def stripes(surface: Surface, frame: int) -> None:
    w, h = _size(surface)
    offset = frame % max(w, 1)
    for y in range(h):
        for x in range(w):
            band = (x + offset) // 3 % 3
            if band == 0:
                _set_index(surface, x, y, (200, 0, 0))
            elif band == 1:
                _set_index(surface, x, y, (0, 200, 0))
            else:
                _set_index(surface, x, y, (0, 0, 200))


def bounce(surface: Surface, frame: int) -> None:
    w, h = _size(surface)
    if _is_matrix(surface):
        surface.fill((0, 0, 0))
    else:
        surface.fill((0, 0, 0))
    t = frame * 0.2
    x = int((math.sin(t) + 1) / 2 * (w - 1))
    y = int((math.cos(t * 1.3) + 1) / 2 * (h - 1))
    for dy in range(-1, 2):
        for dx in range(-1, 2):
            px, py = x + dx, y + dy
            if 0 <= px < w and 0 <= py < h:
                _set_index(surface, px, py, wheel(frame * 5))


def _iter_line(x0: int, y0: int, x1: int, y1: int):
    """Yield pixel coordinates along a Bresenham line."""
    dx = abs(x1 - x0)
    dy = -abs(y1 - y0)
    sx = 1 if x0 < x1 else -1
    sy = 1 if y0 < y1 else -1
    err = dx + dy
    x, y = x0, y0
    while True:
        yield x, y
        if x == x1 and y == y1:
            break
        e2 = 2 * err
        if e2 >= dy:
            err += dy
            x += sx
        if e2 <= dx:
            err += dx
            y += sy


def rolling(surface: Surface, frame: int) -> None:
    """Full-height band rolling left to right across the panel."""
    w, h = _size(surface)
    tail = 10
    head = frame % (w + tail)
    if _is_matrix(surface):
        surface.fill((0, 0, 0))
    else:
        surface.fill((0, 0, 0))
    for y in range(h):
        hue_shift = y * 18 + frame * 3
        for i in range(tail):
            x = head - i
            if 0 <= x < w:
                bright = (tail - i) / tail
                _set_index(surface, x, y, dim(wheel(hue_shift + i * 8), bright))


def spin_line(surface: Surface, frame: int) -> None:
    """Line from panel center spinning like a radar sweep."""
    w, h = _size(surface)
    if _is_matrix(surface):
        surface.fill((0, 0, 0))
    else:
        surface.fill((0, 0, 0))

    if not _is_matrix(surface):
        n = len(surface)
        pos = int((frame * 3) % max(n, 1))
        for i in range(8):
            idx = (pos - i) % n
            surface[idx] = dim(wheel(frame * 5 + i * 20), (8 - i) / 8)
        return

    cx = (w - 1) / 2.0
    cy = (h - 1) / 2.0
    angle = frame * 0.14
    reach = max(w, h) * 1.2
    x1 = int(cx + math.cos(angle) * reach)
    y1 = int(cy + math.sin(angle) * reach)
    color = wheel(frame * 4)

    for x, y in _iter_line(int(cx), int(cy), x1, y1):
        if 0 <= x < w and 0 <= y < h:
            _set_index(surface, x, y, color)

    hub_x, hub_y = int(cx), int(cy)
    if 0 <= hub_x < w and 0 <= hub_y < h:
        _set_index(surface, hub_x, hub_y, (255, 255, 255))


def solid(surface: Surface, frame: int) -> None:
    colors = [(255, 0, 0), (0, 255, 0), (0, 0, 255)]
    if _is_matrix(surface):
        surface.fill(colors[frame % 3])
    else:
        surface.fill(colors[frame % 3])


ANIMATIONS: dict[str, DrawFn] = {
    "solid": solid,
    "rainbow": rainbow_scroll,
    "plasma": plasma,
    "wave": sine_wave,
    "scanner": scanner,
    "comet": comet,
    "pulse": pulse,
    "sparkle": sparkle,
    "fire": fire,
    "rain": rain,
    "stripes": stripes,
    "bounce": bounce,
    "rolling": rolling,
    "spin_line": spin_line,
}

DEFAULT_FPS: dict[str, float] = {
    "solid": 1,
    "rainbow": 30,
    "plasma": 25,
    "wave": 25,
    "scanner": 35,
    "comet": 30,
    "pulse": 30,
    "sparkle": 40,
    "fire": 35,
    "rain": 30,
    "stripes": 25,
    "bounce": 35,
    "rolling": 30,
    "spin_line": 28,
}