import numpy as np
from artiq.experiment import *
from pyLIAF.artiq.controllers import UrukulCh


class PMTCalibration(EnvExperiment):
    """Blinking script plotting for pmt calibration """

    def build(self):
        self.setattr_device("ccb")
        self.setattr_device("scheduler")
        self.setattr_device("core")
        self.pmt = self.get_device("ttl0")

        self.laserUV = UrukulCh(self, ch=2, freq=110.0, amp=0.3, name="UV") #corresponde a 0.7 Vpp
        self.laserIR = UrukulCh(self, ch=1, freq=208.0, amp=0.35, name="IR") #corresponde a 0.8 Vpp

        self.setattr_argument(f"t_readout",
                              NumberValue(100*ms, unit='ms', scale=ms, min=1*ms),
                             "PMT Calibration")

    @rpc
    def create_datasets(self):
        self.set_dataset("pmt_counts_on", np.zeros(1, dtype=int), broadcast=True, archive=False)
        self.set_dataset("pmt_counts_off", np.zeros(1, dtype=int), broadcast=True, archive=False)

    @rpc
    def create_applets(self):
        self.ccb.issue("create_applet", "calibration_pmt_blinking",
                        "${python} -m pyLIAF.artiq.applets.realtime "
                        "200 "
                        "pmt_counts_on "
                        "--y2 pmt_counts_off")

    @kernel
    def init_kernel(self):
        self.core.reset()
        self.laserIR.initialize_channel()
        self.laserUV.initialize_channel()
        self.pmt.input()

        delay(10*us)

        self.laserIR.set_channel()
        self.laserUV.set_channel()

        self.core.wait_until_mu(now_mu())

        delay(1*ms)
        self.laserIR.on()
        self.laserUV.on()


    def run(self):
        self.create_datasets()
        self.create_applets()
        self.init_kernel()

        try:
            while True:
                self.measure_counts()
                while self.scheduler.check_pause():
                    print("PAUSED BLINKING")
                    self.core.comm.close()
                    self.scheduler.pause()
                    # TODO: reset freqs/amps
                    print("RESUMED BLINKING")
        except TerminationRequested:
            self.cleanup()
            print("STOPPED BLINKING")

    @kernel
    def readout(self):
        """Registro de cuentas emitidas"""
       
        here = self.pmt.gate_rising(self.t_readout)

        cuentas = self.pmt.count(here)
        delay(1*us)
        return cuentas

    @kernel
    def measure_counts(self):
        self.core.break_realtime()
        self.laserUV.on()
        self.laserIR.off()

        self.mutate_dataset("pmt_counts_off", 0, self.readout())
        at_mu(self.core.get_rtio_counter_mu() + 10000)

        self.laserIR.on()          
        self.mutate_dataset("pmt_counts_on", 0, self.readout())

    @kernel
    def cleanup(self):
        self.core.break_realtime()
        self.laserIR.off()
        self.laserUV.off()

