import sympy as sp
# import scipy as scp
import numpy as np
import pylab as plt
from sympy import cos, sin
from sympy.physics.mechanics import dynamicsymbols
from pydy.codegen.ode_function_generators import generate_ode_function
from util import coordinate_limiting_force, coriolis_matrix, \
    apply_generalized_force, substitute
from logger import Logger
from functools import reduce

# ------------------------------------------------------------------------
# ArmModel
# ------------------------------------------------------------------------


class ArmModel:
    """A planar toy model composed of 3 degrees of freedom and 9 muscles.

    The forces that act on the model are: gravity (tau_g), Coriolis and
    centrifugal (tau_c), coordinate limiting forces (tau_l) and viscous joints
    (tau_b). We assume the following:

    M qddot + tau_c + tau_g = tau + tau_l + tau_b

    M qddot = forcing

    forcing = tau - f, f = tau_c + tau_g - tau_l - tau_b

    or

    M qddot + f = tau

    Furthermore, we use 9 linear muscles. The musculotendon lengths and their
    derivatives are given by lm, lmd, lmdd. The muscle's moment arm matrix (R)
    is given by:

    R =  d lm(q) / qdot

    such as that

    lmd = R qdot, lmdd = R qddot + Rdot qdot

    tau = -R^T f_m

    Notes
    -----

    (a) The model can represent a 2-link (used for validation) or 3-link system
    by changing nd = 2/3. Unfortunately, the muscle geometry is defined for the
    3-body case.

    (b) The Newton's 3rd law is applied for each force and torque. For example,
    for the first body: tau_1 - tau_2 is applied.

    (c) For the derivation of the equations of motion we used the Euler-Lagrange
    Method assuming:

    M qddot + C qdot + d V / dq = tau, V: potential energy, C: Coriolis Matrix

    """

    def __init__(self, use_gravity=0, use_coordinate_limits=1, use_viscosity=1):
        """
        """
        self.logger = Logger('ArmModel')
        # n: DoFs, md: muscles
        self.nd = 3
        self.md = 9
        # used for selecting since we use 1-based indexing
        self.s = self.nd + 1
        # used for iterating
        self.dim = list(range(1, self.s))
        # enable/disable gravity in EoM [0/1] (simulation too slow)
        self.use_gravity = use_gravity
        # enable/disable coordinate limits in EoM [0/1]
        self.use_coordinate_limits = use_coordinate_limits
        # enable/disable viscous joints in EoM [0/1]
        self.use_viscosity = use_viscosity
        # true if constants are not substituted
        self.sub_constants = True
        # default model state
        self.state0 = np.array([
            np.deg2rad(45.0),  # q1
            np.deg2rad(45.0),  # q2
            np.deg2rad(45.0),  # q3
            np.deg2rad(0.00),  # u1
            np.deg2rad(0.00),  # u2
            np.deg2rad(0.00)   # u3
        ])
        # reference pose used for calculating the optimal fiber length
        self.reference_pose = np.array([
            np.deg2rad(60.0),  # q1
            np.deg2rad(70.0),  # q2
            np.deg2rad(50.0),  # q3
        ])
        self.logger.debug('Constructing model...')
        # construct model
        self.__construct_symbols()
        self.__construct_kinematics()
        self.__construct_coordinate_limiting_forces()
        self.__construct_kinetics()
        self.__construct_drawables()
        self.__construct_muscle_geometry()
        self.__define_muscle_parameters()
        self.__construct_rhs()

    def __construct_symbols(self):
        """Define the symbols used by the analytical model.

        """
        self.logger.debug('__construct_symbols')
        # define model constants 10 a's and b's are used but indexed from 1-9
        # thus 0 is not used (for correspondence with the paper)
        self.a = sp.Matrix(sp.symbols('a0:10'))
        self.b = sp.Matrix(sp.symbols('b0:10'))
        # segment lengths
        self.L = sp.Matrix(sp.symbols('L0:4'))
        # distance to segment's CoM
        self.Lc = sp.Matrix(sp.symbols('Lc0:4'))
        # body parameters
        # inertia
        self.Iz = sp.Matrix(sp.symbols('Iz0:4'))
        # mass
        self.m = sp.Matrix(sp.symbols('m0:4'))
        # time
        self.t = sp.symbols('t')
        # gravity
        self.g = sp.symbols('g')
        # q's are used instead of $\theta$
        self.q = sp.Matrix(dynamicsymbols('theta0:4'))
        self.u = sp.Matrix(dynamicsymbols('u:4'))
        self.dq = sp.Matrix([sp.diff(x, self.t) for x in self.q])
        self.ddq = sp.Matrix([sp.diff(x, self.t) for x in self.dq])
        # tau acting forces
        self.tau = sp.Matrix(dynamicsymbols('tau0:4'))
        # define a dictionary that maps symbols to values
        # parameters are derived from [1]
        self.constants = dict({self.a[1]: 0.055, self.a[2]: 0.055, self.a[3]:
                               0.220, self.a[4]: 0.24, self.a[5]: 0.040,
                               self.a[6]: 0.040, self.a[7]: 0.220, self.a[8]:
                               0.06, self.a[9]: 0.26, self.b[1]: 0.080,
                               self.b[2]: 0.11, self.b[3]: 0.030, self.b[4]:
                               0.03, self.b[5]: 0.045, self.b[6]: 0.045,
                               self.b[7]: 0.048, self.b[8]: 0.050, self.b[9]:
                               0.03, self.L[1]: 0.310, self.L[2]: 0.270,
                               self.L[3]: 0.150, self.Lc[1]: 0.165, self.Lc[2]:
                               0.135, self.Lc[3]: 0.075, self.m[1]: 1.93,
                               self.m[2]: 1.32, self.m[3]: 0.35, self.Iz[1]:
                               0.0141, self.Iz[2]: 0.0120, self.Iz[3]: 0.001,
                               self.g: 9.81})

        # pickle workaround
        for q in self.q:
            q.__class__.__module__ = '__main__'

        for dq in self.dq:
            dq.__class__.__module__ = '__main__'

        for ddq in self.ddq:
            ddq.__class__.__module__ = '__main__'

        for u in self.u:
            u.__class__.__module__ = '__main__'

        for tau in self.tau:
            tau.__class__.__module__ = '__main__'

        # max isometrix force
        # TODO all muscles are equally strong
        fmax = 50
        self.Fmax = sp.diag(fmax, fmax, fmax,
                            fmax, 0.5 * fmax, 0.5 * fmax,
                            fmax, fmax, 0.5 * fmax)

    def __construct_kinematics(self):
        """Define points of interest for the derivation of the EoM. These are used for
        constructing the EoM.

        """
        self.logger.debug('__construct_kinematics')
        L = self.L
        Lc = self.Lc
        q = self.q
        # define the spatial coordinates for the Lc in terms of Lc s' and q's
        # arm
        xc1 = sp.Matrix([Lc[1] * cos(q[1]),
                         Lc[1] * sin(q[1]),
                         0,
                         0,
                         0,
                         q[1]])
        # forearm
        xc2 = sp.Matrix([L[1] * cos(q[1]) + Lc[2] * cos(q[1] + q[2]),
                         L[1] * sin(q[1]) + Lc[2] * sin(q[1] + q[2]),
                         0,
                         0,
                         0,
                         q[1] + q[2]])

        # hand
        xc3 = sp.Matrix([L[1] * cos(q[1]) + L[2] * cos(q[1] + q[2]) +
                         Lc[3] * cos(q[1] + q[2] + q[3]),
                         L[1] * sin(q[1]) + L[2] * sin(q[1] + q[2]) +
                         Lc[3] * sin(q[1] + q[2] + q[3]),
                         0,
                         0,
                         0,
                         q[1] + q[2] + q[3]])
        self.xc = [sp.Matrix([0]), xc1, xc2, xc3]
        # CoM velocities
        self.vc = [sp.diff(x, self.t) for x in self.xc]
        # calculate CoM Jacobian
        self.Jc = [x.jacobian(self.QDot()) for x in self.vc]

    def __construct_kinetics(self):
        """Construct model's dynamics (M, tau_c, tau_g).

        """
        self.logger.debug('__construct_kinetics')
        # generate the mass matrix [6 x 6] for each body
        self.M = [sp.diag(self.m[i], self.m[i], self.m[i], 0, 0,
                          self.Iz[i]) for i in self.dim]
        # dummy 0 for 1-based indexing
        self.M.insert(0, 0)
        # map spatial to generalized inertia
        self.M = [self.Jc[i].T * self.M[i] * self.Jc[i]
                  for i in self.dim]
        # sum the mass product of each body
        self.M = reduce(lambda x, y: x + y, self.M)
        self.M = sp.trigsimp(self.M)
        # Coriolis matrix
        self.C = sp.trigsimp(coriolis_matrix(self.M, self.Q(), self.QDot()))
        # Coriolis forces
        self.tau_c = sp.trigsimp(self.C * sp.Matrix(self.QDot()))
        # potential energy due to gravity force
        self.V = 0
        for i in self.dim:
            self.V = self.V + self.m[i] * self.g * self.xc[i][1]

        self.tau_g = sp.Matrix([sp.diff(self.V, x) for x in self.Q()])

    def __construct_coordinate_limiting_forces(self):
        """Construct coordinate limiting forces.

        """
        self.logger.debug('__construct_coordinate_limiting_forces')

        a = 5
        b = 50
        q_low = [None, np.deg2rad(5), np.deg2rad(5), np.deg2rad(5)]
        q_up = [None, np.deg2rad(175), np.pi, np.deg2rad(100)]

        self.__tau_l = [coordinate_limiting_force(self.q[i], q_low[i], q_up[i], a,
                                                  b) for i in range(1, self.s)]

    def __construct_drawables(self):
        """Construct points of interest (e.g. muscle insertion, CoM, joint centers).

        """
        self.logger.debug('__construct_drawables')
        a = self.a
        b = self.b
        q = self.q
        L = self.L
        # define muscle a
        self.ap = [[-a[1], sp.Rational(0)],  # a1
                   [a[2], sp.Rational(0)],  # a2
                   [a[3] * cos(q[1]), a[3] * sin(q[1])],  # a3
                   [a[4] * cos(q[1]), a[4] * sin(q[1])],  # a4
                   [-a[5], sp.Rational(0)],  # a5
                   [a[6], sp.Rational(0)],  # a6
                   [L[1] * cos(q[1]) + a[7] * cos(q[1] + q[2]),
                    L[1] * sin(q[1]) + a[7] * sin(q[1] + q[2])],  # a7
                   [L[1] * cos(q[1]) + a[8] * cos(q[1] + q[2]),
                    L[1] * sin(q[1]) + a[8] * sin(q[1] + q[2])],  # a8
                   [a[9] * cos(q[1]), a[9] * sin(q[1])]  # a9
                   ]
        # define muscle b
        self.bp = [[b[1] * cos(q[1]), b[1] * sin(q[1])],  # b1
                   [b[2] * cos(q[1]), b[2] * sin(q[1])],  # b2
                   [L[1] * cos(q[1]) + b[3] * cos(q[1] + q[2]),
                    L[1] * sin(q[1]) + b[3] * sin(q[1] + q[2])],  # b3
                   [L[1] * cos(q[1]) - b[4] * cos(q[1] + q[2]),
                    L[1] * sin(q[1]) - b[4] * sin(q[1] + q[2])],  # b4
                   [L[1] * cos(q[1]) + b[5] * cos(q[1] + q[2]),
                    L[1] * sin(q[1]) + b[5] * sin(q[1] + q[2])],  # b5
                   [L[1] * cos(q[1]) - b[6] * cos(q[1] + q[2]),
                    L[1] * sin(q[1]) - b[6] * sin(q[1] + q[2])],  # b6
                   [L[1] * cos(q[1]) + L[2] * cos(q[1] + q[2]) + b[7] * cos(q[1] + q[2] + q[3]),
                    L[1] * sin(q[1]) + L[2] * sin(q[1] + q[2]) + b[7] * sin(q[1] + q[2] + q[3])],  # b7
                   [L[1] * cos(q[1]) + L[2] * cos(q[1] + q[2]) - b[8] * cos(q[1] + q[2] + q[3]),
                    L[1] * sin(q[1]) + L[2] * sin(q[1] + q[2]) - b[8] * sin(q[1] + q[2] + q[3])],  # b8
                   [L[1] * cos(q[1]) + L[2] * cos(q[1] + q[2]) + b[9] * cos(q[1] + q[2] + q[3]),
                    L[1] * sin(q[1]) + L[2] * sin(q[1] + q[2]) + b[9] * sin(q[1] + q[2] + q[3])]  # b9
                   ]
        # define CoM
        self.bc = [[self.xc[1][0], self.xc[1][1]],
                   [self.xc[2][0], self.xc[2][1]],
                   [self.xc[3][0], self.xc[3][1]]
                   ]
        # joint center
        self.jc = [[sp.Rational(0), sp.Rational(0)],
                   [L[1] * cos(q[1]), L[1] * sin(q[1])],
                   [L[1] * cos(q[1]) + L[2] * cos(q[1] + q[2]),
                    L[1] * sin(q[1]) + L[2] * sin(q[1] + q[2])]
                   ]
        # end effector
        if self.nd == 3:
            self.ee = sp.Matrix([L[1] * cos(q[1]) + L[2] * cos(q[1] + q[2]) +
                                 L[3] * cos(q[1] + q[2] + q[3]), L[1] *
                                 sin(q[1]) + L[2] * sin(q[1] + q[2]) + L[3] *
                                 sin(q[1] + q[2] + q[3])])
        elif self.nd == 2:
            self.ee = sp.Matrix([L[1] * cos(q[1]) + L[2] * cos(q[1] + q[2]),
                                 L[1] * sin(q[1]) + L[2] * sin(q[1] + q[2])
                                 ])

    def __construct_muscle_geometry(self):
        """Construct muscle length function and moment arm.

        """
        self.logger.debug('__construct_geometry')
        a = self.a
        b = self.b
        L = self.L
        q = self.q
        # muscle length ($l(q)$)
        self.lm = sp.Matrix([
            (a[1] ** 2 + b[1] ** 2 + 2 * a[1] *
             b[1] * cos(q[1])) ** sp.Rational(1, 2),
            (a[2] ** 2 + b[2] ** 2 - 2 * a[2] *
             b[2] * cos(q[1])) ** sp.Rational(1, 2),
            ((L[1] - a[3]) ** 2 + b[3] ** 2 + 2 * (L[1] - a[3])
             * b[3] * cos(q[2])) ** sp.Rational(1, 2),
            ((L[1] - a[4]) ** 2 + b[4] ** 2 - 2 * (L[1] - a[4])
             * b[4] * cos(q[2])) ** sp.Rational(1, 2),
            (a[5] ** 2 + b[5] ** 2 + L[1] ** 2 + 2 * a[5] * L[1] * cos(q[1]) +
             2 * b[5] * L[1] * cos(q[2]) + 2 * a[5] * b[5] * cos(q[1] + q[2])) ** sp.Rational(1, 2),
            (a[6] ** 2 + b[6] ** 2 + L[1] ** 2 - 2 * a[6] * L[1] * cos(q[1]) -
             2 * b[6] * L[1] * cos(q[2]) + 2 * a[6] * b[6] * cos(q[1] + q[2])) ** sp.Rational(1, 2),
            ((L[2] - a[7]) ** 2 + b[7] ** 2 + 2 * (L[2] - a[7])
             * b[7] * cos(q[3])) ** sp.Rational(1, 2),
            ((L[2] - a[8]) ** 2 + b[8] ** 2 - 2 * (L[2] - a[8])
             * b[8] * cos(q[3])) ** sp.Rational(1, 2),
            ((L[1] - a[9]) ** 2 + b[9] ** 2 + L[2] ** 2 + 2 * (L[1] - a[9]) * L[2] * cos(q[2]) +
             2 * b[9] * L[2] * cos(q[3]) +
             2 * (L[1] - a[9]) * b[9] * cos(q[2] + q[3])) ** sp.Rational(1, 2)
        ])

        self.lmd = sp.diff(self.lm, self.t)
        self.lmdd = sp.diff(self.lmd, self.t)

        # calculate the moment arm matrix and its derivatives
        self.R = sp.trigsimp(self.lm.jacobian(self.Q()))
        self.RDot = sp.diff(self.R, self.t)
        self.RDotQDot = self.RDot * sp.Matrix(self.U())

    def __define_muscle_parameters(self):
        """Define muscle parameters, such as optimal fiber length (reference pose) and
        tendon stiffness.

        """
        self.logger.debug('__define_muscle_parameters')
        parameters = self.model_parameters(q=self.reference_pose,
                                           in_deg=False)
        self.lm0 = self.lm.subs(parameters)
        # the derivative of a matrix with a vector is a rank 3 tensor (3D
        # array), [dM/dq1, dM/dq2, ...]
        self.RTDq = sp.derive_by_array(self.R.transpose(), self.Q())

    def __construct_rhs(self):
        """Construct a callable function that can be used to integrate the mode.

        rhs = rhs(x, t, controller specifieds, parameters values)

        """
        self.logger.debug('__construct_rhs')
        self.logger.debug('Use Gravity: ' + str(self.use_gravity))
        self.logger.debug('Use Coordinate Limits: ' +
                          str(self.use_coordinate_limits))
        self.logger.debug('Use Viscous Joints: ' + str(self.use_viscosity))

        # forces
        # Newton's 3rd law
        b = 0.05
        tau = sp.Matrix(apply_generalized_force(self.Tau()))
        self.tau_l = sp.Matrix(apply_generalized_force(self.__tau_l))
        self.tau_b = -b * sp.Matrix(apply_generalized_force(self.U()))
        # tau = sp.Matrix(self.Tau())
        # self.tau_l = sp.Matrix(self.__tau_l)
        # self.tau_b = -b *sp.Matrix(self.U())

        # M qdd + tau_c + tau_g = tau + tau_l + tau_b-> M qdd = forcing
        # f = tau_c + tau_g - tau_l - tau_b
        self.f = self.tau_c \
            + self.use_gravity * self.tau_g \
            - self.use_coordinate_limits * self.tau_l \
            - self.use_viscosity * self.tau_b
        self.forcing = tau - self.f

        # substitute dq with u (required for code-gen)
        for i in range(0, self.forcing.shape[0]):
            self.forcing = self.forcing.subs(self.dq[i + 1], self.u[i + 1])

        # rhs
        self.coordinates = sp.Matrix(self.Q())
        self.speeds = sp.Matrix(self.U())
        self.coordinates_derivatives = self.speeds
        self.specifieds = sp.Matrix(self.Tau())
        self.rhs = generate_ode_function(
            self.forcing,
            self.coordinates,
            self.speeds,
            list(self.constants.keys()),
            mass_matrix=self.M,
            coordinate_derivatives=self.coordinates_derivatives,
            specifieds=self.specifieds)

# ------------------------------------------------------------------------
# ArmModel public interface
# ------------------------------------------------------------------------

    def Q(self):
        'Get active coordinates (q) [1, s].'
        return self.q[1:self.s]

    def QDot(self):
        'Get active speeds (qdot) [1, s].'
        return self.dq[1:self.s]

    def U(self):
        'Get active speeds (qdot = u) [1, s].'
        return self.u[1:self.s]

    def Tau(self):
        'Get active speeds (tau) [1, s].'
        return self.tau[1:self.s]

    def model_parameters(self, **kwargs):
        """Get the model parameters dictionary given q, u, in_deg=[True/False].

        Parameters
        ----------

        kwargs: q=q, u=u, in_deg=[True/False]

        """

        expected_args = ["q", "u", "in_deg"]
        kwargsdict = {}
        for key in list(kwargs.keys()):
            if key in expected_args:
                kwargsdict[key] = kwargs[key]
            else:
                raise Exception("Unexpected Argument")

        return self.__model_parameters(kwargsdict)

    def __model_parameters(self, dic):
        """A private implementation of model_parameters.

        Parameters
        ----------

        dic: dictionary containing {q, u, in_deg}

        """

        constants = {}
        if self.sub_constants:
            constants = self.constants.copy()

        q = self.Q()
        dq = self.QDot()
        u = self.U()
        in_deg = False
        if "in_deg" in list(dic.keys()):
            in_deg = dic["in_deg"]

        if "q" in list(dic.keys()):
            qs = dic["q"]
            if in_deg:
                qs = np.deg2rad(dic["q"])

            constants.update({q[i]: qs[i] for i in range(0, self.nd)})

        if "u" in list(dic.keys()):
            us = dic["u"]

            constants.update({dq[i]: us[i] for i in range(0, self.nd)})

            constants.update({u[i]: us[i] for i in range(0, self.nd)})

        return constants

    def pre_substitute_parameters(self):
        """Substitute model parameters into the variables of the model to improve speed.

        """

        self.logger.debug('pre_substitute_parameters')
        self.sub_constants = False
        constants = self.constants

        self.M = self.M.subs(constants)
        self.tau_c = self.tau_c.subs(constants)
        self.tau_g = self.tau_g.subs(constants)
        self.tau_l = self.tau_l.subs(constants)
        self.tau_b = self.tau_b.subs(constants)
        self.f = self.f.subs(constants)

        self.lm = self.lm.subs(constants)
        self.lmd = self.lmd.subs(constants)
        self.lmdd = self.lmdd.subs(constants)
        self.R = self.R.subs(constants)
        self.RDot = self.RDot.subs(constants)
        self.RDotQDot = self.RDotQDot.subs(constants)
        self.RTDq = self.RTDq.subs(constants)

        self.ap = substitute(self.ap, constants)
        self.bp = substitute(self.bp, constants)
        self.bc = substitute(self.bc, constants)
        self.jc = substitute(self.jc, constants)
        self.ee = substitute(self.ee, constants)

        # self.L = self.L.subs(constants)
        # self.Lc = self.Lc.subs(constants)
        # self.Iz = self.Iz.subs(constants)

        self.xc = substitute(self.xc, constants)
        self.vc = substitute(self.vc, constants)
        self.Jc = substitute(self.Jc, constants)

        # self.m = self.m.subs(constants)
        # self.g = self.g.subs(constants)

    def draw_model(self, q, in_deg, ax=None, scale=0.8, use_axis_limits=True, alpha=1.0,
                   text=True):
        """Draws the 3D toy model.

        Parameters
        ----------

        q: coordinate values in degrees or rad
        in_deg: True/False
        ax: axis 1D
        scale: if figure is small this helps in visualizing details
        use_axis_limits: use axis limits from max length

        """
        if self.nd == 2:
            self.logger.debug('draw_model supports 3DoFs case')
            return

        constants = self.model_parameters(q=q, in_deg=in_deg)

        joints = substitute(self.jc, constants)
        muscle_a = substitute(self.ap, constants)
        muscle_b = substitute(self.bp, constants)
        end_effector = substitute(self.ee, constants)
        CoM = substitute(self.bc, constants)

        linewidth = 4 * scale
        gd_markersize = 14 * scale
        jc_markersize = 12 * scale
        mo_markersize = 7 * scale
        ef_markersize = 15 * scale
        fontsize = 12 * scale

        if ax == None:
            fig, ax = plt.subplots(1, 1, figsize=(5, 5))

        # arm
        ax.plot([joints[0, 0], joints[1, 0]], [joints[0, 1], joints[1, 1]],
                'r', linewidth=linewidth, alpha=alpha)
        # forearm
        ax.plot([muscle_b[5, 0], joints[2, 0]], [muscle_b[5, 1], joints[2, 1]],
                'r', linewidth=linewidth, alpha=alpha)
        # hand
        ax.plot([muscle_b[7, 0], end_effector[0]], [muscle_b[7, 1], end_effector[1]],
                'r', linewidth=linewidth, alpha=alpha)
        # muscles
        for i in range(0, muscle_a.shape[0]):
            ax.plot([muscle_a[i, 0], muscle_b[i, 0]], [
                muscle_a[i, 1], muscle_b[i, 1]], 'b', alpha=alpha)
            if text:
                ax.text(muscle_a[i, 0], muscle_a[i, 1],
                        r'$a_' + str(i + 1) + '$', fontsize=fontsize, alpha=alpha)
                ax.text(muscle_b[i, 0], muscle_b[i, 1],
                        r'$b_' + str(i + 1) + '$', fontsize=fontsize, alpha=alpha)
                ax.text((muscle_b[i, 0] + muscle_a[i, 0]) / 2.0,
                        (muscle_b[i, 1] + muscle_a[i, 1]) / 2.0,
                        r'$l_' + str(i + 1) + '$', fontsize=fontsize, alpha=alpha)

        # joint centers
        ax.plot(joints[:, 0], joints[:, 1], 'or',
                markersize=gd_markersize, alpha=alpha)
        if text:
            for i in range(0, joints.shape[0]):
                ax.text(joints[i, 0], joints[i, 1], r'$J_' +
                        str(i + 1) + '$', fontsize=fontsize, alpha=alpha)

        # CoM
        ax.plot(CoM[:, 0], CoM[:, 1], 'oy',
                markersize=jc_markersize, alpha=alpha)
        if text:
            for i in range(0, CoM.shape[0]):
                ax.text(CoM[i, 0], CoM[i, 1], r'$Lc_' +
                        str(i + 1) + '$', fontsize=fontsize, alpha=alpha)

        # end effector
        ax.plot(end_effector[0], end_effector[1],
                '<b', markersize=ef_markersize, alpha=alpha)
        # muscle origin
        ax.plot(muscle_a[:, 0], muscle_a[:, 1], 'dy',
                markersize=mo_markersize, alpha=alpha)
        # muscle insertion
        ax.plot(muscle_b[:, 0], muscle_b[:, 1], 'db',
                markersize=mo_markersize, alpha=alpha)

        ax.axis('equal')
        ax.set_title('Model Pose')
        ax.set_xlabel('$x \; (m)$')
        ax.set_ylabel('$y \; (m)$')

        # axis limits
        if use_axis_limits:
            L_max = self.constants[self.L[1]] + \
                self.constants[self.L[2]] + self.constants[self.L[3]]
            ax.set_xlim([-L_max, L_max])
            ax.set_ylim([-L_max / 2, 1.5 * L_max])


# ------------------------------------------------------------------------
# ArmModel tests
# ------------------------------------------------------------------------

    def test_muscle_geometry(self):
        """Evaluate the muscle geometry for a random pose against the ground truth.

        """
        ma = self.ap
        mb = self.bp
        lmt = sp.Matrix([sp.trigsimp(
            sp.sqrt(pow(ma[i][0] - mb[i][0], 2) +
                    pow(ma[i][1] - mb[i][1], 2))) for i in range(0, len(ma))])

        pose = self.model_parameters(q=np.random.random(3), in_deg=False)
        assert_if_same(self.lm.subs(pose), lmt.subs(pose))