from time import sleep, time
from scipy.stats import poisson
import numpy as np
from artiq.experiment import *

class MockPMT(): #{{{
    def __init__(self, mean=20):
        self.mean = mean

    def count(self) -> TInt32:
        medida = poisson.rvs(self.mean)
        return medida

    def change_mean(self, new_mean):
        self.mean = new_mean

    # }}}

class MockLaser():
    def __init__(self, pmt):
        self.pmt = pmt

    def on(self):
        self.pmt.change_mean(500)
        sleep(0.1)

    def off(self):
        self.pmt.change_mean(100)
        sleep(0.1)


class Histograms(EnvExperiment):
    """Histograms mockpmt simulation blinking v2"""

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

        self.setattr_device("ccb")
        self.setattr_device("core")

    def analyze(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
        return self.mean_value_off, self.mean_value_on, self.mean_diff     

    
    def run(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)

        print(self.get_dataset("bins"))

        self.ccb.issue("create_applet", "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", "fluo_diff_histo",
                        "${artiq_applet}histogram "
                        "partial_chunk --x bins "
                        "--update-delay 0.5",
                        group="blinking_applets")

        self.pmt = MockPMT(100)
        self.laser = MockLaser(self.pmt)

        i = 0
        self.lim = 0

        while self.lim<self.Total_Number_measurements:
            self.run_kernel()
            mean_off, mean_on, mean_diff = self.analyze(i)
            self.mutate_dataset("mean_value_off_vec", i, mean_off)
            self.mutate_dataset("mean_value_on_vec", i, mean_on)
            self.mutate_dataset("mean_diff_vec", i, mean_diff)
            sleep(0.1)
            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 run_kernel(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.laser.off()

        for j in range(self.Chunks):
            self.mutate_dataset("partial_chunk", j, self.pmt.count())

        self.laser.on()            
        for j in range(self.Chunks, 2*self.Chunks):
            self.mutate_dataset("partial_chunk", j, self.pmt.count())



sleep(0.3)
