import numpy as np
import nidaqmx as ni
import nidaqmx.stream_readers
import nidaqmx.constants as niconstants
import queue

class CDaq(object):
    def __init__(self):

        self.channels = []
        self.task = None
        self.task_dig = None
        self.stream = None
        self.stream_dig = None
        self.databuff = queue.Queue(maxsize=128)
        self.databuff_dig = queue.Queue(maxsize=128)
        
        self.dataarray = None
        self.dataarray_dig= None
        self.datahandler = self.default_data_handler
        self.datahandler_dig = self.default_data_handler_dig
        self.running = False
        self.nchan = 0
        self.nchan_ana = 0
        self.nchan_dig = 0
        self.nsamples = 0

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.clear_task()

    def add_channel(self, channel: str, type: str, name: str):
        """
        Add a channel to the running list of channels that will be configured in the task before starting it.

        :param channel: string name for nidaqmx channel, see list_devices and NI MAX utility for determining these
        :param type: string type 'voltage' or 'current'
        :param name: string channel reference name
        :return: none
        """

        self.channels.append(dict(chan=channel, type=type, name=name))

    def setup_task(self, samplerate: float, nsamples: int):
        """
        Using `nidaqmx`, setup the acquisition task for continuous acquisition at the given samplerate, and
        break data into segments for a given number of samples.

        :param samplerate: float value of sampling rate in hertz
        :param nsamples: integer number of samples that will trigger a callback event once `start_acquisition` is called.
        :return: none
        """

        # Check object state
        if self.task is not None or self.task_dig is not None:
            raise ReferenceError("task already created, call 'clear_task' first")

        if self.datahandler is self.default_data_handler:
            print("Warning: using default data handler")

        # Determine n channels
        self.nchan = len(self.channels)
        self.nsamples = nsamples
        self.nchan_ana = 0
        self.nchan_dig = 0
        for channel in self.channels:
            if channel['type'] is 'voltage' or channel['type'] is 'current':
                self.nchan_ana += 1
            elif channel['type'] is 'ang_encoder':
                self.nchan_dig += 1
        
        # Analog Task
        self.task = ni.Task()

        for channel in self.channels:
            if channel['type'] is 'voltage':
                self.task.ai_channels.add_ai_voltage_chan(channel['chan'], name_to_assign_to_channel=channel['name'])
            elif channel['type'] is 'current':
                self.task.ai_channels.add_ai_current_chan(channel['chan'], name_to_assign_to_channel=channel['name'])

        # Analog Setup Timing
        self.task.timing.cfg_samp_clk_timing(rate=samplerate, samps_per_chan=nsamples, sample_mode=niconstants.AcquisitionType.CONTINUOUS)

        self.task.register_every_n_samples_acquired_into_buffer_event(nsamples, self.datahandler)
        self.task.in_stream.input_buf_size = 5*nsamples
        self.stream = ni.stream_readers.AnalogMultiChannelReader(self.task.in_stream)
        self.dataarray = np.ndarray([self.nchan_ana, nsamples])

        #Encoder Task
        self.task_dig = ni.Task()

        for channel in self.channels:
            if channel['type'] is 'ang_encoder':
                self.task_dig.ci_channels.add_ci_ang_encoder_chan(channel['chan'],
                    name_to_assign_to_channel=channel['name'],
                    decoding_type=ni.constants.EncoderType.X_1)

        self.task_dig.timing.cfg_samp_clk_timing(rate=samplerate, samps_per_chan=nsamples, sample_mode=niconstants.AcquisitionType.CONTINUOUS)
        self.task_dig.register_every_n_samples_acquired_into_buffer_event(nsamples, self.datahandler_dig)
        self.task_dig.in_stream.input_buf_size = 5*nsamples
        self.stream_dig = ni.stream_readers.CounterReader(self.task_dig.in_stream)
        if self.nchan_dig is 1:
            self.dataarray_dig = np.ndarray([160,])
        else:
            self.dataarray_dig = np.ndarray([self.nchan_dig, nsamples])


    def assign_datahandler(self, callback):
        """
        Assign a callback function that will work with the nidaqmx driver to execute when n samples are in the daq
        buffer. This must be called prior to ``setup_task()``

        :param callback: function reference that matches the prototype set in the ``nidaqmx.task.register_every_n_samples_acquired_into_buffer_event()``
        :return: none
        """
        self.datahandler = callback

    def default_data_handler(self, task_handle, event_type, samps_acq, callback_data):
        """
        This is the default data handler sent to the nidaqmx driver for acquisition of every n samples. See ``nidaqmx.task.register_every_n_samples_acquired_into_buffer_event()``
        """
        if self.stream is None:
            raise ReferenceError("Ingestion stream not initialized, run setup_task() first")
        # off load data
        samps_read = self.stream.read_many_sample(data=self.dataarray, number_of_samples_per_channel=self.nsamples, timeout=10.0)
        # put into queue
        if self.databuff.full() == False:
            #print("[Data] samps acq: ", samps_acq, "samps read: ", samps_read)
            self.databuff.put(self.dataarray)
        else:
            print("Data queue full!")

        return 0

    def default_data_handler_dig(self, task_handle, event_type, samps_acq, callback_data):
        """
        This is the default data handler sent to the nidaqmx driver for acquisition of every n digital samples. See ``nidaqmx.task.register_every_n_samples_acquired_into_buffer_event()``
        """
        if self.stream is None:
            raise ReferenceError("Ingestion stream not initialized, run setup_task() first")
        # off load data
        samps_read = self.stream_dig.read_many_sample_double(data=self.dataarray_dig, number_of_samples_per_channel=self.nsamples, timeout=10.0)
        # put into queue
        if self.databuff_dig.full() == False:
            #print("[Data] samps acq: ", samps_acq, "samps read: ", samps_read)
            self.databuff_dig.put(self.dataarray_dig)
        else:
            print("Data queue full!")

        return 0

    def start_acquisition(self):
        """
        Start the nidaqmx task
        """
        if self.task is not None and self.task_dig is not None:
            self.task.start()
            self.task_dig.start()
            self.running = True
        else:
            raise ReferenceError("task not initialized")

    def stop_acquisition(self):
        """
        Stop the nidaqmx task
        """
        if self.task is not None  and self.task_dig is not None:
            self.task.stop()
            self.task_dig.stop()
            self.running = False
        else:
            raise ReferenceError("task not initialized")

    def clear_task(self):
        """
        Stops data acquisition and clears the task
        """
        if self.task is not None and self.task_dig is not None:
            self.stop_acquisition()
            self.running = False
            self.task.close()
            self.task_dig.close()
        else:
            raise ReferenceError("task not initialized")


if __name__ == "__main__":
    """
    Test to open a task and acquire analog and angular encoder data
    """

    def process_data(data: np.ndarray):
        voltage0 = np.mean(data[0, :])
        voltage1 = np.mean(data[1, :])
        voltage2 = np.mean(data[2, :])
        print("V0: {:5.2f} V1: {:5.2f} V2: {:5.2f}".format(voltage0, voltage1, voltage2))

    def process_data_dig(data: np.ndarray):
        counter0 = np.mean(data)
        print("C0: {:5.2f}".format(counter0))
        
    import time
    
    with CDaq() as cdaq:
        cdaq.add_channel('cDAQ1Mod1/ai0', 'voltage', 'Voltage0')
        cdaq.add_channel('cDAQ1Mod1/ai1', 'voltage', 'Voltage1')
        cdaq.add_channel('cDAQ1Mod1/ai2', 'voltage', 'Voltage2')
        cdaq.add_channel('cDAQ1Mod3/ctr0', 'ang_encoder', 'Encoder0')

        cdaq.setup_task(samplerate=1600.0, nsamples=160)

        cdaq.start_acquisition()

        n = 0
        try:
            while n < 30:
                if not cdaq.databuff.empty():
                    data = cdaq.databuff.get()
                    if data is None:
                        continue
                    process_data(data)
                    n += 1
                if not cdaq.databuff_dig.empty():
                    data = cdaq.databuff_dig.get()
                    process_data_dig(data)
                    n += 1

        except KeyboardInterrupt:
            cdaq.stop_acquisition()

        cdaq.stop_acquisition()