import datetime, glob, os, logging, traceback, time

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

from hdf5.hdf5_tdms import HDF5_tdms

logger = logging.getLogger(__name__)

class CompactDAQ():

    def __init__(self, log_folder=None, enable_hdf5_logging=False):
        self._analog_task = None                 # nidaqmx task object
        self._analog_stream_reader = None        # stream reader function
        
        self._encoder_task = None                # same as above
        self._encoder_stream_reader = 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 = log_folder                   # folder where the tdms files are written
        self._hdf = None
        self.enable_hdf5_logging = enable_hdf5_logging

        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 8 updates per second
        self._buffer_size = 80*self._samples_per_callback  # set the read buffer size to tolerate 10 seconds of delay (10 sec x 8 updates per sec)

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

        if log_folder:
            self.setup_task(log_folder)
        else:
            logger.error('No log folder specified.')

        
    def __enter__(self):
        return self


    def __exit__(self, exc_type, exc_val, exc_tb):
        if exc_type or exc_val or exc_tb:
            traceback.print_exception(exc_type, exc_val, exc_tb)
        logger.debug('In DAQ __exit__.')
        self.stop_logging()
        self.stop_task()
        self.clear_task()
        self.remove_index_files()
        return False # rethrow
    

    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)
        if self.enable_hdf5_logging:
            self._hdf = HDF5_tdms(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)


    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

        if self.enable_hdf5_logging:
            self._hdf.new_log()


    def stop_logging(self):
        logger.debug('Stopping logging.')
        self._analog_task.in_stream.logging_pause = True
        self._encoder_task.in_stream.logging_pause = True
        if self.enable_hdf5_logging:
            self._hdf.stop()
            
        
    def flush_log(self):
        if self.enable_hdf5_logging:
            self._hdf.flush()


    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. '''
        return self.get_time_for_timestamp(self._callback_counter), self._sensor_data
    

    def get_time_for_timestamp(self, ts):
        ''' Returns dT from the start of acquisition in seconds corresponding to input timestamp.'''
        dT = ts * 1/self._actual_sampling_rate * self._samples_per_callback
        return dT        


    def get_time_vector_for_timestamp(self, ts):
        ''' Returns time vector of length samples_per_callback whose last value corresponds to the input timestamp.'''
        v = np.linspace(self.get_time_for_timestamp(ts - 1), 
                        self.get_time_for_timestamp(ts),
                        self._samples_per_callback)
        return v


    def remove_index_files(self):
        logger.debug('Removing index files.')
        prototype = os.path.join(self._log_folder,"*.tdms_index")
        for f in glob.glob(prototype):
            os.remove(f)


    def _callback(self, task_handle, event_type, n_samples, dummy):
        hdftime = None
        try:
            start = time.time()
            # read samples & increment counter
            a_samples_read = self._analog_stream_reader.read_many_sample(data=self._sensor_data[0:3,:], number_of_samples_per_channel=n_samples, timeout=0.1)
            e_samples_read = self._encoder_stream_reader.read_many_sample_double(data=self._sensor_data[3,:], number_of_samples_per_channel=n_samples, timeout=0.1)
            if a_samples_read == self._samples_per_callback and e_samples_read == self._samples_per_callback:
                self._callback_counter += 1.0

                # write to hdf log
                if self.enable_hdf5_logging:
                    hdfstart = time.time()
                    self._hdf.append(self.get_time_vector_for_timestamp(self._callback_counter), self._sensor_data)
                    hdftime = time.time()-hdfstart
            else:
                logger.debug(f'Read {a_samples_read} on the analog channel, and {e_samples_read} on the encoder channel.')

#            logger.debug(f'Callback time {time.time()-start} hdf {hdftime}.')
            return 0
        except:
            logger.exception(f'Exception in DAQ callback. Time {time.time()-start} hdf {hdftime}')
            return 0
