from artiq.experiment import *

class UrukulCh(HasEnvironment):

    def build(self, ch=0, freq=100.0, amp=1.0):
        self.ch = ch
        self.amp = amp
        self.freq = freq
        
        self.channel = self.get_device(f"urukul0_ch{self.ch}")
        
        self.frequency = self.get_argument(f"freq_ch{self.ch}",
                            NumberValue(self.freq*MHz, unit='MHz', scale=MHz, min=1*MHz, max=400*MHz),
                            f"channel_{self.ch}")

        self.amplitude = self.get_argument(f"amp_ch{self.ch}",
                            NumberValue(self.amp, min=0., max=1.),
                            f"channel_{self.ch}")

    @kernel
    def initialize_channel(self):
        # initialises CPLD the selected channel
        self.channel.cpld.init()
        self.channel.init()

    @kernel
    def set_channel(self):
        self.channel.set(self.frequency, amplitude=self.amplitude)

    @kernel
    def set_o(self, state):
        self.channel.sw.set_o(state)

class TestUnderflow(EnvExperiment):
    def build(self):
        self.setattr_device("core")
        self.laserIR = UrukulCh(self, ch=1, freq=208.0, amp=0.35)
        
    @kernel
    def run(self):
        self.core.reset()
        self.laserIR.initialize_channel()
        self.core.break_realtime()
        self.laserIR.set_channel()
        
        delay(10*us)

        # Just try to pulse the laser
        self.laserIR.set_o(True)
        delay(1*ms)
        self.laserIR.set_o(False)
