# Example based on pio_ws2812.py, but modified to use
# a DMA channel to push out the data to the PIO

import array, time
from machine import Pin
import rp2

# Configure the number of WS2812 LEDs.
NUM_LEDS = 9

SM = 0

PIO0_BASE = 0x50200000
PIO1_BASE = 0x50300000
PIO0_BASE_TXF0 = PIO0_BASE+0x10
PIO0_BASE_TXF1 = PIO0_BASE+0x14
PIO0_BASE_TXF2 = PIO0_BASE+0x18
PIO0_BASE_TXF3 = PIO0_BASE+0x1c

PIO1_BASE_TXF0 = PIO1_BASE+0x10
PIO1_BASE_TXF1 = PIO1_BASE+0x14
PIO1_BASE_TXF2 = PIO1_BASE+0x18
PIO1_BASE_TXF3 = PIO1_BASE+0x1c


if SM < 4:
    TFX = PIO0_BASE+0x10 + SM * 4
else:
    TFX = PIO1_BASE+0x10 + (SM-4) * 4
    


@rp2.asm_pio(
    sideset_init=rp2.PIO.OUT_LOW,
    out_shiftdir=rp2.PIO.SHIFT_LEFT,
    autopull=True,
    pull_thresh=24,
)
def ws2812():
    # fmt: off
    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()
    # fmt: on


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

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

# Dummy data
ar = array.array("I", [0 for _ in range(NUM_LEDS)])
r = 0xff
g = 0xff
b = 0xff
ar[0] = 0xff000000
ar[1] = 0x00ff0000
ar[2] = 0x0000ff00
ar[3] = r<<24 + g<<16 + b<<8


# starts the transfer, using DMA channel 0
def dma_out():
    import dma, uctypes
    dma.init_channels()
    d0=dma.CHANNELS[0]
    d0.CTRL_TRIG.EN = 0
    d0.TRANS_COUNT = NUM_LEDS
    d0.READ_ADDR = uctypes.addressof(ar)
    d0.WRITE_ADDR = TFX
    d0.CTRL_TRIG.INCR_WRITE = 0
    d0.CTRL_TRIG.INCR_READ = 1
    d0.CTRL_TRIG.DATA_SIZE = 2
    d0.CTRL_TRIG.EN = 1
    
dma_out()
time.sleep(1)