#! /usr/bin/env python
"""Module for executing control for a given effector"""

from __future__ import division
import os
import json
from math import asin, atan, radians

import numpy as np

from masterloop import mp


FRAME_TORSO = 0
AXIS_XYZ = 7
K = 0.1

_here = os.path.dirname(os.path.realpath(__file__))
with open('{}/../config/default.yaml'.format(_here)) as f:
    _xxx = json.load(f)

ARM = _xxx['arm']
MY_SHOULDERS = {
    'L': _xxx['lsh'],
    'R': _xxx['rsh']
}
JOINTS = ('ShoulderPitch', 'ShoulderRoll', 'ElbowYaw', 'ElbowRoll')
NAO_ARM = 0.22


def xyz(joint):
    return np.array(mp.getPosition(joint, FRAME_TORSO, False),
                    dtype=np.float64)[:3]


crt_xyz = {
    'L': xyz('LArm'),
    'R': xyz('RArm')
}


crt_joint = {
    'L': np.array(mp.getAngles(['L' + j for j in JOINTS], False),
                  dtype=np.float64),
    'R': np.array(mp.getAngles(['R' + j for j in JOINTS], False),
                  dtype=np.float64)
}


best_guess = {
    'L': np.random.randn(3, 4),
    'R': np.random.randn(3, 4)
}


def get_transform(joint):
    """Retrieve a transform matrix of a joint in the torso frame."""
    return np.array(
        mp.getTransform(joint, FRAME_TORSO, False), dtype=np.float64
    ).reshape((4, 4))


def diff_jacobian(side):
    """Update and return the differential Jacobian. NOT GOOD."""
    new_end_xyz = xyz('{}Arm'.format(side))
    delta_r = new_end_xyz - crt_xyz[side]
    crt_xyz[side] = new_end_xyz.copy()

    new_joint = np.array(mp.getAngles([side + j for j in JOINTS], False))
    delta_j = new_joint - crt_joint[side]
    crt_joint[side] = new_joint.copy()

    J = delta_r[:,None] / delta_j
    okval = np.isfinite(J)
    J[~okval] = best_guess[side][~okval]
    best_guess[side][okval] = J[okval]
    return J


def jacobian(side):
    """Calculate the analytical Jacobian for side 'L' or 'R'."""
    end_xyz = xyz('{}Arm'.format(side))
    xyzs = np.array([xyz(side + j) for j in JOINTS])

    x_axis = np.array([[1, 0, 0, 1]]).T
    y_axis = np.array([[0, 1, 0, 1]]).T
    z_axis = np.array([[0, 0, 1, 1]]).T

    ax_order = (y_axis, z_axis, x_axis, z_axis)

    axes = np.concatenate(
        [get_transform(side + j).dot(ax) for j, ax in zip(JOINTS, ax_order)],
        axis=1
    )[:3].T

    J = np.cross(axes, end_xyz - xyzs).T
    return J


def to_nao(my_xyz, side):
    """Transform object coordinates from operator chest frame to NAO torso."""
    sh_offs = np.array(MY_SHOULDERS[side])
    my_sh_xyz = my_xyz - sh_offs
    nao_sh_xyz = my_sh_xyz / ARM * NAO_ARM
    nao_xyz = nao_sh_xyz + xyz(side + 'Shoulder')
    return nao_xyz


def inv_jacobian(j):
    """Singluarity robust inverse Jacobian."""
    u, s, vt = np.linalg.svd(j)
    s_inv = np.zeros((vt.shape[0], u.shape[1]))
    np.fill_diagonal(s_inv, 1 / s)
    s_inv[np.abs(s_inv) > 30] = 0
    return vt.T.dot(s_inv).dot(u.T)


def _some_cartesian(my_xyz, side, jfunc):
    """A generic cartesian controller.

    Parameters
    ----------
    my_xyz : numpy.array
        Coordinates of the object/target in the operator chest frame.
    side : str, 'L' or 'R'
    jfunc : func
        A function that will return a jacobian matrix for the given side.

    Returns
    -------
    numpy.array
        The control to execute in the joint space.

    """
    nao_xyz = to_nao(my_xyz, side)
    delta_r = 0.1 * (nao_xyz - xyz('{}Arm'.format(side)))
    crt_q = mp.getAngles([side + j for j in JOINTS], False)
    delta_q =  inv_jacobian(jfunc(side)).dot(delta_r).flatten()
    return crt_q + delta_q


def our_cartesian(my_xyz, side):
    """Our cartesian controller."""
    return _some_cartesian(my_xyz, side, jacobian)


def our_diff_cartesian(my_xyz, side):
    """Our differential cartesian controller."""
    return _some_cartesian(my_xyz, side, diff_jacobian)


def _elbow(arm_):
    "Dumb way to calculate the elbow roll."""
    quot = min(1.0, arm_ / 0.70)
    return radians(180 * (1 - quot))


def dumb(my_xyz, side):
    """Our dumb controller that directly maps 3D coords into joint space."""
    sign = 1 if side == 'L' else -1
    roll = asin(my_xyz[1] / np.linalg.norm(my_xyz))
    roll -= sign * radians(25)
    pitch = atan(-my_xyz[2] / np.abs(my_xyz[0]))

    eroll = -sign * _elbow(np.linalg.norm(my_xyz))
    eyaw = -sign * radians(40)
    return np.array([pitch, roll, eyaw, eroll])


def movement(my_xyz, side, control):
    """Execute the movement.

    my_xyz : numpy.array
        Coordinates of the object/target in the operator chest frame.
    side : str, 'L' or 'R'
    control : func
        A function to calculate the conrol for the side, depending on the
        target coordinates. It must return returning the `numpy.array`
        joint vector of format
        [ShoulderPitch, ShoulderRoll, ElbowYaw, ElbowRoll].

    """
    ref = control(np.array(my_xyz), side).tolist()
    mp.setAngles([side + j for j in JOINTS], ref, 0.3)


def nao_cart_movement(my_arm_xyz, side, *_):
    """Execute the movement using the NAO cartesian controller."""
    nao_xyz = to_nao(my_arm_xyz, side)
    mp.setPositions('{}Arm'.format(side), FRAME_TORSO,
                    tuple(nao_xyz.tolist()) + (0, 0, 0),
                    1.0, AXIS_XYZ)