# Example using PIO to drive a set of WS2812 LEDs.
import array, rp2
from machine import Pin
from time import sleep

# Configure the number of WS2812 LEDs.
NUM_LEDS = 4
PIN_NUM = 22 

@rp2.asm_pio(sideset_init=rp2.PIO.OUT_LOW, out_shiftdir=rp2.PIO.SHIFT_LEFT, autopull=True, pull_thresh=24)
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()
pass
        
class RGBLed(object): #NeoPixel
    
    def __init__(self,pin=PIN_NUM,num=NUM_LEDS,brightness=0.8):
        self.pin=pin
        self.num=num
        self.brightness = brightness
        
        # Create the StateMachine with the ws2812 program, outputting on pin
        self.sm = rp2.StateMachine(0, ws2812, freq=8_000_000, sideset_base=Pin(PIN_NUM))

        # Start the StateMachine, it will wait for data on its FIFO.
        self.sm.active(1)

        # Display a pattern on the LEDs via an array of LED RGB values.
        self.ar = array.array("I", [0 for _ in range(self.num)])
        
        self.BLACK = (0, 0, 0)
        self.RED = (255, 0, 0)
        self.YELLOW = (255, 150, 0)
        self.GREEN = (0, 255, 0)
        self.CYAN = (0, 255, 255)
        self.BLUE = (0, 0, 255)
        self.PURPLE = (180, 0, 255)
        self.WHITE = (255, 255, 255)
        self.COLORS = (self.BLACK, self.RED, self.YELLOW, self.GREEN, self.CYAN, self.BLUE, self.PURPLE, self.WHITE)
    pass
        
    def pixels_show(self):
        dimmer_ar = array.array("I", [0 for _ in range(self.num)])
        
        for i,c in enumerate(self.ar):
            r = int(((c >> 8) & 0xFF) * self.brightness)
            g = int(((c >> 16) & 0xFF) * self.brightness)
            b = int((c & 0xFF) * self.brightness)
            dimmer_ar[i] = (g<<16) + (r<<8) + b
        pass
    
        self.sm.put(dimmer_ar, 8)
    pass

    def set_color( self, color ):
        self.pixels_set(0, color)
        self.pixels_set(1, color)
        self.pixels_set(2, color)
        self.pixels_set(3, color)
        
        self.pixels_show()
    pass

    def pixels_set(self, i, color):
        self.ar[i] = (color[1]<<16) + (color[0]<<8) + color[2]
    pass

    def pixels_fill(self, color):
        for i in range(len(self.ar)) :
            self.pixels_set(i, color)
        pass
    pass

    def color_chase(self, color, wait):
        for i in range(self.num):
            self.pixels_set(i, color)
            sleep(wait)
            self.pixels_show()
        sleep(0.2)
    pass
     
    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)
    pass
     
     
    def rainbow_cycle(self, wait):
        for j in range(255):
            for i in range(self.num):
                rc_index = (i * 256 // self.num) + j
                self.pixels_set(i, self.wheel(rc_index & 255))
            self.pixels_show()
            sleep(wait)
        pass
    pass

pass

if __name__=='__main__':
    
    rgbLed = RGBLed()
    print("fills")
    for color in rgbLed.COLORS:       
        rgbLed.pixels_fill(color)
        rgbLed.pixels_show()
        sleep(0.5)

    print("chases")
    for color in strip.COLORS:       
        rgbLed.color_chase(color, 0.05)

    print("rainbow")
    while True :
        rgbLed.rainbow_cycle(0.05)
    pass

pass
