from artiq.experiment import *


class OnePulsePerSecond(EnvExperiment):
    def build(self):
        self.setattr_device("core")
        self.ttl0 = self.get_device("ttl0")

    @kernel
    def run(self):
        self.core.reset()
        while True:
            self.ttl0.pulse(500*ms)
            delay(500*ms)


class MorseCode(EnvExperiment):
    def build(self):
        self.setattr_device("core")
        self.led = self.get_device("led0")

    def prepare(self):
        # As of ARTIQ-6, the ARTIQ compiler has limited string handling
        # capabilities, so we pass a list of integers instead.
        message = ".- .-. - .. --.-"
        self.commands = [{".": 1, "-": 2, " ": 3}[c] for c in message]

    @kernel
    def run(self):
        self.core.reset()
        for cmd in self.commands:
            if cmd == 1:
                self.led.pulse(100*ms)
                delay(100*ms)
            if cmd == 2:
                self.led.pulse(300*ms)
                delay(100*ms)
            if cmd == 3:
                delay(700*ms)


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

    @kernel
    def run(self):
        self.core.reset()
        gate_end_mu = self.ttl0.gate_rising(1*ms)
        counts = self.ttl0.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)


class ShortPulse(EnvExperiment):
    def build(self):
        self.setattr_device("core")
        self.ttl0 = self.get_device("ttl0")
    
    @kernel
    def run(self):
        self.core.reset()
        delay(6*ns)             # Coarse RTIO period: 0 - 7  ns
        self.ttl0.pulse(3*ns)   # Coarse RTIO period: 8 - 15 ns


class ClockGen(EnvExperiment):
    def build(self):
        self.setattr_device("core")
        self.ttl0 = self.get_device("ttl0")
    
    @kernel
    def run(self):
        self.core.reset()
        self.ttl0.set(62.5*MHz)