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

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

    def count(self):
        return poisson.rvs(self.mean)

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

    # }}}

class PMTCountRealtime(EnvExperiment):
    """Realtime PMT counts visualization"""
    def build(self):
        self.setattr_argument("no_measures", NumberValue(10, min=1, ndecimals=0, step=1))
        self.setattr_argument("freqs", Scannable(
                                        default=CenterScan(200*MHz, 20*MHz, 1*MHz),
                                        unit="MHz",
                                        scale=MHz,
                                        global_min = 1*MHz,
                                        global_max = 400*MHz
                                       )
                             )
        self.setattr_device("ccb") # To issue the applets
        self.pmt = MockPMT(mean=6)

    def run(self):
        self.set_dataset("current_freq", np.full(1, 0.), broadcast=True, archive=False)

        # Applet creation {{{
        self.ccb.issue("create_applet", "output_frecuency",
                        "${artiq_applet}big_number "
                        "current_freq",
                        group="realtime_pmt")

        self.ccb.issue("create_applet", "ticks_plot",
                        "${artiq_applet}plot_xy "
                        "ticks",
                        group="realtime_pmt")

        # Custom histogram plot creator (needs a reasonable --update-delay)
        self.ccb.issue("create_applet", "ticks_histogram",
                        "${artiq_applet}histogram "
                        "ticks "
                        "--update-delay 0.5",
                        group="realtime_pmt")
        # }}}s

        # # Change freq and read {{{
        # # this emulates a change in frequency, then many pmt readings
        # # and storing (displaying) the mean
        # self.set_dataset("ticks", np.full(len(self.freqs), np.nan),
        #                  broadcast=True, archive=False)
        # counts = 0
        # for fn, curr_freq in enumerate(self.freqs):
        #     self.mutate_dataset("current_freq", 0, curr_freq)
        #     # self.pmt.change_mean(60 + 5*fn)
        #     counts = 0
        #     for i in range(self.no_measures):
        #         counts += self.pmt.count()
        #         time.sleep(20e-6)
        #     self.mutate_dataset("ticks", fn, counts/self.no_measures)
        # # }}}

        # PMT count reading {{{
        # This emulates just reading from the PMT: e.g. dark counts distr
        self.set_dataset("ticks", np.full(self.no_measures, np.nan),
                         broadcast=True, archive=False)
        for i in range(self.no_measures):
            self.mutate_dataset("ticks", i, self.pmt.count())
            time.sleep(20e-3)
        # }}}
