import datetime, glob, os, logging

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

logger = logging.getLogger(__name__)

class CompactDAQ():

    def __init__(self, log_folder):
        self._analog_task = None                 # nidaqmx task object
        self._analog_stream_reader = None        # stream reader function
        self._analog_array = None                # numpy array for sensor data
        
        self._encoder_task = None                # same as above
        self._encoder_stream_reader = None
        self._encoder_array= None

        self._sensor_data = None                # block of most recent sensor data
        self._callback_counter = 0.0            # a counter that is incremented each time the callback is invoked
        self._log_folder = ""                   # folder where the tdms files are written

        self._sampling_rate = 1600.0                 # note: the analog card will round this up to 50,000/31
        self._actual_sampling_rate = 50000/31
        self._samples_per_callback = 200             # shooting for 4 updates per second
        self._buffer_size = 10*self._samples_per_callback  # set the read buffer size 

        self._voltage_channel = 'cDAQ1Mod1/ai0'
        self._current_channel = 'cDAQ1Mod1/ai1'
        self._str_pot_channel = 'cDAQ1Mod1/ai2'
        self._encoder_channel = 'cDAQ1Mod2/ctr0'

        self.setup_task(log_folder)

        
    def __enter__(self):
        return self


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


    def setup_task(self, log_folder):
        sampling_rate = self._sampling_rate
        samples_per_callback = self._samples_per_callback
        buffer_size = self._buffer_size
        self._sensor_data = np.zeros((4, samples_per_callback))

        self._create_analog_task(sampling_rate, buffer_size)
        self._create_encoder_task(sampling_rate, buffer_size)
        self._slave_encoder_timebase()
        self._setup_encoder_start_trigger()
        self._setup_read_callback(samples_per_callback)
        self._setup_tdms_logging(log_folder)

        
    def _create_analog_task(self, sampling_rate, buffer_size):
        # 1) create the analog task
        logger.debug('Creating analog task.')
        try:
            self._analog_task = ni.Task()       
            self._analog_task.ai_channels.add_ai_voltage_chan(self._voltage_channel, name_to_assign_to_channel='voltage', min_val=-10, max_val=10)
            self._analog_task.ai_channels.add_ai_voltage_chan(self._current_channel, name_to_assign_to_channel='current', min_val=-10, max_val=10)
            self._analog_task.ai_channels.add_ai_voltage_chan(self._str_pot_channel, name_to_assign_to_channel='str_pot', min_val=-10, max_val=10)
            
            self._analog_task.timing.cfg_samp_clk_timing(sample_mode=niconstants.AcquisitionType.CONTINUOUS, rate=sampling_rate, samps_per_chan=buffer_size)
            self._analog_task.in_stream.input_buf_size = buffer_size
            self._analog_stream_reader = ni.stream_readers.AnalogMultiChannelReader(self._analog_task.in_stream)
        except:
            logger.exception('Analog task creation failed.')
            raise


    def _create_encoder_task(self, sampling_rate, buffer_size):
        # 2) create the encoder task
        logger.debug('Creating encoder task.')
        try:
            self._encoder_task = ni.Task()
            self._encoder_task.ci_channels.add_ci_ang_encoder_chan(self._encoder_channel, name_to_assign_to_channel='encoder', decoding_type=ni.constants.EncoderType.X_4) # by default: in deg @ 96 tks per rev
            actual_sampling_rate = self._analog_task.timing.samp_clk_rate    # use the actual rate otherwise the two tasks will drift
            self._encoder_task.timing.cfg_samp_clk_timing(sample_mode=niconstants.AcquisitionType.CONTINUOUS, rate=actual_sampling_rate, samps_per_chan=buffer_size)
            self._encoder_task.in_stream.input_buf_size = buffer_size
            self._encoder_stream_reader = ni.stream_readers.CounterReader(self._encoder_task.in_stream)
        except:
            logger.exception('Encoder task creation failed.')
            raise


    def _slave_encoder_timebase(self):
        # 3) slave the encoder timebase
        logger.debug('Slaving the encoder timebase.')
        self._analog_task.control(ni.constants.TaskMode.TASK_RESERVE)                                       # reserve a specific timing engine 
        self._encoder_task.timing.samp_clk_timebase_src  = self._analog_task.timing.samp_clk_timebase_src   # export the timebase source
        self._encoder_task.timing.samp_clk_timebase_rate = self._analog_task.timing.samp_clk_timebase_rate  # export the timebase rate


    def _setup_encoder_start_trigger(self):
        # 4) setup the encoder start trigger
        logger.debug('Setting up the encoder start trigger.')
        trigger_src = self._analog_task.timing.samp_clk_term  # use the rising edge of the analog card sampling clock signal
        self._encoder_task.triggers.start_trigger.cfg_dig_edge_start_trig(trigger_src)
       

    def _setup_read_callback(self, samples_per_callback):
        # 5) setup the read callback
        logger.debug('Setting up the read callback.')
        self._analog_task.register_every_n_samples_acquired_into_buffer_event(samples_per_callback, self._callback)
        self._analog_array = np.ndarray([3, samples_per_callback])
        self._encoder_array = np.ndarray([samples_per_callback,])


    def _setup_tdms_logging(self, log_folder):
        # 6) setup tdms logging
        self._log_folder = log_folder
        filename = os.path.join(log_folder,"dummy_a.tdms")
        filename2 = os.path.join(log_folder,"dummy_e.tdms")
        logger.debug('Setting up TDMS logging to %s and %s.', filename, filename2)
        try:
            self._analog_task.in_stream.configure_logging(file_path=filename, group_name="analog")   # no choice here but to create a dummy file; this will is deleted when the task completes
            self._analog_task.in_stream.logging_pause = True
        except:
            logger.exception('Failed to start TDMS log %s.', filename)
            raise

        try:
            self._encoder_task.in_stream.configure_logging(file_path=filename2, group_name="encoder")
            self._encoder_task.in_stream.logging_pause = True
        except:
            logger.exception('Failed to start TDMS log %s.', filename2)
            raise


    def start_task(self):
        logger.debug('Starting tasks.')
        self._encoder_task.start()   # start the encoder first, then wait for analog sampling clock to rise
        self._analog_task.start()


    def stop_task(self):
        logger.debug('Stopping tasks.')
        self._analog_task.stop()     
        self._encoder_task.stop()


    def clear_task(self):
        logger.debug('Clearing tasks.')
        self._analog_task.close()
        self._encoder_task.close()
        # remove dummy files
        prototype = os.path.join(self._log_folder,"dummy*")
        for f in glob.glob(prototype):
            os.remove(f)


    def start_logging(self):
        time_str = datetime.datetime.now().strftime("-%Y-%m-%d-%H-%M-%S") # append timestamp to the filename
        filename = os.path.join(self._log_folder, "analog" + time_str + ".tdms")
        filename2 = os.path.join(self._log_folder, "encoder" + time_str + ".tdms")
        logger.debug('Starting logging to %s and %s', filename, filename2)
        
        try:
            self._analog_task.in_stream.start_new_file(filename)
            self._analog_task.in_stream.logging_pause = False
        except:
            logger.exception('Failed to start logging to %s.', filename)
            raise

        try:
            self._encoder_task.in_stream.start_new_file(filename2)
            self._encoder_task.in_stream.logging_pause = False
        except:
            logger.exception('Failed to start logging to %s.', filename2)
            raise


    def stop_logging(self):
        logger.debug('Stopping logging.')
        self._analog_task.in_stream.logging_pause = True
        self._encoder_task.in_stream.logging_pause = True


    def get_data(self):
        return self._sensor_data


    def get_index_and_data(self):
        ''' Returns callback index from the start of acquisition, and the data. '''
        return self._callback_counter, self._sensor_data

    
    def get_time_and_data(self):
        ''' Returns dT from the start of acquisition in seconds, and the data. '''
        dT = self._callback_counter * 1/self._actual_sampling_rate * self._samples_per_callback
        return dT, self._sensor_data
    

    def remove_index_files(self):   
        for f in glob.glob(self._log_folder + "*_index"):
            os.remove(f)


    def _callback(self, task_handle, event_type, n_samples, dummy):
        n_samples_read = self._analog_stream_reader.read_many_sample(data=self._analog_array, number_of_samples_per_channel=n_samples, timeout=10.0)
        n_samples_read = self._encoder_stream_reader.read_many_sample_double(data=self._encoder_array, number_of_samples_per_channel=n_samples, timeout=10.0)
        self._sensor_data = np.concatenate((self._analog_array,self._encoder_array[:,None].T), axis=0)  # cat into 4xN array = [ analog 3xN ; encoder 1xN]
        self._callback_counter += 1.0
        return 0