from artiq.experiment import *


class SoftwareEdgeCount(EnvExperiment):
    def build(self):
        self.setattr_device("core")
        self.ttlin = self.get_device("ttl0")
        self.ttlout = self.get_device("ttl7")

    @kernel
    def run(self):
        self.core.reset()
        gate_start_mu = now_mu()
        # Start input gate & advance timeline cursor to gate_end_mu
        gate_end_mu = self.ttlin.gate_rising(1*ms)
        at_mu(gate_start_mu)

        for _ in range(64):
            self.ttlout.pulse(8*ns)
            delay(8*ns)

        counts = self.ttlin.count(gate_end_mu)
        print(counts)


class EdgeCounter(EnvExperiment):
    def build(self):
        self.setattr_device("core")
        self.edgecounter0 = self.get_device("ttl0_counter")

    @kernel
    def run(self):
        self.core.reset()
        self.edgecounter0.gate_rising(1*ms)
        counts = self.edgecounter0.fetch_count()
        print(counts)


class ExternalTrigger(EnvExperiment):
    def build(self):
        self.setattr_device("core")
        self.ttlin = self.get_device("ttl0")
        self.ttlout = self.get_device("ttl4")

    @kernel
    def run(self):
        self.core.reset()
        gate_end_mu = self.ttlin.gate_rising(5*ms)
        timestamp_mu = self.ttlin.timestamp_mu(gate_end_mu)
        at_mu(timestamp_mu + self.core.seconds_to_mu(10*ms))
        self.ttlout.pulse(1*us)


import time


class MeanTimestampDuration(EnvExperiment):
    def build(self):
        self.setattr_device("core")
        self.ttlin = self.get_device("ttl0")
        self.ttlclk = self.get_device("ttl7")
    
    @kernel
    def get_timestamp_duration(self, pulse_num) -> TInt64:
        self.core.break_realtime()
        delay(1*ms)
        gate_start_mu = now_mu()
        # Start input gate & advance timeline cursor to gate_end_mu
        gate_end_mu = self.ttlin.gate_rising(1*ms)
        at_mu(gate_start_mu)

        self.ttlclk.set_mu(0x800000)
        delay(16*pulse_num*ns)
        self.ttlclk.set_mu(0)

        # Guarantee t0 > gate_end_mu
        # Otherwise timestamp_mu may wait for pulses till gate_end_mu
        rtio_time_mu = self.core.get_rtio_counter_mu()
        sleep_mu = float(gate_end_mu - rtio_time_mu)
        self.rpc_sleep(self.core.mu_to_seconds(sleep_mu))

        t0 = self.core.get_rtio_counter_mu()
        while self.ttlin.timestamp_mu(gate_end_mu) >= 0:
            pass
        t1 = self.core.get_rtio_counter_mu()
        return t1 - t0

    @rpc
    def rpc_sleep(self, duration):
        time.sleep(duration)

    @kernel
    def run(self):
        self.core.reset()
        t64 = self.get_timestamp_duration(64)
        t8 = self.get_timestamp_duration(8)
        print("Mean timestamp_mu duration:")
        print(self.core.mu_to_seconds((t64 - t8)/((64 + 1) - (8 + 1))))