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


class PMTCalibration(EnvExperiment):
    """Non blinking script for pmt checking with both lasers on"""

    def build(self):

        self.setattr_device("ccb")
        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_nonblinking",
                        "${python} -m pyLIAF.artiq.applets.realtime "
                        "200 "
                        "pmt_counts_on "
                        "--y2 pmt_counts_off")

    @kernel
    def init_kernel(self):
        self.core.reset()
        self.pmt.input()
        delay(1*s)
        print(1)
        self.laserIR.initialize_channel()
        self.core.break_realtime()
        print(2)
        self.laserUV.initialize_channel()
        self.core.break_realtime()
        print(3)
        self.core.wait_until_mu(now_mu())


    @kernel
    def run(self):

        self.create_datasets()
        self.create_applets()
        self.init_kernel()
        delay(100*ms)
        i = 0
        self.core.break_realtime()
        self.laserIR.on()
        self.core.break_realtime()
        self.laserUV.on()
        self.core.break_realtime()
        while True:
            self.measure_counts()

    @kernel
    def readout(self):
        """Registro cuentas emitidas con el laser UV prendido"""
       
        # Prendo y apago la TTL para ver en el osc.

        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.mutate_dataset("pmt_counts_off", 0, self.readout())
#        self.core.break_realtime()

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

