LuxStats/firmware/code.py

# LuxStats firmware v3 - Pimoroni Tiny 2040 + NeoPixels (CircuitPython 8/9)
#
# Config lives in RAM only; the host configurator owns it and re-pushes it on
# connect and on every edit.
#
# Segment types:
#   percent  0.0-1.0 value -> gradient (optional per-segment "gradient")
#   temp     RAW Celsius value -> gradient positioned by min_c/max_c
#   link     Ethernet speed -> POSITION palette "colors"
#            positions: [down, 10M, 100M, 1G, 2.5G, 10G]
#   wifi     Wi-Fi band -> POSITION palette "colors"
#            positions: [down, 2.4GHz, 5GHz, 6GHz]
#   vpn/bool on_color / off_color
#   custom   pushed [r,g,b,w] value, else fallback "color"
#   ambient  effect zone: off|solid|breathe|rainbow|sparkle + color + speed
#
# Camera glow: an ASSIGNED LED list (host sends expanded indices) with the
# same effect presets as ambient zones, rendered as an overlay when active.
#
# Overlay order: segments -> camera glow -> notification blinker -> test blink
#
# Protocol (host -> device), one JSON object per line:
#   {"cmd":"ping"} {"cmd":"config","config":{...}} {"cmd":"get_config"}
#   {"cmd":"status","values":{...}}          silent unless "ack":1
#   {"cmd":"notify","color":[..],"period":s,"count":n}  n -1 = until cleared
#   {"cmd":"notify_clear"} {"cmd":"camera","active":bool}
#   {"cmd":"test","start":i,"count":n,"active":bool}
#   {"cmd":"raw","pixels":[[..],...]} {"cmd":"raw_off"}

import json
import math
import random
import time

import board
import neopixel
import usb_cdc

LINK_TIERS = [0, 10, 100, 1000, 2500, 10000]
DEFAULT_LINK_COLORS = [
    [0, 0, 0, 0], [255, 0, 0, 0], [255, 90, 0, 0],
    [0, 255, 0, 0], [0, 255, 180, 0], [80, 0, 255, 0],
]
WIFI_BANDS = [0, 2, 5, 6]
DEFAULT_WIFI_COLORS = [
    [0, 0, 0, 0], [255, 120, 0, 0], [0, 255, 60, 0], [40, 80, 255, 0],
]

DEFAULT_CONFIG = {
    "num_pixels": 8,
    "pin": "GP0",
    "pixel_order": "GRBW",
    "brightness": 0.10,
    "fps": 30,
    "segments": [
        {"name": "cpu", "type": "percent", "start": 0, "count": 2,
         "fill": True,
         "gradient": [[0.0, [0, 0, 255, 0]], [0.5, [0, 255, 0, 0]],
                      [1.0, [255, 0, 0, 0]]]},
        {"name": "eth",  "type": "link", "start": 2, "count": 1},
        {"name": "wifi", "type": "wifi", "start": 3, "count": 1},
        {"name": "vpn",  "type": "vpn",  "start": 4, "count": 1},
        {"name": "glow", "type": "ambient", "start": 5, "count": 3,
         "effect": "breathe", "color": [10, 10, 30, 0], "speed": 1.0},
    ],
    "camera": {"pixels": [5, 6, 7], "color": [255, 40, 0, 0],
               "effect": "breathe", "speed": 1.0},
    "notify": {"start": 0, "count": 0, "blinks": 6,
               "default_color": [0, 60, 255, 0], "default_period": 0.4},
}

TEST_TIMEOUT = 180.0


def clamp(x, a, b):
    return a if x < a else (b if x > b else x)


def lerp(a, b, t):
    return a + (b - a) * t


def pad4(c):
    c = list(c)[:4]
    return c + [0] * (4 - len(c))


def scale(color, b):
    return tuple(int(c * b) for c in color)


def default_percent_color(t):
    t = clamp(t, 0.0, 1.0)
    if t < 0.33:
        f = t / 0.33
        r, g, b = 0, int(255 * f), int(255 * (1 - f))
    elif t < 0.66:
        f = (t - 0.33) / 0.33
        r, g, b = int(255 * f), 255, 0
    else:
        f = (t - 0.66) / 0.34
        r, g, b = 255, int(255 * (1 - f)), 0
    return (r, g, b, 0)


def gradient_color(stops, t):
    if not stops:
        return default_percent_color(t)
    stops = sorted(stops, key=lambda s: s[0])
    t = clamp(t, 0.0, 1.0)
    if t <= stops[0][0]:
        return tuple(pad4(stops[0][1]))
    for i in range(1, len(stops)):
        p0, c0 = stops[i - 1][0], pad4(stops[i - 1][1])
        p1, c1 = stops[i][0], pad4(stops[i][1])
        if t <= p1:
            f = 0.0 if p1 == p0 else (t - p0) / (p1 - p0)
            return tuple(int(lerp(a, b, f)) for a, b in zip(c0, c1))
    return tuple(pad4(stops[-1][1]))


def palette_color(seg, defaults, index):
    """Positional palette lookup (link/wifi): one color per position."""
    colors = seg.get("colors") or defaults
    index = int(clamp(index, 0, len(colors) - 1))
    return tuple(pad4(colors[index]))


def wheel(pos):
    pos = int(pos) & 255
    if pos < 85:
        return (255 - pos * 3, pos * 3, 0, 0)
    if pos < 170:
        pos -= 85
        return (0, 255 - pos * 3, pos * 3, 0)
    pos -= 170
    return (pos * 3, 0, 255 - pos * 3, 0)


class LuxStats:
    def __init__(self):
        self.serial = usb_cdc.data or usb_cdc.console
        # Non-blocking IO: a blocked write must never stall the loop
        if self.serial is not None:
            try:
                self.serial.timeout = 0
                self.serial.write_timeout = 0
            except (AttributeError, ValueError):
                pass
        self.buf = b""
        self.pixels = None
        self.statuses = {}
        self.notify = None
        self.camera_active = False
        self.test = None
        self.raw = None
        self.sparkles = {}          # zone key -> [[pixel, t0], ...]
        self.apply_config(DEFAULT_CONFIG)

    # ---------- config ----------

    def apply_config(self, cfg):
        merged = dict(DEFAULT_CONFIG)
        merged.update(cfg or {})
        self.config = merged
        self.sparkles = {}
        if self.pixels is not None:
            try:
                self.pixels.deinit()
            except Exception:
                pass
        pin = getattr(board, merged.get("pin", "GP0"), board.GP0)
        order = merged.get("pixel_order", "GRBW")
        self.bpp = 4 if "W" in order else 3
        self.pixels = neopixel.NeoPixel(
            pin, int(merged["num_pixels"]),
            brightness=1.0, auto_write=False, pixel_order=order,
        )
        self.frame = [(0, 0, 0, 0)] * int(merged["num_pixels"])

    # ---------- serial ----------

    def poll_serial(self):
        s = self.serial
        if s is None:
            return
        n = s.in_waiting
        if n:
            self.buf += s.read(n)
            while b"\n" in self.buf:
                line, self.buf = self.buf.split(b"\n", 1)
                line = line.strip()
                if line:
                    self.handle_line(line)
        if len(self.buf) > 8192:
            self.buf = b""

    def send(self, obj):
        try:
            self.serial.write((json.dumps(obj) + "\n").encode())
        except Exception:
            pass

    def handle_line(self, line):
        try:
            msg = json.loads(line)
        except ValueError:
            self.send({"err": "bad json"})
            return
        cmd = msg.get("cmd", "")

        if cmd == "ping":
            self.send({"pong": 1, "num_pixels": self.config["num_pixels"],
                       "fw": 3})
        elif cmd == "config":
            self.apply_config(msg.get("config", {}))
            self.send({"ok": "config"})
        elif cmd == "get_config":
            self.send({"config": self.config})
        elif cmd == "status":
            vals = msg.get("values", {})
            if isinstance(vals, dict):
                self.statuses.update(vals)
            if msg.get("ack"):
                self.send({"ok": "status"})
        elif cmd == "notify":
            nz = self.config["notify"]
            self.notify = {
                "color": tuple(pad4(msg.get("color",
                                            nz["default_color"]))),
                "period": float(msg.get("period", nz["default_period"])),
                "count": int(msg.get("count", nz.get("blinks", 6))),
                "t0": time.monotonic(),
            }
            self.send({"ok": "notify"})
        elif cmd == "notify_clear":
            self.notify = None
            self.send({"ok": "notify_clear"})
        elif cmd == "camera":
            self.camera_active = bool(msg.get("active", False))
            self.send({"ok": "camera"})
        elif cmd == "test":
            if msg.get("active", True):
                self.test = {"start": int(msg.get("start", 0)),
                             "count": int(msg.get("count", 1)),
                             "t0": time.monotonic()}
            else:
                self.test = None
            self.send({"ok": "test"})
        elif cmd == "raw":
            self.raw = [tuple(pad4(p)) for p in msg.get("pixels", [])]
            self.send({"ok": "raw"})
        elif cmd == "raw_off":
            self.raw = None
            self.send({"ok": "raw_off"})
        else:
            self.send({"err": "unknown cmd '%s'" % cmd})

    # ---------- effect helper (ambient zones + camera glow) ----------

    def effect_color(self, key, effect, base, speed, now, k, idxs):
        """Color for the k-th pixel of an effect zone at time `now`."""
        if effect == "solid" or effect == "off":
            return base
        if effect == "breathe":
            t = 0.5 + 0.5 * math.sin(now * 2.0 * speed + k * 0.4)
            return scale(base, 0.15 + 0.85 * t)
        if effect == "rainbow":
            return wheel(now * 40 * speed + k * (256 // max(len(idxs), 1)))
        if effect == "sparkle":
            return scale(base, 0.10)
        return base

    def apply_sparkles(self, key, idxs, speed, now):
        sp = self.sparkles.setdefault(key, [])
        if idxs and random.random() < 0.15 * speed:
            sp.append([random.choice(idxs), now])
        keep = []
        for i, t0 in sp:
            age = now - t0
            if age < 0.6 and i in idxs:
                self.frame[i] = scale((255, 255, 255, 40), 1.0 - age / 0.6)
                keep.append([i, t0])
        self.sparkles[key] = keep

    # ---------- rendering ----------

    def segment_color(self, seg):
        stype = seg.get("type", "percent")
        val = self.statuses.get(seg.get("name", ""), None)
        grad = seg.get("gradient") or []

        if stype == "percent":
            if val is None:
                return (4, 4, 4, 0)
            return gradient_color(grad, float(val))
        if stype == "temp":
            # raw Celsius, positioned by the segment's min/max
            if val is None:
                return (4, 4, 4, 0)
            lo = float(seg.get("min_c", 20.0))
            hi = float(seg.get("max_c", 90.0))
            t = 0.0 if hi <= lo else (float(val) - lo) / (hi - lo)
            return gradient_color(grad, t)
        if stype == "link":
            speed = int(val or 0)
            tier = 0
            for i, k in enumerate(LINK_TIERS):
                if speed >= k and speed > 0:
                    tier = i
            if speed <= 0:
                tier = 0
            return palette_color(seg, DEFAULT_LINK_COLORS, tier)
        if stype == "wifi":
            band = int(val or 0)
            pos = WIFI_BANDS.index(band) if band in WIFI_BANDS else 0
            return palette_color(seg, DEFAULT_WIFI_COLORS, pos)
        if stype == "vpn":
            on = tuple(pad4(seg.get("on_color", [170, 0, 255, 0])))
            off = tuple(pad4(seg.get("off_color", [6, 6, 6, 0])))
            return on if val else off
        if stype == "bool":
            on = tuple(pad4(seg.get("on_color", [0, 255, 0, 0])))
            off = tuple(pad4(seg.get("off_color", [4, 0, 0, 0])))
            return on if val else off
        if stype == "custom":
            if isinstance(val, (list, tuple)) and len(val) >= 3:
                return tuple(pad4(val))
            return tuple(pad4(seg.get("color", [0, 0, 0, 0])))
        return (0, 0, 0, 0)

    def render_status_segment(self, seg, n):
        c = self.segment_color(seg)
        start, count = int(seg.get("start", 0)), int(seg.get("count", 1))
        fill = seg.get("fill", False) and seg.get("type") == "percent"
        val = self.statuses.get(seg.get("name", ""), 0) or 0
        for j in range(count):
            i = start + j
            if not (0 <= i < n):
                continue
            if fill:
                lit = float(val) * count
                if j + 1 <= lit:
                    self.frame[i] = c
                elif j < lit:
                    self.frame[i] = scale(c, lit - j)
                else:
                    self.frame[i] = (2, 2, 2, 0)
            else:
                self.frame[i] = c

    def render_ambient_segment(self, seg, key, now, n):
        effect = seg.get("effect", "breathe")
        speed = float(seg.get("speed", 1.0))
        color = tuple(pad4(seg.get("color", [10, 10, 30, 0])))
        start, count = int(seg.get("start", 0)), int(seg.get("count", 1))
        idxs = [start + j for j in range(count) if 0 <= start + j < n]
        if effect == "off":
            for i in idxs:
                self.frame[i] = (0, 0, 0, 0)
            return
        for k, i in enumerate(idxs):
            self.frame[i] = self.effect_color(
                key, effect, color, speed, now, k, idxs)
        if effect == "sparkle":
            self.apply_sparkles(key, idxs, speed, now)

    def render(self):
        now = time.monotonic()
        n = self.config["num_pixels"]

        if self.raw is not None:
            for i in range(n):
                self.frame[i] = self.raw[i] if i < len(self.raw) \
                    else (0, 0, 0, 0)
        else:
            self.frame = [(0, 0, 0, 0)] * n

            for idx, seg in enumerate(self.config.get("segments", [])):
                if seg.get("type") == "ambient":
                    self.render_ambient_segment(seg, "seg%d" % idx, now, n)
                else:
                    self.render_status_segment(seg, n)

            # camera glow overlay: assigned LED list + effect preset
            if self.camera_active:
                cam = self.config["camera"]
                color = tuple(pad4(cam.get("color", [255, 40, 0, 0])))
                effect = cam.get("effect", "breathe")
                speed = float(cam.get("speed", 1.0))
                idxs = [i for i in cam.get("pixels", []) if 0 <= i < n]
                for k, i in enumerate(idxs):
                    self.frame[i] = self.effect_color(
                        "cam", effect, color, speed, now, k, idxs)
                if effect == "sparkle":
                    self.apply_sparkles("cam", idxs, speed, now)

            if self.notify:
                nt = self.notify
                cycle = int((now - nt["t0"]) / nt["period"])
                if nt["count"] >= 0 and cycle >= nt["count"] * 2:
                    self.notify = None
                elif cycle % 2 == 0:
                    nz = self.config["notify"]
                    s, cnt = int(nz.get("start", 0)), int(nz.get("count", 0))
                    rng = range(n) if cnt <= 0 else range(s, s + cnt)
                    for i in rng:
                        if 0 <= i < n:
                            self.frame[i] = nt["color"]

            if self.test:
                if now - self.test["t0"] > TEST_TIMEOUT:
                    self.test = None
                else:
                    on = int(now * 4) % 2 == 0
                    c = (255, 140, 0, 0) if on else (255, 255, 255, 60)
                    s, cnt = self.test["start"], self.test["count"]
                    for i in range(s, s + cnt):
                        if 0 <= i < n:
                            self.frame[i] = c

        b = float(self.config.get("brightness", 0.1))
        for i in range(n):
            c = scale(self.frame[i], b)
            self.pixels[i] = c if self.bpp == 4 else c[:3]
        self.pixels.show()

    def run(self):
        next_frame = time.monotonic()
        while True:
            self.poll_serial()
            now = time.monotonic()
            if now >= next_frame:
                self.render()
                next_frame = now + 1.0 / max(
                    int(self.config.get("fps", 30)), 5)
            time.sleep(0.002)


LuxStats().run()