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

class MockPMT():
    def __init__(self):
        pass

    def count_on(self):
        return poisson.rvs(200)

    def count_off(self):
        return poisson.rvs(10)

class MedicionesReferencia(EnvExperiment):
    """TEST DE REALTIME APPLET"""
    def build(self):
        self.setattr_device("ccb") # PAra armar el applet

        self.setattr_argument("tau",
                              NumberValue(1., min=1e-3, max=1., unit="s", scale=s),
                              tooltip="Tiempo de recolección de cuentas")
        self.pmt = MockPMT() # Emulo el PMT prendido/apagado

    def create_applets(self):
        self.ccb.issue("create_applet", "cuentas_segundo",
                        "${artiq_applet}realtime "
                        "200 "
                        "count_on "
                        "--y2 count_off")

    def create_datasets(self):
        self.set_dataset('count_on', np.full(1, np.nan), broadcast=True, archive=False)
        self.set_dataset('count_off', np.full(1, np.nan), broadcast=True, archive=False)
        self.set_dataset('xs', np.full(1, 0), broadcast=True, archive=False)

    def run(self):
        self.create_datasets()
        self.create_applets()
        while True:
            self.mutate_dataset("count_on", 0, self.pmt.count_on())
            self.mutate_dataset("count_off", 0, self.pmt.count_off())
            time.sleep(self.tau)

