from artiq.experiment import *
from scipy import signal
import numpy


class Voltage(EnvExperiment):
    def build(self):
        self.setattr_device("core")
        self.zotino = self.get_device("zotino0")

    def prepare(self):
        self.channels = [0, 1, 2, 3]
        self.voltages = [1.0, 2.0, 3.0, 4.0]

    @kernel
    def run(self):
        self.core.reset()
        self.core.break_realtime()
        self.zotino.init()

        delay(1*ms)
        self.zotino.set_dac(self.voltages, self.channels)


class TriangularWave(EnvExperiment):
    def build(self):
        self.setattr_device("core")
        self.zotino = self.get_device("zotino0")

    def prepare(self):
        self.period = 0.1*s
        self.sample = 128
        t = numpy.linspace(0, 1, self.sample)
        self.voltages = 8*signal.sawtooth(2*numpy.pi*t, 0.5)
        self.interval = self.period/self.sample

    @kernel
    def run(self):
        self.core.reset()
        self.core.break_realtime()
        self.zotino.init()

        delay(1*ms)

        counter = 0
        while True:
            self.zotino.set_dac([self.voltages[counter]], [0])
            counter = (counter + 1) % self.sample
            delay(self.interval)
