from artiq.experiment import *

class UrukulCh(HasEnvironment):
    """Urukul single freq class

    Set the frecuencies/amplitudes of each Urukul channel
    """

    def build(self, ch=0, freq=100.0, amp=1.0, name=None):
#        super().build()  ch, amp, freq
        self.ch = ch
        self.amp = amp
        self.freq = freq
        self.name = name

        # sets core device drivers as attributes
        self.setattr_device("core")
        
        if name==None:
            self.channel = self.get_device(f"urukul0_ch{self.ch}")

            ### This two attributes will be shown in the GUI grouped by channel
            # use/don't use each channel
            self.state = self.get_argument(f"state_ch{self.ch}", BooleanValue(1==0), f"canal_{self.ch}")
            # each channel's frequency
            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"canal_{self.ch}")
            # each channel's amplitude
            self.amplitude = self.get_argument(f"amp_ch{self.ch}",
                                NumberValue(self.amp, min=0., max=1.),
                                f"canal_{self.ch}")

        else:
            self.channel = self.get_device(f"urukul0_ch{self.ch}")

            ### This two attributes will be shown in the GUI grouped by channel
            # use/don't use each channel
            self.state = self.get_argument(f"state_ch{self.ch}", BooleanValue(1==0), name)
            # each channel's frequency
            self.frequency = self.get_argument(f"freq_ch{self.ch}",
                                NumberValue(self.freq*MHz, unit='MHz', scale=MHz, min=1*MHz, max=400*MHz),
                                name)
            # each channel's amplitude
            self.amplitude = self.get_argument(f"amp_ch{self.ch}",
                                NumberValue(self.amp, min=0., max=1.),
                                name)

    @rpc
    def generate_dataset(self):
        self.set_dataset(f"laser_{self.name}", '{'+f'"freq": {self.frequency}, "amp": {self.amplitude}"' + '}')

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

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

        if self.state == True:
            self.channel.sw.on()
        else:
            self.channel.sw.off()

    @kernel
    def on(self):
        self.channel.sw.on()

    @kernel
    def off(self):
        self.channel.sw.off()

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

    @kernel
    def pulse_mu(self, delay_time_mu):
        self.channel.sw.pulse_mu(delay_time_mu)

    @kernel
    def pulse(self, delay_time):
        self.channel.sw.pulse(delay_time)
