portal/leds/matrix.py

"""2D WS2812 matrix mapped onto one or two strips."""

from __future__ import annotations

from typing import Literal, Sequence, Tuple

from leds.array_config import (
    MATRIX_BRIGHTNESS,
    MATRIX_HEIGHT,
    MATRIX_ROWS_PER_STRIP,
    MATRIX_SERPENTINE,
    MATRIX_WIDTH,
)
from leds.config import StripConfig, default_dual_strip_configs
from leds.strip import Color, LedStrip, _normalize_color
from leds.strips import LedStrips

Serpentine = Literal["rows", "none"]


class LedMatrix:
    """
    45×9 matrix on one WS2812 data line (405 LEDs).

    Any size works: LedMatrix(width=50, height=12) → 600 LEDs, etc.
    """

    def __init__(
        self,
        width: int = MATRIX_WIDTH,
        height: int = MATRIX_HEIGHT,
        *,
        rows_per_strip: Sequence[int] | None = MATRIX_ROWS_PER_STRIP,
        serpentine: Serpentine = MATRIX_SERPENTINE,
        brightness: float = MATRIX_BRIGHTNESS,
        config0: StripConfig | None = None,
        config1: StripConfig | None = None,
    ) -> None:
        if width < 1 or height < 1:
            raise ValueError("width and height must be positive")

        self.width = width
        self.height = height
        self.serpentine = serpentine
        self._strip: LedStrip | None = None
        self._strips: LedStrips | None = None

        if rows_per_strip is None:
            self._rows_per_strip = (height,)
            count0 = width * height
            if config0 is None:
                from leds.config import default_strip_config

                config0 = default_strip_config(count0, brightness=brightness)
            self._strip = LedStrip(
                count0,
                backend=config0.backend,
                pin=config0.pin,
                channel=config0.channel,
                device=config0.device,
                brightness=config0.brightness,
                spi_bus=config0.spi_bus,
                spi_device=config0.spi_device,
                spi_max_speed_hz=config0.spi_max_speed_hz,
                wire_order=config0.wire_order,
                passthrough=config0.passthrough,
                swap_rg=config0.swap_rg,
            )
        else:
            total_rows = sum(rows_per_strip)
            if total_rows != height:
                raise ValueError(
                    f"rows_per_strip sums to {total_rows}, expected {height}"
                )
            counts = [width * r for r in rows_per_strip]
            self._rows_per_strip = tuple(rows_per_strip)
            if config0 is None or config1 is None:
                defaults = default_dual_strip_configs(
                    counts[0], counts[1], brightness=brightness
                )
                config0 = config0 or defaults[0]
                config1 = config1 or defaults[1]
            self._strips = LedStrips(
                counts[0],
                counts[1],
                config0=config0,
                config1=config1,
                brightness=brightness,
            )

    def _locate(self, x: int, y: int) -> Tuple[int, int]:
        if not (0 <= x < self.width and 0 <= y < self.height):
            raise IndexError(f"({x}, {y}) out of range for {self.width}×{self.height}")

        row_offset = 0
        for strip_id, num_rows in enumerate(self._rows_per_strip):
            if y < row_offset + num_rows:
                local_y = y - row_offset
                col = x
                if self.serpentine == "rows" and y % 2 == 1:
                    col = self.width - 1 - x
                return strip_id, local_y * self.width + col
            row_offset += num_rows

        raise IndexError(f"({x}, {y}) out of range")

    def index_at(self, x: int, y: int) -> int:
        """Strip index for (x, y) on a single-strip matrix."""
        strip_id, index = self._locate(x, y)
        if strip_id != 0 or self._strip is None:
            raise RuntimeError("index_at requires a single-strip matrix")
        return index

    @property
    def strip(self) -> LedStrip | None:
        return self._strip

    def __getitem__(self, pos: Tuple[int, int]) -> Tuple[int, int, int]:
        x, y = pos
        strip_id, index = self._locate(x, y)
        if self._strip is not None:
            return self._strip[index]
        assert self._strips is not None
        return self._strips[strip_id][index]

    def __setitem__(self, pos: Tuple[int, int], color: Color) -> None:
        x, y = pos
        strip_id, index = self._locate(x, y)
        rgb = _normalize_color(color)
        if self._strip is not None:
            self._strip[index] = rgb
        else:
            assert self._strips is not None
            self._strips[strip_id][index] = rgb

    def set_pixel(self, x: int, y: int, color: Color) -> None:
        self[x, y] = color

    def fill(self, color: Color) -> None:
        rgb = _normalize_color(color)
        if self._strip is not None:
            self._strip.fill(rgb)
        else:
            assert self._strips is not None
            self._strips[0].fill(rgb)
            self._strips[1].fill(rgb)

    def show(self) -> None:
        if self._strip is not None:
            self._strip.show()
        else:
            assert self._strips is not None
            self._strips.show()

    def show_rgb_bytes(self, rgb: bytes) -> None:
        """Push a linear RGB frame (row-major) straight to the strip."""
        if self._strip is None:
            raise RuntimeError("show_rgb_bytes requires a single-strip matrix")
        self._strip.show_rgb_bytes(rgb)

    def clear(self) -> None:
        if self._strip is not None:
            self._strip.clear()
        else:
            assert self._strips is not None
            self._strips.clear()

    def close(self) -> None:
        if self._strip is not None:
            self._strip.close()
        elif self._strips is not None:
            self._strips.close()

    def __enter__(self) -> LedMatrix:
        return self

    def __exit__(self, *_exc) -> None:
        self.close()