from artiq import *


my_range = range


class CompilerTest(AutoContext):
    parameters = "a b A B"

    def print_done(self):
        print("Done!")

    def set_some_slowdev(self, n):
        print("Slow device setting: {}".format(n))

    @kernel
    def run(self, n, t2):
        for i in my_range(n):
            self.set_some_slowdev(i)
            delay(100*ms)
            with parallel:
                with sequential:
                    for j in my_range(3):
                        self.a.pulse((j+1)*100*MHz, 20*us)
                        self.b.pulse(100*MHz, t2)
                with sequential:
                    self.A.pulse(100*MHz, 10*us)
                    self.B.pulse(100*MHz, t2)
        self.print_done()


def main():
    from artiq.coredevice import comm_dummy, core, dds

    coredev = core.Core(comm_dummy.Comm())
    exp = CompilerTest(
        core=coredev,
        a=dds.DDS(core=coredev, dds_sysclk=1*GHz,
                  reg_channel=0, rtio_switch=0),
        b=dds.DDS(core=coredev, dds_sysclk=1*GHz,
                  reg_channel=1, rtio_switch=1),
        A=dds.DDS(core=coredev, dds_sysclk=1*GHz,
                  reg_channel=2, rtio_switch=2),
        B=dds.DDS(core=coredev, dds_sysclk=1*GHz,
                  reg_channel=3, rtio_switch=3)
    )
    exp.run(3, 100*us)

if __name__ == "__main__":
    main()