from artiq.experiment import *
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())