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):
        return poisson.rvs(self.mean)

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

    # }}}

class Histograms(EnvExperiment):
    """Histograms mockpmt simulation"""
    def build(self):
        
        self.setattr_argument("npoints", NumberValue(100, ndecimals=0, step=5))
        self.setattr_argument("chunks", NumberValue(5, ndecimals=0, step=1))

        self.setattr_argument("bin_init", NumberValue(-10, ndecimals=0, step=1), "Binning parameters")
        self.setattr_argument("bin_end", NumberValue(30, ndecimals=0, step=1), "Binning parameters")
        self.setattr_argument("bin_step", NumberValue(1, ndecimals=0, step=1), "Binning parameters")

        self.setattr_argument("nbins", NumberValue(100, ndecimals=0, step=10))
        self.setattr_argument("histo_delay", NumberValue(0.6, ndecimals=2, step=0.01))

        self.setattr_device("ccb")

    def run(self):

        self.pmt = MockPMT()

        # Custom histogram plot creator (needs a reasonable --update-delay)
        self.ccb.issue("create_applet", "mockpmt_histogram",
                        "${artiq_applet}plot_hist "
                        "histo --x hd_bins --title mockpmt "
                        f"--update-delay {self.histo_delay}",
                        group="ExampleHistogram")

        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.set_dataset("histo", np.zeros(len(bin_boundaries)-1, dtype=int), broadcast=True, archive=False)

        self.set_dataset("medidas", np.empty(self.npoints*self.chunks), broadcast=True, archive=False)
        self.set_dataset("partial_chunk", np.empty(self.chunks), broadcast=True, archive=False)

        for i in range(self.npoints):

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

            print(1)
            histogram_partial, _ = np.histogram(self.partial_chunk,
                                        bin_boundaries)
            print(2)
            for k in range(len(histo)):
                self.mutate_dataset("histo", k, self.histo[k]+histogram_partial[k])
            print(3)
            sleep(0.5)



sleep(0.3)
