#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""package micro5125a
author    Benoit Dubois
copyright FEMTO-Engineering, 2019
license   GPL v3.0+
brief     Handle Microsemi (former Symmetricom, former Timing Solutions)
          5125a/5120a/5115a/5110a Phase Noise Test Set device.
"""

import os
import re
import logging
import argparse
import time
import datetime
import textwrap

import micro5125a.micro5125a as micro5125a
import micro5125a.version as version


# Current logging level
LOG_LEVEL = logging.INFO

# Default output directory basename to write data
DIR_BASENAME = '5125a_experiments'

# Method to extract floats in a string.
# See https://stackoverflow.com/questions/4703390/how-to-extract-a-floating-number-from-a-string
# Return a list of floats contained in the string.
NUMERIC_CONST_PATTERN = r"""
     [-+]? # optional sign
     (?:
         (?: \d* \. \d+ ) # .1 .12 .123 etc 9.1 etc 98.1 etc
         |
         (?: \d+ \.? ) # 1. 12. 123. etc 1 12 123 etc
     )
     # followed by optional exponent part if desired
     (?: [Ee] [+-]? \d+ ) ?
     """
RX = re.compile(NUMERIC_CONST_PATTERN, re.VERBOSE)

DQ_LIST = list(micro5125a.DATA_LIST).append('getall')


# =============================================================================
def cont_acq(dev, filename):
    """Do continous acquisition from data output port of 'dev' and save data
    in 'filename'.
    :param dev: instance of device object (object)
    ;param filename: filename of file to write data (str)
    :returns: None
    """
    while True:
        try:
            data = wait_for_data(dev)
            with open(filename, 'a') as fd:
                fd.write(data)
        except KeyboardInterrupt:
            return

def wait_for_data(dev):
    """Polling data on data output port.
    :param dev: instance of device object (object)
    :returns: data read on output port (array)
    """
    while True:
        try:
            data = dev.read_very_eager()
            return data
        except KeyboardInterrupt:
            raise KeyboardInterrupt
        except Exception as ex:
            logging.warning("Exception: %r", ex)
            logging.warning("Possible missing data samples")
            time.sleep(0.4)


# =============================================================================
def parse_data(cmd, data):
    """Format raw data returned by 5125a device.
    :param cmd: command query used to download data (str)
    :param data: raw data read from device (str)
    :returns: dictionary with type of data as keys and data as values (dict)
    """
    if cmd == 'adev':
        data500, data = data.split('TAU0: 1E-2 (NEQ BW: 50 Hz)')
        data50, data = data.split('TAU0: 1E-1 (NEQ BW: 5 Hz)')
        data5, data05 = data.split('TAU0: 1E0 (NEQ BW: 0.5 Hz)')
        adev500, noise500 = data500.split('Noisefloor')
        adev50, noise50 = data50.split('Noisefloor')
        adev5, noise5 = data5.split('Noisefloor')
        adev05, noise05 = data05.split('Noisefloor')
        adev500 = adev500.split('\r\n')[1:-1]
        adev50 = adev50.split('\r\n')[1:-1]
        adev5 = adev5.split('\r\n')[1:-1]
        adev05 = adev05.split('\r\n')[1:-1]
        noise500 = noise500.split('\r\n')[1:-2]
        noise50 = noise50.split('\r\n')[1:-2]
        noise5 = noise5.split('\r\n')[1:-2]
        noise05 = noise05.split('\r\n')[1:-3]
        adev_noise = (zip(adev500, noise500),
                      zip(adev50, noise50),
                      zip(adev5, noise5),
                      zip(adev05, noise05))
        data_out = {'adev_1ms': '', 'adev_10ms': '',
                    'adev_100ms': '', 'adev_1s': ''}
        for dok, an in zip(data_out.keys(), adev_noise):
            for a, n in an:
                [tau, adev, tau_, noise] = RX.findall(a + n)
                data_out[dok] += '{}\t{}\t{}\n'.format(tau, adev, noise)
        return data_out
    if cmd == 'spectrum':
        spectrum, noise = data.split('Noise Floor')
        spectrum = spectrum.split('\r\n')[3:-2]
        noise = noise.split('\r\n')[2:-2]
        spectrum_out = ''
        noise_out = ''
        for s in spectrum:
            [freq, psd] = RX.findall(s)
            spectrum_out += '{}\t{}\n'.format(freq, psd)
        for n in noise:
            [freq, psd] = RX.findall(n)
            noise_out += '{}\t{}\n'.format(freq, psd)
        data_out = {'spectrum': spectrum_out, 'spectrum_noise': noise_out}
        return data_out
    data = data[1:].replace('\r', '')
    return {cmd: data}


# =============================================================================
def parse_args():
    """Parse scriptargument.
    """
    parser = argparse.ArgumentParser(
        formatter_class=argparse.RawDescriptionHelpFormatter,
        description='''
   Acquire data from Microsemi (former Symmetricom) 5125A Phase Noise
   Test Set device.''',
        epilog=textwrap.dedent('''
   For more details:
      $ micro-5125a {rt|dq|ca} -h
'''))
    cmdsubparser = parser.add_subparsers(title='subcommands',
                                         dest='cmd',
                                         description='valid sub commands',
                                         help='Query sub command')
    #
    actparser = cmdsubparser.add_parser(
        'rt',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        help='Real time action from {}'.format(micro5125a.ACTION_LIST),
        epilog=textwrap.dedent('''\
   Example:
      $ micro-5125a rt pause 192.168.0.2
   Connect to 5125a device @ip 192.168.0.2 and pause acquisition
'''))
    actparser.add_argument('action',
                           choices=micro5125a.ACTION_LIST,
                           help='Action to be done')
    #
    datparser = cmdsubparser.add_parser(
        'dq',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        help='Data query from {}'.format(DQ_LIST),
        epilog=textwrap.dedent('''
   Example 1:
      $ micro-5125a dq getall -o ~/myexperiment 192.168.0.2
   Connect to 5125a device @ip 192.168.0.2 and save a snapshot of all data
   to directory '~/myexperiment/'

   Example 2:
      $ micro-5125a dq spectrum -dd 192.168.0.2
   Connect to 5125a device @ip 192.168.0.2 and save spectrum data
   to directory './{}/date_of_snapshot/'
'''.format(DIR_BASENAME)))
    datparser.add_argument('data_type',
                           choices=DQ_LIST,
                           help='Data to save')
    datparser.add_argument('-o', '--output-directory',
                           action='store', dest='odir',
                           default=DIR_BASENAME,
                           help='Output data directory (default: {})' \
                           .format(DIR_BASENAME))
    datparser.add_argument('-dd', '--dated-dir', action='store_true',
                           dest='dated_dir',
                           help='Add subdirectory with the date as name ' \
                           '(can help to sort data)')
    #
    acqparser = cmdsubparser.add_parser(
        'ca',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        help='Continous Acquisition from device.',
        epilog=textwrap.dedent('''\
   Example:
      $ micro-5125a ca 100 -o ~/myexperiment 192.168.0.2
   Connect to 5125a device @ip 192.168.0.2, start acquisition and save
   data to directory '~/myexperiment/5125a-YYYYMMDD-hhmmss.dat'.

   To stop acquisition, simply stop application.
'''))
    acqparser.add_argument('-o', '--output-directory',
                           action='store', dest='odir',
                           default=DIR_BASENAME,
                           help='Output data directory ' \
                           '(default: {})'.format(DIR_BASENAME))
    #
    parser.add_argument('ip', help='IP address of device')
    parser.add_argument('--version', action='version',
                        version=version.__version__)
    #
    return parser.parse_args()


# =============================================================================
def main():
    """Main part of script
    """
    args = parse_args()
    dev = micro5125a.Micro5125A(args.ip)
    cmd = args.cmd
    if cmd == 'rt':
        action = args.action
        if action not in micro5125a.ACTION_LIST:
            raise ValueError("Bad 'action' parameter: {}".format(action))
        dev.connect()
        dev.write(action)
        return
    if cmd == 'ca':
        date = datetime.datetime.utcnow().strftime("%Y%m%d-%H%M%S")
        full_filename = args.odir + '/5125a-' + date + '.dat'
        dev.connect_data_port()
        cont_acq(dev, full_filename)
        return
    if cmd == 'dq':
        data_type = args.data_type
        if data_type == 'getall':
            data_types = micro5125a.DATA_LIST
        elif data_type in micro5125a.DATA_LIST:
            data_types = (data_type, )
        else:
            raise ValueError("Bad 'data_type' parameter: {}".format(action))
        dirname = args.odir
        if args.dated_dir is True:
            date = datetime.datetime.utcnow().strftime("%Y%m%d-%H%M%S")
            dirname += '/' + date
        if not os.path.exists(dirname):
            os.makedirs(dirname)
        dev.connect()
        for dt in data_types:
            raw_data = dev.show_data(dt)
            data = parse_data(dt, raw_data)
            for key in data:
                with open(dirname + '/' + key + ".txt", 'w') as fd:
                    fd.write(data[key])
        return
    raise ValueError("Bad 'cmd' parameter: {}".format(cmd))


# =============================================================================
if __name__ == "__main__":
    LOG_FORMAT = '%(levelname) -8s %(filename)s (%(lineno)d): %(message)s'
    logging.basicConfig(format=LOG_FORMAT, level=LOG_LEVEL)
    #
    main()
