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)

from artiq.coredevice import spi2 as spi

SPI_CONFIG = (0 * spi.SPI_OFFLINE | 0 * spi.SPI_END |
              0 * spi.SPI_INPUT | 0 * spi.SPI_CS_POLARITY |
              0 * spi.SPI_CLK_POLARITY | 0 * spi.SPI_CLK_PHASE |
              0 * spi.SPI_LSB_FIRST | 0 * spi.SPI_HALF_DUPLEX)

CLK_DIV = 125

class SPIWrite(EnvExperiment):
    def build(self):
        self.setattr_device("core")
        self.spi = self.get_device("dio_spi0")

    @kernel
    def run(self):
        self.core.reset()
        self.spi.set_config_mu(SPI_CONFIG, 8, CLK_DIV, 0b110)
        self.spi.write(0x13 << 24)      # Shift the bits to the MSBs.
                                        # Since SPI_LSB_FIRST is NOT set,
                                        # SPI Machine will shift out bits from
                                        # the MSB of the `data` register.`
        self.spi.set_config_mu(SPI_CONFIG | spi.SPI_END, 32, CLK_DIV, 0b110)
        self.spi.write(0xDEADBEEF)


class SPIRead(EnvExperiment):
    def build(self):
        self.setattr_device("core")
        self.spi = self.get_device("dio_spi0")

    @kernel
    def run(self):
        self.core.reset()
        self.spi.set_config_mu(SPI_CONFIG, 8, CLK_DIV, 0b001)
        self.spi.write(0x81 << 24)      # Shift the bits to the MSBs.
                                        # Since SPI_LSB_FIRST is NOT set,
                                        # SPI Machine will shift out bits from
                                        # the MSB of the `data` register.`
        self.spi.set_config_mu(SPI_CONFIG | spi.SPI_END | spi.SPI_INPUT,
                                32, CLK_DIV, 0b001)
        self.spi.write(0)               # write() performs the SPI transfer.
                                        # As suggested by the timing diagram,
                                        # the exact value of this argument
                                        # does not matter.
        print(self.spi.read())