import machine
from rp2 import asm_pio, PIO
from machine import Pin
from FastADC import FastADC
import time

@asm_pio()
def glitch():
    wait(1, irq, 1)
    irq(block, 0)
    wrap_target()
    # emit the glitches here
    wrap()
    push(block)

@asm_pio()
def tio_trigger():
    wait(1, irq, 2)
    irq(block, 1)
    push(block)
    
@asm_pio()
def pre_trigger():
    wait(0, pin, 0)
    wait(1, pin, 0)
    irq(block, 2)
    push(block)

class Test:
    def __init__(self):
        self.fastadc = FastADC()
        self.adc_samples_captured = False
        self.fastsamples = None
        self.pio_base = 1
        self.sm0 = PIO(self.pio_base).state_machine(0)
        self.sm1 = PIO(self.pio_base).state_machine(1)
        self.sm2 = PIO(self.pio_base).state_machine(2)
        self.pin_trigger = Pin(14, Pin.IN)
        self.led_pin = Pin(7, Pin.OUT)
        
    def arm(self):
        self.sm0.init(glitch, freq=machine.freq())
        self.sm1.init(tio_trigger, freq=machine.freq())
        self.sm2.init(pre_trigger, freq=machine.freq(), in_base=self.pin_trigger)
        self.sm0.active(1)
        self.sm1.active(1)
        self.sm2.active(1)
        self.sm0.irq(handler=self.__nop_handler)
        
    def arm_adc(self):
        self.adc_samples_captured = False
        self.sm0.irq(handler=self.__poll_fast_adc)

    @micropython.native
    def __nop_handler(self, sm):
        pass

    @micropython.native
    def __poll_fast_adc(self, sm):
        self.fastsamples = self.fastadc.read()
        self.adc_samples_captured = True
        self.led_pin.value(not self.led_pin.value())
    
    def block(self):
        # wait for the statemachine to finish
        while self.sm0.rx_fifo() == 0:
            pass
        
    def get_adc_samples(self):
        # wait for the adc to capture samples
        while not self.adc_samples_captured:
            pass
        print(self.fastsamples)
        
reset_pin = Pin(2, Pin.OUT)
test = Test()
while True:
    test.arm()
    print("armed")
    #test.arm_adc()
    reset_pin.value(0)
    reset_pin.value(1)
    test.block()
    print("triggered")
    #test.get_adc_samples()
    time.sleep(0.1)

