import utime


class WS2812B:
    """RP2350的WS2812B驱动"""

    def __init__(self, pin=16):
        """
        初始化WS2812B驱动
        :param pin: WS2812B的DIN引脚编号（RP2350ZERO默认为16）
        """

        import rp2
        from machine import Pin

        # 基于RP2350的可编程输入输出（PIO）生成输入数据信号
        @rp2.asm_pio(
            sideset_init=rp2.PIO.OUT_LOW,  # 初始化侧置引脚为低电平
            out_shiftdir=rp2.PIO.SHIFT_LEFT,  # 输出移位方向为左移（高位先出）
            autopull=True,  # 自动从缓存区拉取数据
            pull_thresh=24,  # 拉取数据阈值
        )
        def ws2812b():
            # 时序参数
            T1 = 3  # 比特0或比特1由低电平上升为高电平前的低电平周期数
            T2 = 2  # 比特0高电平周期数，或比特1前置高电平周期数
            T3 = 5  # 比特0由高电平下降为低电平后的低电平周期数，或比特1后置高电平周期数

            wrap_target()
            label("bit_loop")
            # 从输出移位寄存器（OSR）移出一个比特到X寄存器，同时将侧置引脚置为低电平，延迟T1-1周期数
            out(x, 1).side(0)[T1 - 1]
            # 判断X寄存器中为比特0，同时将侧置引脚置为高电平，延迟T2-1周期数。若为比特0则跳转至bit_0，否则跳转至bit_loop，同时将侧置引脚置为高电平，延迟T3-1周期数
            jmp(not_x, "bit_0").side(1)[T2 - 1]
            jmp("bit_loop").side(1)[T3 - 1]
            label("bit_0")
            # 无操作，同时将侧置引脚置为低电平，延迟T3-1周期数
            nop().side(0)[T3 - 1]
            wrap()

        try:
            # 初始化状态机
            self.state_mechine = rp2.StateMachine(
                0, ws2812b, freq=8_000_000, sideset_base=Pin(pin)
            )
            # 启动状态机
            self.state_mechine.active(1)

        except Exception as exception:
            raise

    def send(self, rgb):
        """
        发送RGB颜色数据
        :param rgb: RGB颜色数据（元组）
        """
        r, g, b = rgb
        # 组合RGB颜色数据
        rgb = (r << 24) | (g << 16) | (b << 8)

        self.state_mechine.put(rgb)

    def iridesce(self):
        """炫彩"""
        # 色相
        H = 0
        # 饱和度
        S = 255
        # 亮度
        V = 16

        while True:
            H = H % 360
            # 色域
            region = H // 60
            # 色域相对位置
            remainder = (H - region * 60) * 256 // 60

            # 最小颜色分量值
            P = (V * (255 - S)) >> 8
            # 中间颜色分量值
            Q = (V * (255 - ((S * remainder) >> 8))) >> 8
            # 最大颜色分量值
            T = (V * (255 - ((S * (256 - remainder)) >> 8))) >> 8

            state_mechine.send(
                {
                    0: (V, T, P),
                    1: (Q, V, P),
                    2: (P, V, T),
                    3: (P, Q, V),
                    4: (T, P, V),
                }.get(region, (V, P, Q))
            )
            utime.sleep_ms(20)
            
            H = H + 1