from __future__ import print_function
import os
import datetime
import subprocess
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec

from uldaq import (get_daq_device_inventory, DaqDevice, AInScanFlag, ScanStatus,
                   ScanOption, create_float_buffer, InterfaceType, AiInputMode)


class plot():
    def __init__(self):
        self.n_rows = None
        self.n_cols = None
        self.max_v = None
        self.channel_handle = []

        self.fig = plt.figure(figsize=(20 / 2.54, 12 / 2.54), facecolor='white')
        self.axs = []
        plt.show(block=False)

    def create_axis(self):
        gs = gridspec.GridSpec(self.n_rows, self.n_cols)
        for x in range(self.n_cols):
            for y in range(self.n_rows):
            # for x in range(self.n_cols):
                ax = plt.subplot(gs[y, x])

                if not y == self.n_rows - 1:
                    ax.tick_params(axis='x', which='both', bottom=False, top=False, labelbottom=False)

                if not x == 0:
                    ax.tick_params(axis='y', which='both', left=False, right=False, labelleft=False)

                ax.set_ylim(-self.max_v, self.max_v)

                self.axs.append(ax)
        gs.update(left=0.1, bottom=0.05, top=1, right=1, hspace=0, wspace=0)

    def creat_axis_competition(self):
        gs = gridspec.GridSpec(3, 8, left =0.1, bottom = 0.05, top = 1, right=1)
        pos = [[0, 1],
               [0, 2],
               [0, 3],
               [0, 4],
               [1, 0],
               [1, 1],
               [1, 2],
               [1, 3],
               [1, 4],
               [1, 5],
               [2, 1],
               [2, 2],
               [2, 3],
               [2, 4],
               [1, 7]]

        show_x = np.array([0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1], dtype=bool)
        show_y = np.array([1, 0, 0, 0, 1, 0, 0, 0, 0, 0, 1, 0, 0, 0, 1], dtype=bool)
        for enu, p in enumerate(pos):
            ax = self.fig.add_subplot(gs[p[0], p[1]]) # elec 1
            ax.set_ylim(-self.max_v, self.max_v)
            if show_x[enu] == False:
                ax.tick_params(axis='x', which='both', bottom=False, top=False, labelbottom=False)
            if show_y[enu] == False:
                ax.tick_params(axis='y', which='both', left=False, right=False, labelleft=False)

            self.axs.append(ax)


def read_cfg(cfg_file, now, init_read=False):
    cfg_f = open(cfg_file, 'r+')
    cfg = cfg_f.readlines()

    ### read cfg information ###
    if init_read:
        for line in cfg:
            if "PathFormat" in line:
                path_format = ':'.join(line.split(':')[1:]).strip().replace('"', '').replace("'", "")
                cfg_f.close()
                return path_format

    for line in cfg:
        if 'Columns1' in line:
            n_cols = int(line.split(':')[1].strip())
        elif 'Rows1' in line:
            n_rows = int(line.split(':')[1].strip())
        elif "AISampleRate" in line:
            samplerate = int(float(line.split(':')[-1].strip().replace('kHz', '')) * 1000)
        elif "AIMaxVolt" in line:
            max_v = float(line.split(':')[1].strip().replace('mV', ''))
        elif 'Gain' in line:
            gain = int(line.split(':')[1].strip())
    channels = n_rows * n_cols

    ### alter information and re-write ###
    for enu, line in enumerate(cfg):
        if "StartDate" in line:
            cfg[enu] = ('          StartDate        : %s\n' % now.strftime('%Y-%m-%d'))
        elif "StartTime" in line:
            cfg[enu] = ('          StartTime        : %s\n' % (now.strftime('%H:%M:%S') + now.strftime(".%f")[:4]))
    cfg_f.close()

    cfg_f = open(cfg_file, 'w+')
    for line in cfg:
        cfg_f.write(line)
    cfg_f.close()

    return channels, samplerate, n_cols, n_rows, max_v, gain


def main():
    now = datetime.datetime.now()

    # get init cfg
    if os.path.exists('/media/pi/data1'):
        init_path = '/media/pi/data1'
    else:
        init_path = '/home/raab/data/rasp_test'

    init_cfgfile = os.path.join(init_path, 'fishgrid.cfg')
    if os.path.exists(init_cfgfile):
        path_format = read_cfg(init_cfgfile, now, init_read=True)
    else:
        print('cfg file missing !!!')
        quit()

    # read and edit config file
    channels, rate, n_cols, n_rows, max_v, gain = read_cfg(init_cfgfile, now)

    if True:

        status = ScanStatus.IDLE

        descriptor_index = 0  # ToDo: ????
        range_index = 0  # ToDo: ????

        interface_type = InterfaceType.USB
        low_channel = 0
        high_channel = channels - 1

        samples_per_channel = rate * 2  # * channels = Buffer size
        # rate = 20000
        scan_options = ScanOption.CONTINUOUS
        flags = AInScanFlag.DEFAULT

        # Get descriptors for all of the available DAQ devices.
        devices = get_daq_device_inventory(interface_type)
        number_of_devices = len(devices)
        if number_of_devices == 0:
            raise Exception('Error: No DAQ devices found')

        print('Found', number_of_devices, 'DAQ device(s):')
        for i in range(number_of_devices):
            print('  ', devices[i].product_name, ' (', devices[i].unique_id, ')', sep='')

        # Create the DAQ device object associated with the specified descriptor index.
        daq_device = None
        daq_device = DaqDevice(devices[descriptor_index])

        # Get the AiDevice object and verify that it is valid.
        ai_device = None
        ai_device = daq_device.get_ai_device()
        if ai_device is None:
            raise Exception('Error: The DAQ device does not support analog input')

        # Verify that the specified device supports hardware pacing for analog input.
        ai_info = ai_device.get_info()
        if not ai_info.has_pacer():
            raise Exception('\nError: The specified DAQ device does not support hardware paced analog input')

        # Establish a connection to the DAQ device.
        descriptor = daq_device.get_descriptor()
        print('\nConnecting to', descriptor.dev_string, '- please wait...')
        daq_device.connect()

        # The default input mode is SINGLE_ENDED.
        input_mode = AiInputMode.SINGLE_ENDED
        # If SINGLE_ENDED input mode is not supported, set to DIFFERENTIAL.
        if ai_info.get_num_chans_by_mode(AiInputMode.SINGLE_ENDED) <= 0:
            input_mode = AiInputMode.DIFFERENTIAL

        # Get the number of channels and validate the high channel number.
        number_of_channels = ai_info.get_num_chans_by_mode(input_mode)
        if high_channel >= number_of_channels:
            high_channel = number_of_channels - 1
        channel_count = high_channel - low_channel + 1

        # Get a list of supported ranges and validate the range index.
        ranges = ai_info.get_ranges(input_mode)
        int_ranges = []
        for r in ranges:
            int_ranges.append(int(r.name.replace('BIP', '').replace('VOLTS', '')))

        for idx in np.argsort(int_ranges):
            if max_v * gain / 1000 <= int_ranges[idx]:
                range_index = idx
                break
        print(ranges[range_index])

        data = create_float_buffer(channel_count, samples_per_channel)

        rate = ai_device.a_in_scan(low_channel, high_channel, input_mode, ranges[range_index], samples_per_channel,
                                   rate, scan_options, flags, data)
        last_idx = 0

    Plot = plot()
    Plot.max_v = max_v
    Plot.n_rows = n_rows
    Plot.n_cols = n_cols

    #Plot.create_axis()
    Plot.creat_axis_competition()

    init_fig = True
    try:
        while True:
            status, transfer_status = ai_device.get_scan_status()

            index = transfer_status.current_index

            if (last_idx > index) and (index != -1):
                channel_array = np.arange(channels)
                channel_data = list(map(lambda x : data[x::channels][:250], channel_array))
                channel_std = list(map(lambda x : np.std(data[x::channels][:250]), channel_array))
                power_channel = int(np.argmax(channel_std))

                if init_fig == True:
                    yspan = (np.min(channel_data[power_channel]) / gain, np.max(channel_data[power_channel]) / gain)
                    ylim = (yspan[0] - np.abs(np.diff(yspan)) * 0.2, yspan[1] + np.abs(np.diff(yspan))* 0.2)

                    for ch in channel_array:
                        h, = Plot.axs[ch].plot(np.arange(250)[:len(channel_data[ch])] / rate, np.array(channel_data[ch]) / gain, color='k')
                        Plot.axs[ch].set_ylim(ylim)
                        Plot.channel_handle.append(h)

                    Plot.fig.canvas.draw()

                    init_fig = False
                else:
                    yspan = [np.min(channel_data[power_channel]) / gain, np.max(channel_data[power_channel]) / gain]
                    ylim = [yspan[0] - np.abs(np.diff(yspan)) * 0.2, yspan[1] + np.abs(np.diff(yspan)) * 0.2]
                    for ch in channel_array:
                        Plot.channel_handle[ch].set_data(np.arange(250)[:len(channel_data[ch])] / rate, np.array(channel_data[ch]) / gain)
                        Plot.axs[ch].set_ylim(ylim)
                    Plot.fig.canvas.draw()

            if index == -1:
                last_idx = len(data)
            else:
                last_idx = index

    except KeyboardInterrupt:
        plt.close()
        pass

    # f.close()
    if daq_device:
        # Stop the acquisition if it is still running.
        if status == ScanStatus.RUNNING:
            ai_device.scan_stop()
        if daq_device.is_connected():
            daq_device.disconnect()
        daq_device.release()


if __name__ == '__main__':
    main()