#!python
# -*- 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 collections
import argparse
import time
import datetime
import textwrap
import itertools
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'

# Default phase rate output
PHASE_RATE = 1000

# List of allowable phase rate output value
ORATE_LIST = [1, 10, 100, 1000]

# 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)
DQ_LIST.append('getall')


# =============================================================================
def change_output_phase_rate(dev, rate):
    """Change output phase rate over data port of test set device.
    The phase rate cannot be changed if device is not the 5125A model.
    :param dev: instance of device object (object)
    :param filename: filename of file to write data (str)
    :returns: None
    """
    dev.connect()
    dev.write("show version")
    retval = dev.read_until("Noise Spectrum", timeout=1)
    if not "5125A" in retval:
        print("Cannot change output phase rate, not a 5215A device")
        dev.disconnect()
        return
    dev.write("show state")
    if "Collecting" in dev.read_very_eager():
        print("Cannot change output phase rate while collecting data")
        dev.disconnect()
        return
    dev.write("set phaserate {}".format(rate))
    dev.disconnect()


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
    """
    dev.connect_data_port()
    while True:
        try:
            data = wait_for_data(dev)
            with open(filename, 'a') as fd:
                fd.write(data)
        except KeyboardInterrupt:
            dev.disconnect()
            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 = (itertools.zip_longest(adev500, noise500),
                      itertools.zip_longest(adev50, noise50),
                      itertools.zip_longest(adev5, noise5),
                      itertools.zip_longest(adev05, noise05))
        data_out = collections.OrderedDict({'adev_1ms': '', 'adev_10ms': '',
                                            'adev_100ms': '', 'adev_1s': ''})
        for dok, an in zip(data_out.keys(), adev_noise):
            for a, n in an:
                if n is not None:
                    [tau, adev, tau_, noise] = RX.findall(a + n)
                    data_out[dok] += '{}\t{}\t{}\n'.format(tau, adev, noise)
                else:
                    [tau, adev] = RX.findall(a)
                    data_out[dok] += '{}\t{}\n'.format(tau, adev)
        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]
        data_out = {'spectrum': '', 'spectrum_noise': ''}
        for s in spectrum:
            [freq, psd] = RX.findall(s)
            data_out['spectrum'] += '{}\t{}\n'.format(freq, psd)
        for n in noise:
            [freq, floor] = RX.findall(n)
            data_out['spectrum_noise'] += '{}\t{}\n'.format(freq, floor)
        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 of phase 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('-r', '--output-rate',
                           choices=ORATE_LIST,
                           action='store',
                           dest='orate',
                           type=int,
                           default=PHASE_RATE,
                           help='Output phase rate of acquisition ' \
                           '!!Only effective with 5125 test set!! ' \
                           '(default: {})'.format(PHASE_RATE))
    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':
        dirname = args.odir
        if not os.path.exists(dirname):
            os.makedirs(dirname)
        orate = args.orate
        date = datetime.datetime.utcnow().strftime("%Y%m%d-%H%M%S")
        full_filename = dirname + '/microsemi-ca-rate-' \
            + str(orate) + 'sps-' \
            + date + '.dat'
        change_output_phase_rate(dev, orate)
        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()
