from slic.core.acquisition import SFAcquisition
from slic.core.acquisition.sfacquisition import BSChannels, transpose_dicts, print_response
from slic.core.task import DAQTask
from slic.devices.timing.events import CTASequencer
import epics
import numpy as np
import time 


class CTAAcquisition(SFAcquisition):

    cta = CTASequencer("SAT-CCTA-ESE")

    def acquire(self, filename, data_base_dir=None, detectors=None, channels=None, pvs=None, scan_info=None, n_pulses=100, n_repeat=1, is_scan_step=False, wait=True):
        if not is_scan_step:
            run_number = self.client.next_run()
            print(f"Advanced run number to {run_number}.")
        else:
            run_number = self.client.run_number
            print(f"Continuing run number {run_number}.")

        if not filename or filename == "/dev/null":
            print("Skipping retrieval since no filename was given.")
            return

        if detectors is None:
            print("No detectors specified, using default detector list.")
            detectors = self.default_detectors

        if pvs is None:
            print("No PVs specified, using default PV list.")
            pvs = self.default_pvs

        if channels is None:
            print("No channels specified, using default channel list.")
            channels = self.default_channels
        print(filename, epics.caget('DPO:AcquireData'))
        bschs = BSChannels(*channels)
        bschs.check()

        client = self.client
        client.set_config(n_pulses, filename, detectors=detectors, channels=channels, pvs=pvs, scan_info=scan_info)
        
        
        def _acquire():
            # set n_pulses to scope PV / sleep?
            # set filename to scope PV
            epics.caput('DPO:DemodulateData', 0)
            while epics.caget('DPO:DemodulateData'):
                time.sleep(0.1)

            epics.caput('DPO:Repetitions', 1)
            epics.caput('DPO:FrameNumber', n_pulses)

            filename_for_scope = f"run{run_number:04}-{filename}"
            epics.caput('DPO:RunFolder', filename_for_scope)

            epics.caput('DPO:AcquireData', 1)
            epics.caput('DPO:DemodulateData', 1)
            time.sleep(1)
            
            self.cta.stop()
            self.cta.cfg.repetitions = n_pulses
            self.cta.run()
            
            start_pid = self.cta.get_start_pid()
            #print('CTA finished', start_pid)
            while epics.caget('DPO:AcquireData'):
                time.sleep(0.1)
            
            #start_pid_new = self.cta.get_start_pid()

            #print('Scope finished', start_pid_new, (start_pid_new-start_pid)/100, self.cta.cfg.repetitions)
            stop_pid = start_pid + n_pulses
            pids = np.arange(start_pid, stop_pid)

            bs = list(pids[::1000])
            bs.append(stop_pid)
            for ii in range(len(bs)-1):

                block = np.arange(bs[ii], bs[ii+1]-1)
                res = self.retrieve(filename, block, run_number=run_number)

                res = transpose_dicts(res)
                filenames = res.pop("filenames")
                print_response(res)

            return filenames

        def stopper():
            client.stop()
            self.cta.stop()
            epics.caput('DPO:AcquireData', 0)

        task = DAQTask(_acquire, stopper=stopper, filename=filename, hold=False)
        self.current_task = task

        if wait:
            try:
                task.wait()
            except KeyboardInterrupt:
                print("Stopped current DAQ task:")

        return task