import array, time
from machine import Pin
import rp2
from time import sleep
import dma
import uasyncio as asyncio

# Configure the number of WS2812 LEDs.

@rp2.asm_pio(sideset_init=rp2.PIO.OUT_LOW, out_shiftdir=rp2.PIO.SHIFT_LEFT, autopull=True, pull_thresh=8)
def ws2812():
    T1 = 2
    T2 = 5
    T3 = 3
    wrap_target()
    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]
    wrap()

class WS2812B: 
    def __init__(self, num_leds, pin, state_machine, brightness=0.1, invert=False):
        self.sm = rp2.StateMachine(state_machine, ws2812, freq=8_000_000, sideset_base=Pin(pin))
        self.sm.active(1)
        self.ar = bytearray(num_leds*3)
        self.num_leds = num_leds
        self.brightness = brightness
        self.invert = invert
        self.pio_dma = dma.PIO_DMA_Transfer(state_machine+4, state_machine, 8, num_leds*3)

    def show(self):
        self.pio_dma.start_transfer(self.ar)
        
    def set(self, i, color):
        self.ar[i*3] = int(color[1]*self.brightness)
        self.ar[i*3+1] = int(color[0]*self.brightness)
        self.ar[i*3+2] = int(color[2]*self.brightness)
    
    def fill(self, color):
        for i in range(self.num_leds):
            self.set(i, color)

    def busy(self):
        return self.pio_dma.busy()
    
    async def color_chase(self, color, wait):
        for i in range(self.num_leds):
            self.set(i, color)
            await asyncio.sleep(wait)
            self.show()
        time.sleep(0.2)
     
    def wheel(self, pos):
        # Input a value 0 to 255 to get a color value.
        # The colours are a transition r - g - b - back to r.
        if pos < 0 or pos > 255:
            return (0, 0, 0)
        if pos < 85:
            return (255 - pos * 3, pos * 3, 0)
        if pos < 170:
            pos -= 85
            return (0, 255 - pos * 3, pos * 3)
        pos -= 170
        return (pos * 3, 0, 255 - pos * 3)
     
     
    async def rainbow_cycle(self, wait):
        for j in range(255):
            for i in range(self.num_leds):
                rc_index = (i * 256 // self.num_leds) + j
                self.set(i, self.wheel(rc_index & 255))
            self.show()
            await asyncio.sleep(wait)

    BLACK = (0, 0, 0)
    RED = (255, 0, 0)
    YELLOW = (255, 150, 0)
    GREEN = (0, 255, 0)
    CYAN = (0, 255, 255)
    BLUE = (0, 0, 255)
    PURPLE = (180, 0, 255)
    WHITE = (255, 255, 255)
    COLORS = (BLACK, RED, YELLOW, GREEN, CYAN, BLUE, PURPLE, WHITE)

if __name__ == "__main__":
    num_leds, pin, sm, brightness = 10, 0, 0, 1
    ws0 = WS2812B(num_leds, pin, sm, brightness)
    while True:
        for color in ws0.COLORS:
            ws0.fill(color)
            ws0.show()
            time.sleep(1)