#!/usr/bin/env python
# coding: utf-8

# # Introduction

# We now have defined all of the necessary components to derive the equations of motion of the system using Kane's Method:
# 
# - Linear velocity of each mass center.
# - Angular velocity of each rigid body.
# - All contributing (external) forces.
# 
# We will use the `KanesMethod` class which provides an automated computation of the first order ordinary differential equations given the above quantities. (Any method can be used, such as Newton-Euler or Hamilton's method, but there are only automated methods for Kane's equations and Lagrangian dynamics at the moment).

# # Setup

# Import the solutions from the previous notebooks:

# In[ ]:


from __future__ import print_function, division
from solution.kinetics import *


# We will need the `KanesMethod` class and `trigsimp` to give a reasonably compact result.

# In[ ]:


from sympy import trigsimp
from sympy.physics.mechanics import KanesMethod


# Enable nice math printing:

# In[ ]:


from sympy.physics.vector import init_vprinting
init_vprinting(use_latex='mathjax', pretty_print=False)


# In[ ]:


get_ipython().run_line_magic('pinfo', 'KanesMethod')


# # Equations of Motion

# At the bare minimum for unconstrained systems, the `KanesMethod` class needs to know the generalized coordinates, the generalized speeds, the kinematical differential equations, the loads, the bodies, and a Newtonian reference frame. First, make a list of the generalized coordinates, i.e. the three joint angles:

# In[ ]:


coordinates = [theta1, theta2, theta3]
coordinates


# Now list the generalized speeds, i.e. the joint rates:

# In[ ]:


speeds = [omega1, omega2, omega3]
speeds


# The `KanesMethod` object can now be created with the inertial reference frame, coordinates, speeds, and kinematical differential equations:

# In[ ]:


kinematical_differential_equations


# In[ ]:


kane = KanesMethod(inertial_frame, coordinates, speeds, kinematical_differential_equations)


# Now list the six loads we defined and the three bodies:

# In[ ]:


loads = [lower_leg_grav_force,
         upper_leg_grav_force,
         torso_grav_force, 
         lower_leg_torque,
         upper_leg_torque,
         torso_torque]
loads


# In[ ]:


bodies = [lower_leg, upper_leg, torso]
bodies


# The equations of motion can now be computed using the `kanes_equations` method, which takes the list of loads and bodies. It returns the equations of motion (i.e. first order ordinary differential equations) in Kane's form:
# 
# $$ F_r + F_r^* = 0$$
# 
# which is essentially equivalent to the classic Newton_Euler form:
# 
# $$ \mathbf{F} = \mathbf{m}\mathbf{a} $$
# 
# $$ \mathbf{T} = \mathbf{I} \mathbf{\alpha} $$

# In[ ]:


fr, frstar = kane.kanes_equations(bodies, loads)


# In[ ]:


trigsimp(fr + frstar)


# Keep in mind that out utlimate goal is to have the equations of motion in first order form:
# 
# $$ \dot{\mathbf{x}} = \mathbf{g}(\mathbf{x}, t) $$
# 
# The equations of motion are linear in terms of the derivatives of the generalized speeds and the `KanesMethod` class automatically puts the equations in a more useful form for the next step of numerical simulation:
# 
# $$ \mathbf{M}(\mathbf{x}, t)\dot{\mathbf{x}} = \mathbf{f}(\mathbf{x}, t) $$
# 
# Note that
# 
# $$ \mathbf{g} = \mathbf{M}^{-1}(\mathbf{x}, t) \mathbf{f}(\mathbf{x}, t) $$
# 
# and that $\mathbf{g}$ can be computed analytically but for non-toy problems, it is best to do this numerically. So we will simply generate the $\mathbf{M}$ and $\mathbf{f}$ matrices for later use.
# 
# The mass matrix, $\mathbf{M}$, can be accessed with the `mass_matrix` method (use `mass_matrix_full` to include the kinematical differential equations too. We can use `trigsimp` again to make this relatively compact: 

# In[ ]:


mass_matrix = trigsimp(kane.mass_matrix_full)
mass_matrix


# The right hand side, $\mathbf{f}$, is a vector function of all the non-inertial forces (gyroscopic, external, coriolis, etc):

# In[ ]:


forcing_vector = trigsimp(kane.forcing_full)
forcing_vector