# Author: Michael Sachs
# Date: 11/25/2023
# Version: 1.0
# Description:
#   This code allows a Raspberry Pi Pico to communicate with a ws2812 LED light ring.
#   When a button is pressed, more LEDS light up until the full ring is lit, at which point
#   the ring cycles through the rainbow a few times. Then, the button causes the lights
#   to turn off. When all LEDs are off, the light ring flashes between off/white.
#
# Credits:
#   Portions of this code have been adapted from an example program located here:
#       https://github.com/geeekpi/picokitadv/blob/main/demo_codes/Project_10_Light_up_ws2812_light_Ring.py


import array, time
from machine import Pin
import rp2

# Configure the WS2812 LEDs.
NUM_LEDS = 16           # Total number of LEDs present on the ring
LED_PIN_NUM = 6         # The pin on the PI being used for LED PIO
ALL_LEDS = -1           # An indicator constant for all LEDs needing to be updated.

BUTTON_PIN_NUM = 13     # The pin on the PI being used for button input.

# RGB specifications of the colors
OFF = (0, 0, 0)
RED = (255, 0, 0)
ORANGE = (255, 103, 0)
YELLOW = (255, 150, 0)
GREEN = (0, 255, 0)
BLUE = (0, 0, 255)
INDIGO = (75, 0, 130)
VIOLET = (180, 0, 255)
WHITE = (255, 255, 255)
RAINBOW = (RED, ORANGE, YELLOW, GREEN, BLUE, INDIGO, VIOLET)

# Setup the PIO loop to be used for communicating 3 8-bit values to each LED
# This is a standard ws2812 loop, and is found in multiple places in the
# Raspberry Pi Pico micropython documentation.
@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()

# Create the StateMachine with the ws2812 program, outputting on LED_PIN_NUM
sm = rp2.StateMachine(0, ws2812, freq=8_000_000, sideset_base=Pin(LED_PIN_NUM))

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

# This is an array integers representing 0GRB for each LED
active_array = array.array("I", [0 for _ in range(NUM_LEDS)])
backup_array = array.array("I", [0 for _ in range(NUM_LEDS)])

# This function shifts the active array of 0GRB values into the state machine FIFO queue
# as 24-bit GRB values, which will be read by the PIO state machine and passed to the ws2812 LED ring.
def display():
    # Put each 4 byte integer into the state machine register
    # First, left shift out the unused 8 bits so only 3 bytes are put in the state machine
    sm.put(active_array, 8)
    # Pause for 10ms to make it easier to see the changes
    time.sleep_ms(10)

# The ws2812 LED expects the color in GRB format, but we usually specify RGB format.
# This function converts an RGB tuple into a single 32-bit 0GRB value.
def grb_from_rgb(color):
    return ((color[1]<<16) + (color[0]<<8) + color[2])

# This function sets a specified LED in the active array to a 0GRB value specified by
# the RGB tuple passed in. If the LED is -1, all values in the active array are set.
def set_led(color, led):
    # Convert the color to 0GRB
    grb_color = grb_from_rgb(color)

    # Determine if all LEDs, need to be set.
    if led != ALL_LEDS:
        active_array[led] = grb_color
    else:
        for i in range(NUM_LEDS):
            active_array[i] = grb_color

# This function backs up the colors of the active_array, in order to restore them later on
def store_active_array():
    for i in range(NUM_LEDS):
        backup_array[i] = active_array[i]

# This function restores the active_array from the previously backed up array.
def restore_active_array():
    for i in range(NUM_LEDS):
        active_array[i] = backup_array[i]

# This function turns on a specific LED
# The color depends on the number of LEDs already lit
# It returns the number of LEDs lit after the current one is turned on
def turn_on_led(led):
    color = None
    # If less than 33% lit, choose GREEN
    if led / NUM_LEDS <= 0.33:
        color = GREEN

    # If less than 67% lit, choose YELLOW
    elif led / NUM_LEDS <= .67:
        color = YELLOW
    
    # Otherwise, choose RED
    else:
        color = RED

    # Set the LED to the specified color, and indicate
    # how many LEDs are currently lit
    set_led((color), led)
    return led + 1

# This function turns off a specified LED
# It returns the number of LEDs lit after the current one is turned off
def turn_off_led(led):
    set_led(OFF, led - 1)
    return led - 1

# This function flashes all LEDs.
# If rainbow is true, it flashes through the colors of the rainbow 3 times.
# Otherwise, it flashes black and white.
def flash_leds(rainbow):

    # Back up the active_array to restore after flashing
    store_active_array()

    # Determine which flash to use
    if rainbow:
        # Flash through the rainbow 3 times (7 colors)
        for i in range(0,21):
            # Determine which color to flash
            color = RAINBOW[i % 7]

            # Set ALL LEDs to the color, and display it for .3 seconds
            set_led(color, ALL_LEDS)
            display()
            time.sleep(.3)
    else:
        # Flash black and white 21 times
        for i in range(0, 21):

            # Choose whether the LEDs should be black (off) or on (white)
            color = None
            if i % 2 == 0:
                color = OFF
            else:
                color = WHITE

            # Set ALL LEDs to the color, and display it for .3 seconds
            set_led(color, ALL_LEDS)
            display()
            time.sleep(.3)

    # Restore the backup array
    restore_active_array()

# State data for current LED state
leds_lit = 0
light_led = True

# Button used to trigger LEDs
button = Pin(BUTTON_PIN_NUM, Pin.IN, Pin.PULL_UP)

# Indefinitely loop for input
while True:

    # If the button is pressed, trigger light sequence
    if button.value() == 0:
        # If LEDs are being lit, turn on the next one
        if light_led:
            leds_lit = turn_on_led(leds_lit)
        # If LEDs are being turned off, turn off the next one
        else:
            leds_lit = turn_off_led(leds_lit)

        # If all LEDs are off/on, switch the order
        if leds_lit == 0 or leds_lit == NUM_LEDS:
            light_led = not light_led
            
        # If all LEDs are off, flash black/white
        if leds_lit == 0:
            flash_leds(False)
        # If all LEDs are on, flash rainbow
        elif leds_lit == NUM_LEDS:
            flash_leds(True) 

        # Display the current LED sequence
        display()
        
    # Wait .2 seconds before seeing if the button is still pressed.
    time.sleep(0.2)