import sympy as sp

import numpy as np

from logger import Logger

from util import to_np_mat

from analysis import construct_muscle_space_inequality

from scipy.optimize import minimize

# ------------------------------------------------------------------------

# Task Space

# ------------------------------------------------------------------------

class TaskSpace:

    """Task space EoM of a given task.

    Assuming the following joint space convention:

    M qddot + f = tau

    task projection is given by:

    ft = Lt (xddot - JDotQDot) + JtBarT f

    Lt = (R M^{-1} R^T)^{-1}

    JtBarT = Lt R M^{-1}

    NtT = (I - JtT JtBarT)

    and the resultant task space forces are given by:

    tau = JtT ft + NtT f

    """

    def __init__(self, model, xt):

        """ Constructor.

        Parameters

        ----------

        model: model

        xt: [d x 1] task positions

        """

        self.logger = Logger('TaskSpace')

        self.model = model

        self.xt = xt

        self.__construct_task_space_variables()

    def __construct_task_space_variables(self):

        """Calculate task Jacobian Jt and JtDot * qDot

        """

        self.Jt = self.xt.jacobian(self.model.Q())

        self.JtT = self.Jt.transpose()

        JtDot = sp.diff(self.Jt, self.model.t)

        self.JtDotQDot = JtDot * sp.Matrix(self.model.U())

        # the derivative of a matrix with a vector is a rank 3 tensor (3D

        # array), [dM/dq1, dM/dq2, ...]

        self.JtTDq = sp.derive_by_array(self.JtT, self.model.Q())

    def calculate_force(self, xddot, pose):

        """Calculate task forces.

        For a given end effector acceleration compute the joint space

        torques:

        ft = Lt (xddot - JDotQDot) + JtBarT f

        tau = JtT ft + NtT f

        Parameters

        ----------

        xddot: [2 x 1] a sympy Matrix containing the desired accelerations in 2D

        pose: system constants, coordinates and speeds as dictionary

        Returns

        -------

        (tau, ft): required torque and task forces to track the desired

        acceleration

        """

        M = to_np_mat(self.model.M.subs(pose))

        f = to_np_mat(self.model.f.subs(pose))

        Jt = to_np_mat(self.Jt.subs(pose))

        JtDotQDot = to_np_mat(self.JtDotQDot.subs(pose))

        JtT = to_np_mat(self.JtT.subs(pose))

        MInv = np.linalg.inv(M)

        LtInv = Jt * MInv * JtT

        Lt = np.linalg.pinv(LtInv)

        JtBarT = Lt * Jt * MInv

        NtT = np.asmatrix(np.eye(len(JtT))) - JtT * JtBarT

        ft = Lt * (xddot - JtDotQDot) + JtBarT * f

        return JtT * ft + 0 * NtT * f, ft

    def x(self, pose):

        """For a given pose (q) evaluate the task position.

        Parameters

        ----------

        pose: dictionary of model parameters and q's

        Returns

        -------

        x: sympy Matrix

        """

        return self.xt.subs(pose)

    def u(self, pose, qdot):

        """For a given pose (q, u) evaluate the task velocity.

        Parameters

        ----------

        pose: dictionary of model parameters, q's and u's

        qdot: an array of u(t)

        Returns

        ------

        u: sympy Matrix

        """

        return self.Jt.subs(pose) * sp.Matrix(qdot)

# ------------------------------------------------------------------------

# Muscle Space

# ------------------------------------------------------------------------

class MuscleSpace:

    """ Muscle space EoM.

    """

    def __init__(self, model, use_optimization=False):

        """ Constructor

        Parameters

        ----------

        model: model

        """

        self.logger = Logger('MuscleSpace')

        self.model = model

        self.use_optimization = use_optimization

        self.Fmax = to_np_mat(self.model.Fmax)

        self.x_max = np.max(self.Fmax)

    def calculate_force(self, lmdd, pose):

        """Calculate muscle forces.

        For a given muscle length acceleration compuyte the muscle space EoM of

        motion and required muscle force to track the goal acceleration.

        fm_par = -Lm (lmdd - RDotQDot) - RBarT f

        tau = -R^T fm_par

        Parameters

        ----------

        lmdd: [m x 1] a sympy Matrix containing the desired muscle length

        accelerations

        pose: dictionary

        Returns

        -------

        (tau, fm_par + fm_perp) required torque to track the desired acceleration

        """

        M = to_np_mat(self.model.M.subs(pose))

        f = to_np_mat(self.model.f.subs(pose))

        R = to_np_mat(self.model.R.subs(pose))

        RT = R.transpose()

        RDotQDot = to_np_mat(self.model.RDotQDot.subs(pose))

        MInv = np.linalg.inv(M)

        LmInv = R * MInv * RT

        Lm = np.linalg.pinv(LmInv)

        RBarT = np.linalg.pinv(RT)

        NR = np.asmatrix(np.eye(len(RBarT)) - RBarT * RT)

        fm_par = -Lm * (lmdd - RDotQDot) - RBarT * f

        # Ensure fm_par > 0 not required for simulation, but for muscle analysis

        # otherwise muscle forces will be negative. Since RT * NR = 0 the null

        # space term does not affect the resultant torques.

        m = fm_par.shape[0]

        fm_0 = np.zeros((m, 1))

        if self.use_optimization:

            Z, B = construct_muscle_space_inequality(NR, fm_par, self.Fmax)

            def objective(x):

                return np.sum(x**2)

            def inequality_constraint(x):

                return np.array(B - Z * (x.reshape(-1, 1))).reshape(-1,)

            x0 = np.zeros(m)

            bounds = tuple([(-self.x_max, self.x_max) for i in range(0, m)])

            constraints = ({'type': 'ineq', 'fun': inequality_constraint})

            sol = minimize(objective, x0,  method='SLSQP',

                           bounds=bounds,

                           constraints=constraints)

            fm_0 = sol.x.reshape(-1, 1)

            if sol.success == False:

                raise RuntimeError('Some muscles are too week for this action')

        fm_perp = NR * fm_0

        return -RT * fm_par, fm_par + fm_perp