# -*- coding: utf-8 -*-
# =============================================================================
# module : awg_context.py
# author : Pierre Heidmann and Matthieu Dartiailh
# license : MIT license
# =============================================================================
"""

"""
from atom.api import Float, observe, set_default
import numpy as np
from .base_context import BaseContext, TIME_CONVERSION


class AWGContext(BaseContext):
    """
    """

    channels = ('Ch1_A', 'Ch1_M1', 'Ch1_M2',
                'Ch2_A', 'Ch2_M1', 'Ch2_M2',
                'Ch3_A', 'Ch3_M1', 'Ch3_M2',
                'Ch4_A', 'Ch4_M1', 'Ch4_M2')

    # Sampling frequency in Hz
    sampling_frequency = Float(1e9)

    time_unit = set_default('mus')

    analogical_channels = set_default(('Ch1_A', 'Ch2_A', 'Ch3_A', 'Ch4_A'))

    logical_channels = set_default(('Ch1_M1', 'Ch2_M1', 'Ch3_M1', 'Ch4_M1',
                                    'Ch1_M2', 'Ch2_M2', 'Ch3_M2', 'Ch4_M2'))

    def compile_sequence(self, pulses, **kwargs):
        """ Transform a sequence of pulse to a dict of waveform.

        Parameters
        ----------
        pulses : list(Pulse)
            List of pulse generated by the compilation of a sequence.

        Returns
        -------
        result : bool
            Boolean indicating whether or not the compilation succeeded.

        to_send or traceback : dict
            Dict of {channel: bytearray} to send to the AWG in case of success
            or the traceback of the issues in case of failure.

        """
        sequence_duration = max([pulse.stop for pulse in pulses])
        # Total length of the sequence to send to the AWG
        if 'sequence_duration' in kwargs:
            if sequence_duration <= kwargs['sequence_duration']:
                sequence_duration = kwargs['sequence_duration']
            else:
                return False, {'Sequence_duration':
                               'Not all pulses fit in given duration'}

        # Collect the channels used in the pulses' sequence
        used_channels = set([pulse.channel[:3] for pulse in pulses])

        # Coefficient to convert the start and stop of pulses in second and
        # then in index integer for array
        time_to_index = TIME_CONVERSION[self.time_unit]['s'] * \
            self.sampling_frequency

        # Length of the sequence
        sequence_length = int(round(sequence_duration * time_to_index))

        # create 3 array for each used_channels
        array_analog = {}
        array_M1 = {}
        array_M2 = {}
        for channel in used_channels:
            # numpy array for analog channels int16 init 2**13
            array_analog[channel] = np.ones(sequence_length,
                                            dtype=np.uint16)*(2**13)
            # numpy array for marker1 init False. For AWG M1 = 0 = off
            array_M1[channel] = np.zeros(sequence_length, dtype=np.int8)
            # numpy array for marker2 init False. For AWG M2 = 0 = off
            array_M2[channel] = np.zeros(sequence_length, dtype=np.int8)

        for pulse in pulses:

            waveform = pulse.waveform
			# ZL
			#waveform = np.ones(len(waveform), dtype=np.int8)
            channel = pulse.channel[:3]
            channeltype = pulse.channel[4:]

            start_index = int(round(pulse.start*time_to_index))
            stop_index = start_index + len(waveform)

            if channeltype == 'A' and pulse.kind == 'Analogical':
#                array_analog[channel][start_index:stop_index] +=\
#                    np.rint(8191*waveform)
                array_analog[channel][start_index:stop_index] +=\
                    (np.rint(8191*waveform)).astype(np.uint16)
            elif channeltype == 'M1' and pulse.kind == 'Logical':
                array_M1[channel][start_index:stop_index] += waveform
            elif channeltype == 'M2' and pulse.kind == 'Logical':
                array_M2[channel][start_index:stop_index] += waveform
            else:
                msg = 'Selected channel does not match kind for pulse {} ({}).'
                return False, {'Kind issue':
                               msg.format(pulse.index,
                                          (pulse.kind, pulse.channel))}

        # Check the overflows
        traceback = {}
        for channel in used_channels:
            analog = array_analog[channel]
            if analog.max() > 16383 or analog.min() < 0:
                mes = 'Analogical values out of range.'
                traceback['{}_A'.format(channel)] = mes

            elif array_M1[channel].max() > 1 or array_M1[channel].min() < 0:
                mes = 'Overflow in marker 1.'
                traceback['{}_M1'.format(channel)] = mes

            elif array_M2[channel].max() > 1 or array_M2[channel].min() < 0:
                mes = 'Overflow in marker 2.'
                traceback['{}_M2'.format(channel)] = mes

        if traceback:
            return False, traceback

        # Invert marked logical channels.
        for i_ch in self.inverted_log_channels:
            ch, m = i_ch.split('_')
            if m == 'M1':
                np.logical_not(array_M1[ch], array_M1[ch])
            else:
                np.logical_not(array_M2[ch], array_M2[ch])

        # Byte arrays to send to the AWG
        to_send = {}
        for channel in used_channels:
            # Convert to sixteen bits integers
            array = array_analog[channel] +\
                array_M1[channel]*(2**14) + array_M2[channel]*(2**15)
            # Creating and filling a byte array for each channel.
            aux = np.empty(2*sequence_length, dtype=np.uint8)
            aux[::2] = array % 2**8
            aux[1::2] = array // 2**8
            to_send[int(channel[-1])] = bytearray(aux)

        return True, to_send
        
    def merge_intervals(self,intervals, sequence_length):
        if intervals[0][0] <= 256:
            intervals[0] = (1, intervals[0][1])
            
        for i, interval in enumerate(intervals):
            if interval[1] - interval[0] < 256:
                if interval[0] + 256 >= sequence_length:
                    intervals[i] = (interval[1] - 256, interval[1])
                else:
                    intervals[i] = (interval[0], interval[0]+256)
        i = 1                    
        while i < len(intervals):
            while i < len(intervals) and intervals[i][0] - intervals[i-1][1] < 256:
                intervals[i-1] = (intervals[i-1][0], max(intervals[i][1], intervals[i-1][1]))
                del intervals[i]
            i+=1
            
        if sequence_length - intervals[-1][1] < 256:
            intervals[-1] = (intervals[-1][0], sequence_length + 1)
        
    def compile_loop(self, pulses, **kwargs):
        """ Transform a sequence of pulse to a dict of waveform suitable for 
        sequence mode of the AWG

        Parameters
        ----------
        pulses : list(Pulse)
            List of pulse generated by the compilation of a sequence.

        Returns
        -------
        result : bool
            Boolean indicating whether or not the compilation succeeded.

        to_send or traceback : dict
            Dict of {channel: array of (bytearrays,intarray,intarray)} to send to the AWG in case of success
            containing waveform, position(s) and repeat(s)
            or the traceback of the issues in case of failure.
            

        """
        sequence_duration = max([pulse.stop for pulse in pulses])
        # Total length of the sequence to send to the AWG
        if 'sequence_duration' in kwargs:
            if sequence_duration <= kwargs['sequence_duration']:
                sequence_duration = kwargs['sequence_duration']
            else:
                return False, {'Sequence_duration':
                               'Not all pulses fit in given duration'}

        # Collect the channels used in the pulses' sequence
        used_channels = set([pulse.channel[:3] for pulse in pulses])

        # Coefficient to convert the start and stop of pulses in second and
        # then in index integer for array
        time_to_index = TIME_CONVERSION[self.time_unit]['s'] * \
            self.sampling_frequency

        # Length of the sequence
        sequence_length = int(round(sequence_duration * time_to_index))

        # intervals are places that have pulses on at least one  channel
        intervals = []     
        for pulse in pulses:
            waveform = pulse.waveform
            start_index = int(round(pulse.start*time_to_index))
            stop_index = start_index + len(waveform)
            intervals.append((start_index, stop_index))
        
        # sort and merge intervals
        intervals.sort()
        i = 1
        while i < len(intervals):
            while i < len(intervals) and intervals[i][0] < intervals[i-1][1]:
                newlast = intervals[i][1]
                del intervals[i]
                intervals[i-1] = (intervals[i-1][0], max(newlast, intervals[i-1][1]))
            i+=1
            
        # check that the minimum length of the each pulse and space between them are > 256 samples
        # fix if the condition is not satisfied
        self.merge_intervals(intervals, sequence_length)

        # make zero space a multiple repeat of 256 zero samples
        zero_length = 256
        azeros = np.ones(zero_length,dtype=np.uint16)*(2**13)
        mzeros = np.zeros(zero_length, dtype=np.int8)
        array_analog = {}
        array_M1 = {}
        array_M2 = {}
        repeats = []
        for channel in used_channels:
            array_analog[channel] = []
            array_M1[channel] = []
            array_M2[channel] = []
        
        # fill up the space between start of sequence and first sample of pulse with zero samples
        if intervals[0][0] != 1:
            zcount = (intervals[0][0] - 1) / zero_length - 1
            zrem = zero_length + (intervals[0][0] - 1) % zero_length
            azrem = np.ones(zrem,dtype=np.uint16)*(2**13)
            mzrem = np.zeros(zrem, dtype = np.int8)
            if zcount > 0:
                repeats.append(zcount)
            repeats.append(1)
            for channel in used_channels:
                if zcount > 0:
                    array_analog[channel].append(azeros)
                    array_M1[channel].append(mzeros)
                    array_M2[channel].append(mzeros)
                array_analog[channel].append(azrem)
                array_M1[channel].append(mzrem)
                array_M2[channel].append(mzrem)
                
        intervals_to_pulses_a = {}
        intervals_to_pulses_m1 = {}
        intervals_to_pulses_m2 = {}
        for channel in used_channels:
            intervals_to_pulses_a[channel] = {}
            intervals_to_pulses_m1[channel] = {}
            intervals_to_pulses_m2[channel] = {}           
            
        for idx,interval in enumerate(intervals):
            for channel in used_channels:
                afill = np.ones(interval[1] - interval[0],dtype=np.uint16)*(2**13)
                mfill1 = np.zeros(interval[1] - interval[0],dtype=np.int8)
                mfill2 = np.zeros(interval[1] - interval[0],dtype=np.int8)
                array_analog[channel].append(afill)
                array_M1[channel].append(mfill1)
                array_M2[channel].append(mfill2)
                intervals_to_pulses_a[channel][interval] = afill
                intervals_to_pulses_m1[channel][interval] = mfill1
                intervals_to_pulses_m2[channel][interval] = mfill2
                
            remaining = 0
            if idx + 1 == len(intervals):
                remaining = sequence_length - interval[1] + 1
            else:
                remaining = intervals[idx+1][0] - intervals[idx][1]
            zcount = remaining / zero_length - 1
            zrem = zero_length + remaining % zero_length
            azrem = np.ones(zrem,dtype=np.uint16)*(2**13)
            mzrem = np.zeros(zrem, dtype = np.int8)
            repeats.append(1)
            if zcount > 0:
                repeats.append(zcount)
            repeats.append(1)
            
            if remaining != 0:
                for channel in used_channels:    
                    if zcount > 0:                    
                        array_analog[channel].append(azeros)
                        array_M1[channel].append(mzeros)
                        array_M2[channel].append(mzeros)
                    array_analog[channel].append(azrem)                    
                    array_M1[channel].append(mzrem)                    
                    array_M2[channel].append(mzrem)                 

        for pulse in pulses:
            waveform = pulse.waveform
            channel = pulse.channel[:3]
            channeltype = pulse.channel[4:]

            start_index = int(round(pulse.start*time_to_index))
            stop_index = start_index + len(waveform)
            interval_i = 0
            while intervals[interval_i][1] < start_index:
                interval_i += 1
            interval = intervals[interval_i]

            if channeltype == 'A' and pulse.kind == 'Analogical':
                intervals_to_pulses_a[channel][interval][start_index - interval[0]:stop_index - interval[0]] +=\
                    (np.rint(8191*waveform)).astype(np.uint16)
            elif channeltype == 'M1' and pulse.kind == 'Logical':
                intervals_to_pulses_m1[channel][interval][start_index - interval[0]:stop_index - interval[0]] += waveform
            elif channeltype == 'M2' and pulse.kind == 'Logical':
                intervals_to_pulses_m2[channel][interval][start_index - interval[0]:stop_index - interval[0]] += waveform
            else:
                msg = 'Selected channel does not match kind for pulse {} ({}).'
                return False, {'Kind issue':
                               msg.format(pulse.index,
                                          (pulse.kind, pulse.channel))}

#        # Check the overflows
        traceback = {}
        for channel in used_channels:
            for i in range(len(array_analog[channel])):
                analog = array_analog[channel][i]
                m1 = array_M1[channel][i]
                m2 = array_M2[channel][i]                
                if analog.max() > 16383 or analog.min() < 0:
                    mes = 'Analogical values out of range.'
                    traceback['{}_A'.format(channel)] = mes
            
                elif m1.max() > 1 or m1.min() < 0:
                    mes = 'Overflow in marker 1.'
                    traceback['{}_M1'.format(channel)] = mes
    
                elif m2.max() > 1 or m2.min() < 0:
                    mes = 'Overflow in marker 2.'
                    traceback['{}_M2'.format(channel)] = mes

        if traceback:
            return False, traceback, None

        # Invert marked logical channels.
        for i_ch in self.inverted_log_channels:
            ch, m = i_ch.split('_')
            if m == 'M1':
                for waveform in array_M1[ch]:
                    np.logical_not(waveform, waveform)
            else:
                for waveform in array_M2[ch]:
                    np.logical_not(waveform, waveform)

        # Byte arrays to send to the AWG
        bytes = {}
        already_added = {}
        for channel in used_channels:
            bytes[int(channel[-1])] = []
            for i in range(len(array_analog[channel])):
                addr = id(array_analog[channel][i])
                if addr in already_added:
                    added = already_added[addr]
                    bytes[int(channel[-1])].append(added)
                    
                else:
                    waveform_new = array_analog[channel][i] + array_M1[channel][i]*(2**14) +\
                    array_M2[channel][i]*(2**15)
                    aux = np.empty(2*len(waveform_new), dtype=np.uint8)
                    aux[::2] = waveform_new % 2**8
                    aux[1::2] = waveform_new // 2**8
                    byteadded = bytearray(aux)
                    bytes[int(channel[-1])].append(byteadded)
                    already_added[addr] = byteadded

        return True, bytes, repeats

    def _get_sampling_time(self):
        """ Getter for the sampling time prop of BaseContext.

        """
        return 1/self.sampling_frequency*TIME_CONVERSION['s'][self.time_unit]

    @observe('sampling_frequency', 'time_unit')
    def _reset_sampling_time(self, change):
        """ Observer resetting the sampling_time property.

        """
        member = self.get_member('sampling_time')
        member.reset(self)

CONTEXTS = {'AWG': AWGContext}
