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

class BlinkAndWait(EnvExperiment):
    """Promedio Blinqueos y espero"""
    def prepare(self):
        # Set the name of the experiment to run in background
        expid_1 = {
                "file": "Experiments/turn_lasers_on.py",
                "class_name": "TurnOnLasers",
                "arguments": {},
                "log_level": self.scheduler.expid["log_level"],
                "repo_rev": self.scheduler.expid["repo_rev"],
                }

    def build(self):
        self.setattr_device("core")
        self.setattr_device("sheduler")

        self.pmt = self.get_device("ttl0")
        self.laserUV = UrukulCh(self, ch=2, freq=110.0, amp=0.3, name="UV") #corresponde a 0.7 Vpp
        self.laserIR = UrukulCh(self, ch=1, freq=208.0, amp=0.35, name="IR") #corresponde a 0.8 Vpp

        self.setattr_argument("t_exposure", NumberValue(50*us, unit='us'))
        self.setattr_argument("no_measures", NumberValue(100, min=1, ndecimals=0, step=1))

    @kernel
    def create_datasets(self):
        self.set_dataset("counts_on", [], broadcast=True, archive=True)
        self.set_dataset("counts_off", [], broadcast=True, archive=True)

    @kernel
    def init_kernel(self):
        self.core.reset()
        self.laserUV.initialize_channel()
        self.laserIR.initialize_channel()
        self.pmt.input()
        delay(100*ns)
        self.laserUV.set_channel()
        self.laserIR.set_channel()
        self.core_break_realtime()

    @kernel{flags="fast-math"}
    def measure(self):
        partial_sum_On = 0
        partial_sum_Off = 0

        for n in range(self.no_measures):
            partial_sum_On += self.pmt.count(t_exposure)
            delay(1*us)
            self.laserIR.off()
            partial_sum_Off += self.pmt.count(t_exposure)
            delay(1*us)
            self.laserIR.on()
            delay(10*us)

        partial_sum_On /= self.no_measures
        partial_sum_Off /= self.no_measures

        return [partial_sum_On, partial_sum_Off]

    @kernel
    def turn_both_on(self):
        self.core.break_realtime()
        self.laserUV.on()
        self.laserIR.on()

    def run(self):
        self.create_datasets():
        self.init_kernel()
        try:
            while True:
                if not self.scheduler.check_pause():
                    self.turn_both_on()
                    datos = self.measure()
                    self.append_to_dataset("counts_on", datos[0])
                    self.append_to_dataset("counts_off", datos[1])
                    self.turn_both_on()

                    # Opcion 1: (Input en terminal)
                    # self.wait_for_go()

                    # Opcion 2: (Scheduler)
                    self.scheduler.pause()
                    self.scheduler.submit("main", expid_1, priority=10)
                    sleep(0.5) # quizas lo necesito para que no se programen muchos
                    self.scheduler.pause()

        except TerminationRequested:
            print("BLINKING END")

    @rpc
    def wait_for_go(self):
        while input("GO?") != 'go':
            pass
        self.scheduler.pause()
