# -*- coding: utf-8 -*-
"""

Data readers & processing utils, Nexus specific

@author: Jussi (jnu@iki.fi)
"""

from __future__ import print_function, division
from collections import defaultdict
import sys
import numpy as np
import os.path as op
import psutil
import glob
import logging
import time
import multiprocessing

from .numutils import change_coords
from . import GaitDataError, cfg


logger = logging.getLogger(__name__)


def _find_nexus_path(vicon_path=None):
    """Return path to most recent Nexus version"""

    if vicon_path is None:
        vicon_path = r'C:\Program Files (x86)\Vicon'  # educated guess

    if not op.isdir(vicon_path):
        return None

    nexus_glob = op.join(vicon_path, 'Nexus?.*')
    nexus_dirs = glob.glob(nexus_glob)
    if not nexus_dirs:
        return None
    nexus_vers = [op.split(dir_)[1][5:] for dir_ in nexus_dirs]
    idx = nexus_vers.index(max(nexus_vers))
    return nexus_dirs[idx]


def _add_nexus_path(vicon_path):
    """Add Nexus SDK dir to sys.path"""

    nexus_path = _find_nexus_path(vicon_path)
    if nexus_path is None:
        logger.debug('cannot locate Nexus installation directory under %s'
              % vicon_path)
        return

    sdk_path = op.join(nexus_path, 'SDK', 'Python')
    if sdk_path not in sys.path:
        sys.path.append(sdk_path)
    else:
        logger.debug('%s already in sys.path' % sdk_path)

    # import from Win32 or Win64 according to bitness of Python interpreter
    bitness = '64' if sys.maxsize > 2**32 else '32'
    win = 'Win' + bitness
    _win_sdk_path = op.join(nexus_path, 'SDK', win)

    # check that the path for the wrong architecture has not already been
    # added to path (this may happen when running inside Nexus)
    win_other = 'Win32' if win == 'Win64' else 'Win64'
    _win_sdk_other = op.join(nexus_path, 'SDK', win_other)
    if _win_sdk_other in sys.path:
        logger.debug('%s already in sys.path, removing' % _win_sdk_other)
        sys.path.remove(_win_sdk_other)

    if _win_sdk_path not in sys.path:
        logger.debug('using Nexus SDK from %s' % _win_sdk_path)
        sys.path.append(_win_sdk_path)
    else:
        logger.debug('%s already in sys.path' % _win_sdk_path)


if sys.version_info.major >= 3:
    logger.debug('running on Python 3 or newer, cannot import Nexus API (yet)')
else:
    vicon_path = op.normpath(cfg.general.vicon_path)
    _add_nexus_path(vicon_path)
    try:
        import ViconNexus
    except ImportError:
        logger.debug('cannot import Vicon Nexus SDK')

sys.stdout.flush()  # make sure import warnings get printed


def pid():
    """ Tries to return the PID of the running Nexus process. """
    PROCNAME = "Nexus.exe"
    for proc in psutil.process_iter():
        try:
            if proc.name() == PROCNAME:
                return proc.pid
        # catch NoSuchProcess for procs that disappear inside loop
        except (psutil.AccessDenied, psutil.NoSuchProcess):
            pass
    return None


def true_ver():
    """ Tries to return the actual version of the running Nexus process
    (API does not do that). Hackish and probably unreliable """
    PROCNAME = "Nexus.exe"
    for proc in psutil.process_iter():
        try:
            if proc.name() == PROCNAME:
                exname = proc.exe()
                vind = exname.find('2.')  # assumes ver >2.
                if vind == -1:
                    return None
                try:
                    return float(exname[vind:vind+3])
                except ValueError:
                    return None
        except psutil.AccessDenied:
            pass
    return None


def viconnexus():
    """ Return a ViconNexus instance. """
    check_nexus()
    return ViconNexus.ViconNexus()


def get_subjectnames(single_only=True):
    """ Get subject name(s) from Nexus """
    vicon = viconnexus()
    get_sessionpath()  # check whether we can get data
    names_ = vicon.GetSubjectNames()
    if not names_:
        raise GaitDataError('No subject defined in Nexus')
    if single_only:
        if len(names_) > 1:
            raise GaitDataError('Nexus returns multiple subjects')
    """ Workaround a Nexus 2.6 bug (?) that creates extra names with
    weird unicode strings """
    names_ = [name for name in names_ if u'\ufffd1' not in name]
    return names_[0] if single_only else names_


def check_nexus():
    if not pid():
        raise GaitDataError('Vicon Nexus does not seem to be running')


def get_sessionpath():
    """ Get path to current session """
    try:
        vicon = viconnexus()
        sessionpath = vicon.GetTrialName()[0]
    except IOError:  # may be raised if Nexus was just terminated
        sessionpath = None
    if not sessionpath:
        raise GaitDataError('Cannot get Nexus session path, '
                            'no session or maybe in Live mode?')
    return op.normpath(sessionpath)


def _run_pipeline(pipeline, foo, timeout):
    """Wrapper needed for multiprocessing due to pickle limitations"""
    vicon = viconnexus()
    return vicon.Client.RunPipeline(pipeline, foo, timeout)


def run_pipelines(pipelines):
    """Run given Nexus pipeline(s)"""
    if type(pipelines) != list:
        pipelines = [pipelines]
    for pipeline in pipelines:
        logger.debug('running pipeline: %s' % pipeline)
        result = _run_pipeline(pipeline.encode('utf-8'), '',
                               cfg.autoproc.nexus_timeout)
        if result.Error():
            logger.warning('error while trying to run Nexus pipeline: %s'
                           % pipeline)


def run_pipelines_multiprocessing(pipelines):
    """Run given Nexus pipeline(s) via multiprocessing module. The idea
    is to work around the GIL, since the Nexus API does not release it.
    This version causes the invoking thread to sleep and thus release the GIL
    while the pipeline is running"""
    if type(pipelines) != list:
        pipelines = [pipelines]
    for pipeline in pipelines:
        logger.debug('running pipeline: %s' % pipeline)
        args = (pipeline.encode('utf-8'), '', cfg.autoproc.nexus_timeout)
        p = multiprocessing.Process(target=_run_pipeline, args=args)
        p.start()
        while p.exitcode is None:
            time.sleep(.1)


def get_trialname():
    """ Get trial name without session path """
    vicon = viconnexus()
    trialname_ = vicon.GetTrialName()
    return trialname_[1]


def is_vicon_instance(obj):
    """ Check if obj is an instance of ViconNexus """
    return obj.__class__.__name__ == 'ViconNexus'


def _get_nexus_subject_param(vicon, name, param):
    """Get subject parameter from Nexus."""
    value = vicon.GetSubjectParam(name, param)
    # for unknown reasons, above method may return tuple or float
    # depending on whether script is run from Nexus or outside
    if type(value) == tuple:
        value = value[0] if value[1] else None
    return value


def _get_marker_names(vicon, trajs_only=True):
    """Return marker names (only ones with trajectories, if trajs_only)"""
    subjname = get_subjectnames()
    markers = vicon.GetMarkerNames(subjname)
    # only get markers with trajectories - excludes calibration markers
    if trajs_only:
        markers = [mkr for mkr in markers if vicon.HasTrajectory(subjname,
                                                                 mkr)]
    return markers


# FIXME: most of get_ methods below are intended to be called via read_data
# so underscore them
def get_metadata(vicon):
    """ Read trial and subject metadata """
    check_nexus()
    logger.debug('reading metadata from Vicon Nexus')
    subjname = get_subjectnames()
    params_available = vicon.GetSubjectParamNames(subjname)
    subj_params = defaultdict(lambda: None)
    subj_params.update({par: _get_nexus_subject_param(vicon, subjname, par) for
                        par in params_available})
    trialname = get_trialname()
    if not trialname:
        raise GaitDataError('No trial loaded in Nexus')
    sessionpath = get_sessionpath()
    markers = _get_marker_names(vicon)
    # get events - GetEvents() indices seem to often be 1 frame less than on
    # Nexus display - only happens with ROI?
    lstrikes = vicon.GetEvents(subjname, "Left", "Foot Strike")[0]
    rstrikes = vicon.GetEvents(subjname, "Right", "Foot Strike")[0]
    ltoeoffs = vicon.GetEvents(subjname, "Left", "Foot Off")[0]
    rtoeoffs = vicon.GetEvents(subjname, "Right", "Foot Off")[0]
    # Offset will be subtracted from event frame numbers to get correct
    # 0-based index for frame data. For Nexus, it is always 1 (Nexus uses
    # 1-based frame numbering)
    offset = 1
    length = vicon.GetFrameCount()
    framerate = vicon.GetFrameRate()
    # Get analog rate. This may not be mandatory if analog devices
    # are not used, but currently it needs to succeed.
    devids = vicon.GetDeviceIDs()
    if not devids:
        raise GaitDataError('Cannot determine analog rate')
    else:
        devid = devids[0]
        _, _, analograte, _, _, _ = vicon.GetDeviceDetails(devid)
    samplesperframe = analograte / framerate
    logger.debug('offset @ %d, %d frames, framerate %d Hz, %d samples per '
                 'frame' % (offset, length, framerate, samplesperframe))
    # get n of forceplates
    fp_devids = [id_ for id_ in devids if
                 vicon.GetDeviceDetails(id_)[1].lower() == 'forceplate']

    # sort events (may be in wrong temporal order, at least in c3d files)
    for li in [lstrikes, rstrikes, ltoeoffs, rtoeoffs]:
        li.sort()

    return {'trialname': trialname, 'sessionpath': sessionpath,
            'offset': offset, 'framerate': framerate, 'analograte': analograte,
            'name': subjname, 'subj_params': subj_params, 'lstrikes': lstrikes,
            'rstrikes': rstrikes, 'ltoeoffs': ltoeoffs, 'rtoeoffs': rtoeoffs,
            'length': length, 'samplesperframe': samplesperframe,
            'n_forceplates': len(fp_devids), 'markers': markers}


def get_emg_data(vicon):
    """ Read EMG data from Nexus. This uses the configured EMG device name. """
    return _get_analog_data(vicon, cfg.emg.devname)


def get_accelerometer_data(vicon):
    """ Read EMG data from Nexus. This uses the configured EMG device name. """
    return _get_analog_data(vicon, cfg.analog.accelerometer_devname)


def _get_analog_data(vicon, devname):
    """ Read analog data from Nexus """
    ids = [id_ for id_ in vicon.GetDeviceIDs() if
           vicon.GetDeviceDetails(id_)[0].lower() == devname.lower()]
    if len(ids) > 1:
        raise GaitDataError('Multiple matching analog devices')
    elif len(ids) == 0:
        raise GaitDataError('No matching analog devices')
    dev_id = ids[0]
    dname, dtype, drate, outputids, _, _ = vicon.GetDeviceDetails(dev_id)
    # not handling multiple output ids yet
    if len(outputids) != 1:
        raise GaitDataError('Expected single output for device')
    outputid = outputids[0]
    # get list of channel names and IDs
    _, _, _, _, chnames, chids = vicon.GetDeviceOutputDetails(dev_id, outputid)
    data = dict()
    for chid in chids:
        chdata, _, chrate = vicon.GetDeviceChannel(dev_id, outputid, chid)
        chname = chnames[chid-1]  # chids start from 1
        data[chname] = np.array(chdata)
    return {'t': np.arange(len(chdata)) / drate, 'data': data}


def _get_1_forceplate_data(vicon, devid):
    """ Read data of single forceplate from Nexus.
    Data is returned in global coordinate frame """
    # get forceplate ids
    logger.debug('reading forceplate data from devid %d' % devid)
    dname, dtype, drate, outputids, nfp, _ = vicon.GetDeviceDetails(devid)
    # outputs should be force, moment, cop. select force
    outputid = outputids[0]
    chid = vicon.GetDeviceChannelIDFromName(devid, outputid, 'Fx')
    fx, chready, chrate = vicon.GetDeviceChannelGlobal(devid, outputid, chid)
    chid = vicon.GetDeviceChannelIDFromName(devid, outputid, 'Fy')
    fy, chready, chrate = vicon.GetDeviceChannelGlobal(devid, outputid, chid)
    chid = vicon.GetDeviceChannelIDFromName(devid, outputid, 'Fz')
    fz, chready, chrate = vicon.GetDeviceChannelGlobal(devid, outputid, chid)
    # moments
    outputid = outputids[1]
    chid = vicon.GetDeviceChannelIDFromName(devid, outputid, 'Mx')
    mx, chready, chrate = vicon.GetDeviceChannelGlobal(devid, outputid, chid)
    chid = vicon.GetDeviceChannelIDFromName(devid, outputid, 'My')
    my, chready, chrate = vicon.GetDeviceChannelGlobal(devid, outputid, chid)
    chid = vicon.GetDeviceChannelIDFromName(devid, outputid, 'Mz')
    mz, chready, chrate = vicon.GetDeviceChannelGlobal(devid, outputid, chid)
    # center of pressure
    outputid = outputids[2]
    chid = vicon.GetDeviceChannelIDFromName(devid, outputid, 'Cx')
    copx, chready, chrate = vicon.GetDeviceChannelGlobal(devid, outputid, chid)
    chid = vicon.GetDeviceChannelIDFromName(devid, outputid, 'Cy')
    copy, chready, chrate = vicon.GetDeviceChannelGlobal(devid, outputid, chid)
    chid = vicon.GetDeviceChannelIDFromName(devid, outputid, 'Cz')
    copz, chready, chrate = vicon.GetDeviceChannelGlobal(devid, outputid, chid)
    cop_w = np.array([copx, copy, copz]).transpose()
    F = np.array([fx, fy, fz]).transpose()
    M = np.array([mx, my, mz]).transpose()
    Ftot = np.linalg.norm(F, axis=1)
    # translation and rotation matrices -> world coords
    # suspect that Nexus wR is wrong (does not match displayed plate axes)?
    wR = np.array(nfp.WorldR).reshape(3, 3)
    wT = np.array(nfp.WorldT)
    # plate corners -> world coords
    cor = np.stack([nfp.LowerBounds, nfp.UpperBounds])
    cor_w = change_coords(cor, wR, wT)
    cor_full = np.array([cor_w[0, :],
                        [cor_w[0, 0], cor_w[1, 1], cor_w[0, 2]],
                        cor_w[1, :],
                        [cor_w[1, 0], cor_w[0, 1], cor_w[0, 2]]])

    lb = np.min(cor_w, axis=0)
    ub = np.max(cor_w, axis=0)
    # check that CoP stays inside plate boundaries
    cop_ok = np.logical_and(cop_w[:, 0] >= lb[0], cop_w[:, 0] <= ub[0]).all()
    cop_ok &= np.logical_and(cop_w[:, 1] >= lb[1], cop_w[:, 1] <= ub[1]).all()
    if not cop_ok:
        logger.warning('center of pressure outside plate boundaries, '
                       'clipping to plate')
        cop_w[:, 0] = np.clip(cop_w[:, 0], lb[0], ub[0])
        cop_w[:, 1] = np.clip(cop_w[:, 1], lb[1], ub[1])
    return {'F': F, 'M': M, 'Ftot': Ftot, 'CoP': cop_w,
            'wR': wR, 'wT': wT, 'lowerbounds': lb, 'upperbounds': ub,
            'cor_w': cor_w, 'cor_full': cor_full}


def get_forceplate_data(vicon):
    """ Read all forceplate data from Nexus. """
    # get forceplate ids
    logger.debug('reading forceplate data from Vicon Nexus')
    devids = [id_ for id_ in vicon.GetDeviceIDs() if
              vicon.GetDeviceDetails(id_)[1].lower() == 'forceplate']
    if len(devids) == 0:
        logger.debug('no forceplates detected')
        return None
    logger.debug('detected %d forceplate(s)' % len(devids))
    return [_get_1_forceplate_data(vicon, devid) for devid in devids]


def _swap_markers(vicon, marker1, marker2):
    """Swap two marker trajectories in currently loaded trial."""
    subj = get_subjectnames()
    m1 = vicon.GetTrajectory(subj, marker1)
    m2 = vicon.GetTrajectory(subj, marker2)
    vicon.SetTrajectory(subj, marker2, m1[0], m1[1], m1[2], m1[3])
    vicon.SetTrajectory(subj, marker1, m2[0], m2[1], m2[2], m2[3])


def _get_marker_data(vicon, markers, ignore_edge_gaps=True,
                     ignore_missing=False):
    """Get position, velocity and acceleration for specified markers."""
    if not isinstance(markers, list):
        markers = [markers]
    subj = get_subjectnames()
    mkrdata = dict()
    for marker in markers:
        x, y, z, _ = vicon.GetTrajectory(subj, marker)
        if len(x) == 0:
            if ignore_missing:
                logger.warning('Cannot read trajectory %s from Nexus'
                               % marker)
                continue
            else:
                raise GaitDataError('Cannot read marker trajectory '
                                    'from Nexus: \'%s\'' % marker)
        mP = np.array([x, y, z]).transpose()
        mkrdata[marker] = mP
        mkrdata[marker + '_P'] = mP
        mkrdata[marker + '_V'] = np.gradient(mP)[0]
        mkrdata[marker + '_A'] = np.gradient(mkrdata[marker+'_V'])[0]
        # find gaps
        allzero = np.any(mP, axis=1).astype(int)
        if ignore_edge_gaps:
            nleading = allzero.argmax()
            allzero_trim = np.trim_zeros(allzero)
            gap_inds = np.where(allzero_trim == 0)[0] + nleading
        else:
            gap_inds = np.where(allzero == 0)[0]
        mkrdata[marker + '_gaps'] = gap_inds
    return mkrdata


def get_roi(vicon):
    """ Return array of frames corresponding to Nexus ROI """
    roi = vicon.GetTrialRegionOfInterest()
    return np.arange(roi[0], roi[1])


def get_model_data(vicon, model):
    """ Read model output variables (e.g. Plug-in Gait) """
    modeldata = dict()
    var_dims = (3, vicon.GetFrameCount())
    subj = get_subjectnames()
    for var in model.read_vars:
        nums, bools = vicon.GetModelOutput(subj, var)
        if nums:
            data = np.squeeze(np.array(nums))
        else:
            if model.is_optional_var(var):
                logger.info('cannot read optional variable %s, returning nans'
                            % var)
                data = np.empty(var_dims)
                data[:] = np.nan
            else:
                raise GaitDataError('Cannot read model variable %s. Make sure '
                                    'that the appropriate model has been run.'
                                    % var)
        modeldata[var] = data
    return modeldata


def create_events(vicon, context, strikes, toeoffs):
    """ Create foot strike and toeoff events in Nexus. """
    logger.debug('marking events in Nexus')
    side_str = 'Right' if context == 'R' else 'Left'
    subjectname = get_subjectnames()
    for fr in strikes:
        vicon.CreateAnEvent(subjectname, side_str, 'Foot Strike',
                            int(fr+1), 0)
    for fr in toeoffs:
        vicon.CreateAnEvent(subjectname, side_str, 'Foot Off',
                            int(fr+1), 0)
