"""
This is a collection of functions for loading files generated by the test UI
(TAC test data) and calculating TAC test metrics, including:
    - completion time
    - path efficiency
    - reaction time
"""

import json
import numpy as np

# check for success by adding this to the last target entered timestamp and
# checking if the final timestamp in the trial is greater than that (basically
# it is dwell time plus some margin)
success_check = 2.5


def completion_time(data):
    """
    Gets the movement time for the given trial, defined as the amount of time
    from movement initiation to the last time the target was entered. On
    unsuccessful trials, the last timestamp of the trial is returned.
    """
    ts = movement_initiation(data)
    te = last_target_enter(data)
    if te is None:
        return end_time(data)
    else:
        return te - ts


def cumulative_completion_rate(completion_times, inc, top):
    """
    Gets the cumulative completion rate data from an array of completion times.
    Starting from zero, time is incremented by `inc` until `top` is reached
    (inclusive) and the number of timestamps in `completion_times` under the
    current time is added to `counts`. The timestamps can be obtained from
    `completions()`.

    Parameters
    ----------
    completion_times : array
        List of completion times (one per trial).
    inc : float
        Amount of time increment per check.
    top : float
        Largest time at which to check the completion rate.

    Returns
    -------
    cutoffs : ndarray
        Array of cutoff times at which the completion rate is checked.
    counts : ndarray
        The number of trials completed faster than the corresponding cutoff.
    """
    ts = np.array(completion_times)
    cutoffs = np.arange(0, top+inc, inc)
    counts = np.zeros(cutoffs.shape)
    for i, t in enumerate(cutoffs):
        counts[i] = np.sum(ts < t)

    return cutoffs, counts


def path_efficiency(data):
    """
    Calculates path efficiency for a TAC test trial. The path efficiency is the
    optimal distance to the final point of a successful trial divided by the
    distance travelled in the trial. The distances are calculated as Euclidean
    distances in joint space.

    Parameters
    ----------
    pose : ndarray, shape = (n_timestamps, n_joints)
        Array of joint angles in time, where each row is the pose at a given
        time, and each column corresponds to a single joint.

    Returns
    -------
    perc : float
        The path efficiency for the given trial. `None` if the trial was not
        successful.
    """
    pose = pose_array(data)
    optdist = _trajectory_dist(np.vstack([pose[0], pose[-1]]))
    dist = _trajectory_dist(pose)
    return optdist / dist


def _trajectory_dist(pose):
    """
    Calculates the cumulative Euclidean distance travelled over a trajectory.

    Parameters
    ----------
    pose : 2D ndarray
        Array of poses in time where each row is a pose at a specific time
        and each column is a component (coordinate) of the pose, i.e. the shape
        is (n_timestamps, n_joints).

    Returns
    -------
    dist : float
        The cumulative distance travelled over the array of poses.
    """
    distances = np.sqrt(np.sum(np.square(np.diff(pose, axis=0)), axis=1))
    total = np.sum(distances)
    return total


def loadfile(filename):
    """
    Gets the data from a single log file.
    """
    data = None
    with open(filename, 'r') as fid:
        data = json.load(fid)

    return data


def movement_initiation(data):
    """
    Gets the timestamp corresponding to movement initiation, i.e. the first
    time a command other than 'no-contraction' was output.
    """
    commands = data['trial_data']['command']
    first = [i for i, c in enumerate(commands) if c != 'no-contraction'][0]
    ts = data['trial_data']['timestamp'][first]
    return ts


def last_target_enter(data):
    """
    Gets the timestamp corresponding to the last time the target was entered.
    Returns None if the target wasn't hit.
    """
    entered = data['trial_data']['target_entered']
    et = end_time(data)
    if entered:
        te = entered[-1]
        if et > te + success_check:
            te = None
        return te
    else:
        return None


def end_time(data):
    """
    Gets the final timestamp of the trial.
    """
    try:
        et = data['trial_data']['timestamp'][-1]
    except IndexError:
        print(data)
        et = 0
    return et


def check_success(data):
    """
    Returns whether or not the given trial was successful.
    """
    te = last_target_enter(data)
    if te is not None:
        return True
    else:
        return False


def get_target(data):
    """
    Gets the target pose.
    """
    return data['target']['pose']


def get_complexity(data):
    """
    Gets the complexity of the given trial data -- the number of degrees of
    freedom that are active, or under control of the subject.
    """
    return int(len(data['active_classes'])/2)


def get_dofs(data):
    """
    Gets the number of target DOFs (i.e. the number of different motion
    classes required to complete a trial).
    """
    return len(data['target']['pose'])


def pose_array(data):
    """
    Gets a 2D array of poses over time. The shape is (n_timestamps,
    n_joints).
    """
    pose = data['trial_data']['pose']
    timestamps = data['trial_data']['timestamp']

    n_timestamps = len(timestamps)
    n_joints = len(pose)
    pose_array = np.zeros((n_timestamps, n_joints))

    for i, (joint, angles) in enumerate(sorted(pose.items())):
        pose_array[:, i] = angles

    pose_array -= pose_array[0, :]
    return pose_array
