# ws2812_ctrl.py

import array
import time
from machine import Pin
import rp2
from rp2 import PIO, StateMachine, asm_pio

@asm_pio(
    sideset_init=PIO.OUT_LOW,
    out_shiftdir=PIO.SHIFT_LEFT,
    autopull=True,
    pull_thresh=24
)
def _ws2812_pio():
    T1 = 2
    T2 = 5
    T3 = 3
    label("bitloop")
    out(x, 1)               .side(0)    [T3 - 1]
    jmp(not_x, "do_zero")   .side(1)    [T1 - 1]
    jmp("bitloop")          .side(1)    [T2 - 1]
    label("do_zero")
    nop()                   .side(0)    [T2 - 1]

class WS2812:
    def __init__(self, pin_num, num_leds=1, freq=8_000_000, sm_id=0):
        """
        :param pin_num: WS2812 数据引脚号
        :param num_leds: LED 数量
        :param freq: PIO 运行频率
        :param sm_id: 使用的 StateMachine ID (0–7)
        """
        self.num_leds = num_leds
        self._array = array.array("I", [0] * num_leds)
        self.sm = StateMachine(
            sm_id,
            _ws2812_pio,
            freq=freq,
            sideset_base=Pin(pin_num)
        )
        self.sm.active(1)

    @staticmethod
    def create_color(r, g, b):
        """把 (r, g, b) 三通道转换成 24bit 整型"""
        return (r & 0xFF) << 16 | (g & 0xFF) << 8 | (b & 0xFF)

    def show(self, color):
        """立即将所有 LED 显示为同一个颜色"""
        packed = color
        for i in range(self.num_leds):
            self._array[i] = packed
        self.sm.put(self._array, 8)

    def fade(self, from_color, to_color, steps=100, delay_ms=10):
        """
        在 from_color 和 to_color 之间渐变
        :param from_color: 起始颜色（24bit）
        :param to_color: 结束颜色（24bit）
        :param steps: 渐变步数
        :param delay_ms: 每步延迟（毫秒）
        """
        r1 = (from_color >> 16) & 0xFF
        g1 = (from_color >> 8) & 0xFF
        b1 = from_color & 0xFF

        r2 = (to_color >> 16) & 0xFF
        g2 = (to_color >> 8) & 0xFF
        b2 = to_color & 0xFF

        for step in range(steps):
            r = r1 + (r2 - r1) * step // (steps - 1)
            g = g1 + (g2 - g1) * step // (steps - 1)
            b = b1 + (b2 - b1) * step // (steps - 1)
            color = self.create_color(r, g, b)
            self.show(color)
            time.sleep_ms(delay_ms)

    def cycle(self, color_list, steps_per_transition=100, delay_ms=10):
        """
        在给定的颜色列表中循环渐变
        :param color_list: 一个整型颜色列表
        :param steps_per_transition: 每次渐变步数
        :param delay_ms: 每步延迟（毫秒）
        """
        n = len(color_list)
        if n < 2:
            return
        idx = 0
        while True:
            from_c = color_list[idx]
            to_c = color_list[(idx + 1) % n]
            self.fade(from_c, to_c, steps_per_transition, delay_ms)
            idx = (idx + 1) % n
