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


class PMTCalibration(EnvExperiment):
    """Not working - Chunk blinking histograms measurement for pmt calibration """

    def build(self):

        self.setattr_device("ccb")
        self.setattr_device("core")
        
        self.pmt = self.get_device("ttl0")

        self.laserUV = UrukulCh(self, ch=0, freq=110.0, amp=0.3) #0.7 Vpp
        self.laserIR = UrukulCh(self, ch=1, freq=208.0, amp=0.35) #0.8 Vpp

        self.setattr_argument("Total_Number_measurements", NumberValue(20, ndecimals=0, step=1))
        self.setattr_argument("Shown_measurements", NumberValue(5, ndecimals=0, step=1))
        self.setattr_argument("Chunks", NumberValue(5, ndecimals=0, step=1))

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

        self.lim = 0

    @rpc
    def process(self, i=0):

        self.measurement = self.get_dataset("partial_chunk")
        self.mean_value_off = np.mean(self.measurement[0:self.Chunks])
        self.mean_value_on = np.mean(self.measurement[self.Chunks:])
        self.mean_diff = self.mean_value_on-self.mean_value_off

        self.mutate_dataset("mean_value_off_vec", i, self.mean_value_off)
        self.mutate_dataset("mean_value_on_vec", i, self.mean_value_on)
        self.mutate_dataset("mean_diff_vec", i, self.mean_diff)

#        return [self.mean_value_off, self.mean_value_on, self.mean_diff]

    @rpc
    def create_datasets(self):

        self.set_dataset("partial_chunk", np.zeros(2*self.Chunks, dtype=int), broadcast=True, archive=False)

        self.set_dataset("mean_value_off_vec", np.zeros(self.Shown_measurements, dtype=int), broadcast=True, archive=False)
        self.set_dataset("mean_value_on_vec", np.zeros(self.Shown_measurements, dtype=int), broadcast=True, archive=False)
        self.set_dataset("mean_diff_vec", np.zeros(self.Shown_measurements, dtype=int), broadcast=True, archive=False)

        self.set_dataset("bins", np.arange(0, 600, 5), broadcast=True, archive=False)


    @rpc
    def create_applets(self):

        self.ccb.issue("create_applet", "calibration_fluo_diff",
                        "${artiq_applet}plot_xy "
                        "mean_diff_vec",
                        group="blinking_applets")

        # Custom histogram plot creator (needs a reasonable --update-delay)
        self.ccb.issue("create_applet", "calibration_fluo_diff_histo",
                        "${python} -m pyLIAF.artiq.applets.histogram "
                        "partial_chunk --x bins "
                        "--update-delay 0.5",
                        group="blinking_applets")        

    @kernel
    def run(self):

        self.create_datasets()
        self.create_applets()
        self.init_kernel()
        delay(100*ms)
        i = 0
        while self.lim < self.Total_Number_measurements:
            self.measure_counts()
            self.process(i)
            self.lim = self.lim+1
                                      
            if i < self.Shown_measurements - 1:
                i = i + 1
            else:  
                i = 0
            #print(f'Done {self.lim}')
        print("That's all folks!")
    

    @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)

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

    @kernel
    def init_kernel(self):
        self.core.reset()
        self.pmt.input()
        self.pmt_state.output()
        delay(1*s)
        self.laserIR.initialize_channel()
        self.laserUV.initialize_channel()
        self.pmt_state.off()
        self.core.wait_until_mu(now_mu())

    @kernel
    def measure_counts(self):

#        bin_boundaries = np.arange(self.bin_init, self.bin_end, self.bin_step)
#        self.set_dataset("hd_bins", bin_boundaries,
#                         broadcast=True, archive=False)

        self.core.break_realtime()
        self.laserIR.off()
        for j in range(self.Chunks):
            self.core.break_realtime()
            self.mutate_dataset("partial_chunk", j, self.readout())
        self.core.break_realtime()

        self.laserIR.on()          
  
        for j in range(self.Chunks, 2*self.Chunks):
            self.core.break_realtime()
            self.mutate_dataset("partial_chunk", j, self.readout())

