import rp2
from machine import Pin

import pio

STRAND0_PIN=6
STRAND1_PIN=7
STRAND2_PIN=8
INPUT_PIN=15
RESET_PIN=20


def main():
  f = 32_000_000

  in_pin = Pin(INPUT_PIN, Pin.IN, Pin.PULL_DOWN)
  reset_pin = Pin(RESET_PIN, Pin.IN, Pin.PULL_DOWN)
  reset_out_pin = Pin(RESET_PIN+1, Pin.OUT, value=0)

  strand0_pin = Pin(STRAND0_PIN, Pin.OUT, value=0)
  strand0_sm = rp2.StateMachine(0, pio.strand_switcher, in_base=in_pin, jmp_pin=reset_pin, set_base=strand0_pin, freq=f)
  strand0_sm.put(9*24-1)

  strand1_pin = Pin(STRAND1_PIN, Pin.OUT, value=0)
  strand1_sm = rp2.StateMachine(1, pio.strand_switcher, in_base=in_pin, jmp_pin=reset_pin, set_base=strand1_pin, freq=f)
  strand1_sm.put(9*24-1)

  strand2_pin = Pin(STRAND2_PIN, Pin.OUT, value=0)
  strand2_sm = rp2.StateMachine(2, pio.strand_switcher, in_base=in_pin, jmp_pin=reset_pin, set_base=strand2_pin, freq=f)
  strand2_sm.put(9*24-1)

  reset_out_pin = Pin(RESET_PIN+1, Pin.OUT, value=0)
  rst_sm = rp2.StateMachine(3, pio.reset_monitor, jmp_pin=in_pin, set_base=reset_out_pin, freq=f)
  rst_sm.put(int(50e-6 * f // 2)-1)

  strand0_sm.active(1)
  strand1_sm.active(1)
  strand2_sm.active(1)
  rst_sm.active(1)

main()





# Testing code
@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()


def test_input(pin, cycles=1):
  import array, time
  fake = rp2.StateMachine(7, ws2812, freq=8_000_000, sideset_base=pin)
  NUM_LEDS = 9*3
  ar = array.array("I", [0 for _ in range(NUM_LEDS)])
  fake.active(1)

  for _ in range(cycles):
    # Cycle colours.
    for i in range(4 * NUM_LEDS):
      for j in range(NUM_LEDS):
        r = j * 100 // (NUM_LEDS - 1)
        b = 100 - j * 100 // (NUM_LEDS - 1)
        if j != i % NUM_LEDS:
          r >>= 3
          b >>= 3
        ar[j] = r << 16 | b
      fake.put(ar, 8)
      time.sleep_ms(50)
    #Fade
    for i in range(24):
      for j in range(NUM_LEDS):
          ar[j] >>= 1
      fake.put(ar, 8)
      time.sleep_ms(50)
  fake.active(0)

# All debug from here
#start = lambda: rst_sm.exec('irq(clear, 0)')
test_pin = Pin(14, Pin.OUT, value=0)
#import time
#time.sleep_ms(int(10e6))
test_input(test_pin)