# NOTES:
# * Check DRTIO channel list by inspecting Metlino's log.
# * Prepare an artiq_route .CFG to route Sayma#1 and #2 as LINK#1,2 and #3,4.

from artiq.experiment import *

def build_st1st2(exp):
    assert isinstance(exp, EnvExperiment)

    exp.setattr_device("core")
    # Test with SyncDDS channels
    exp.basemods_1 = [exp.get_device("basemod_att"+str(i)) for i in range(2)]
    exp.rfsws_1 = [exp.get_device("sawg_sw"+str(i)) for i in range(8)]
    exp.basemods_2 = [exp.get_device("basemod_att"+str(i)) for i in range(2, 4)]
    exp.rfsws_2 = [exp.get_device("sawg_sw"+str(i)) for i in range(8, 16)]
    # Test with 2 TTL channels
    exp.ttl_1 = exp.get_device("ttl0")
    exp.ttl_2 = exp.get_device("ttl2")
    # Bundled DUTs
    exp.basemods = exp.basemods_1 + exp.basemods_2
    exp.rfsws = exp.rfsws_1 + exp.rfsws_2
    exp.ttls = [exp.ttl_1, exp.ttl_2]


class ST1(EnvExperiment):
    def build(self):
        build_st1st2(self)

    @kernel
    def drtio_is_up(self, drtio_index):
        if not self.core.get_rtio_destination_status(drtio_index):
            return False
        print("DRTIO #", drtio_index, "is ready\n")
        return True

    @kernel
    def run(self):
        print("*** Waiting for DRTIO ready...")
        drtio_indices = [1, 2, 3, 4]
        for i in drtio_indices:
            while not self.drtio_is_up(i):
                pass
        print("*** All DRTIO ready !")
        
        self.core.reset()

        for basemod in self.basemods:
            basemod.reset()
            delay(10*ms)
            basemod.set(6.0, 6.0, 6.0, 6.0)
            delay(10*ms)
            # Printing shows that the BaseMod is being controlled properly.
            print(basemod.get_mu())
            delay(500*ms)

        self.core.break_realtime()
        for rfsw in self.rfsws:
            rfsw.on()
            delay(1*ms)

        print("All RF switches are on.")

        while True:
            pass


class ST2(EnvExperiment):
    def build(self):
        build_st1st2(self)
        # Option to switch TTL output source from RTIO PHY to FPGA (bypassing RTIO)
        self.setattr_argument("ttl_use_fpga", BooleanValue(False))

    @kernel
    def drtio_is_up(self, drtio_index):
        if not self.core.get_rtio_destination_status(drtio_index):
            return False
        print("DRTIO #", drtio_index, "is ready\n")
        return True

    @kernel
    def run(self):
        print("*** Waiting for DRTIO ready...")
        drtio_indices = [1, 2, 3, 4]
        for i in drtio_indices:
            while not self.drtio_is_up(i):
                pass
        print("*** All DRTIO ready !")
        
        self.core.reset()

        for basemod in self.basemods:
            basemod.reset()
            delay(10*ms)
            basemod.set(6.0, 6.0, 6.0, 6.0)
            delay(10*ms)
            # Printing shows that the BaseMod is being controlled properly.
            print(basemod.get_mu())
            delay(500*ms)

        self.core.break_realtime()
        for rfsw in self.rfsws:
            rfsw.on()
            delay(1*ms)

        print("All RF switches are on.")

        if self.ttl_use_fpga:
            while True:
                pass
        else:
            # test pulse
            self.core.break_realtime()
            for ttl in self.ttls:
                ttl.output()
                delay(1*ms)

            print("TTLs are now outputting pulses.")
            self.core.break_realtime()
            while True:
                with parallel:
                    for ttl in self.ttls:
                        ttl.on()
                delay(111*ns)
                with parallel:
                    for ttl in self.ttls:
                        ttl.off()
                delay(4444*ns)