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

# TODO:
#     [ ] Revisar los tiempos
#     [ ] Ver por que da overflow intermitentemente en las corrids
#     [ ] Ver como esta guardando los resultados y guardar lo que falta
#     [ ] Cambiarle los parametros a los laseres cuando arranca el exp

class SingleLine(EnvExperiment):
    """SP Transition Emission Experiment"""
    def build(self):
        self.setattr_device("core")
        self.setattr_device("ccb")       

        self.pmt = self.get_device("ttl0")
        self.pmt_state = self.get_device("ttl4")

        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("bin", NumberValue(50e-9, unit='us'), "Binning params")

        self.setattr_argument("no_measures",
                              NumberValue(100, min=1, ndecimals=0, step=1),
                              "Experiment params")

        self.setattr_argument(f"t_prepS",
                              NumberValue(10*us, unit='us', scale=us, min=1*us),
                             "Experiment params")

        self.setattr_argument(f"t_cool",
                              NumberValue(10*us, unit='us', scale=us, min=1*us),
                             "Experiment params")

        self.setattr_argument(f"t_readout",
                              NumberValue(5*us, unit='us', scale=us, min=1*us),
                             "Experiment params")


    @rpc
    def create_datasets(self):
        self.set_dataset("counts", [],
                            broadcast=True, archive=True)
        self.set_dataset("t_readout", self.t_readout, broadcast=False, archive=True)
        self.set_dataset("t_prep_S", self.t_prepS, broadcast=False, archive=True)
        self.set_dataset("t_enfriar_ion", self.t_cool, broadcast=False, archive=True)
        self.set_dataset("no_measures", self.no_measures, broadcast=False, archive=True)
        self.laserIR.generate_dataset()
        self.laserUV.generate_dataset()
        #self.set_dataset("frec_UV",self.frec_UV, broadcast=False,archive=True)
        #self.set_dataset("frec_IR",self.frec_IR, broadcast=False,archive=True)
        # TODO: Agregar forma de guardar los datos de los canales del Urukul.
        # o bien guardando todos aca, o armando un metodo apropiado en su controlador
        
        #self.set_dataset("binvector", np.arange(0, self.t_readout, self.bin*1e-6), broadcast=True, archive=False)


    @rpc
    def create_applets(self):
        self.ccb.issue("create_applet", "cuentas",
                        "${python} -m pyLIAF.artiq.applets.histogram "
                        "counts "
                        "--update-delay 0.2")


    @kernel
    def run(self):
        self.create_datasets()
        self.create_applets()
        self.init_kernel()
        delay(1*ms)

        self.enfriar_ion()
        for runN in range(self.no_measures):
            # Aprovecho el tiempo que necesito para extraer datos
            # par enfriar al ion con los dos laseres prendidos
            at_mu(self.core.get_rtio_counter_mu() + self.core.seconds_to_mu(self.t_cool) )
            self.prep_S()
            t_end = self.readout()
            self.save_counts(t_end)
            #self.mutate_dataset("counts", runN, counts)

        self.cleanup()

    @kernel
    def init_kernel(self):
        self.core.reset()
        self.pmt.input()
        self.pmt_state.output()
        self.laserIR.initialize_channel()
        self.laserUV.initialize_channel()
        # Quizas haya errores de tiempo.
        self.pmt_state.off()
        self.laserIR.set_channel()
        self.laserUV.set_channel()
        self.core.wait_until_mu(now_mu())


    @kernel
    def enfriar_ion(self):
        """Preparo el ion prendiendo ambos laseres"""
        self.laserIR.on()
        self.laserUV.on()
        #delay(self.t_cool)

    @kernel
    def prep_S(self):
        """Preparo el estado S prendiendo solamente el laser IR"""
        self.laserIR.on()
        self.laserUV.off()
        delay(self.t_prepS)


    @kernel
    def readout(self):
        """Registro cuentas emitidas con el laser UV prendido"""
        self.laserIR.off()
        self.laserUV.on()
        
        # Prendo y apago la TTL para ver en el osc.
        self.pmt_state.on()
        here = self.pmt.gate_rising(self.t_readout)
        self.pmt_state.off()

        #self.pmt_state.pulse(self.t_readout)
        #cuentas = self.pmt.count(here)
        #delay(1*us)
        self.enfriar_ion()
        return here

    @kernel
    def save_counts(self, t_end):
        count = self.pmt.timestamp_mu(t_end)
        t0 = t_end - self.core.seconds_to_mu(self.t_readout)
        while count > 0:
            self.append_to_dataset("counts", self.core.mu_to_seconds(count - t0) )
            count = self.pmt.timestamp_mu(t_end)


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