############ Humanoid Robot ########

from numpy import array, zeros, eye, asarray, dot, rad2deg
from numpy.linalg import inv
from matplotlib.pyplot import plot, xlabel, ylabel, legend, rcParams
rcParams['figure.figsize'] = (14, 8)
from sympy import simplify, Matrix, matrix2numpy
from sympy.physics.vector import init_vprinting, vlatex
init_vprinting(use_latex='mathjax')
from Utilities.Controller import controllable
from scipy.linalg import solve_continuous_are
from scipy.integrate import odeint

A = np.array([[   0.   ,    0.   ,    0.   ,    1.   ,    0.   ,    0.   ],
               [   0.   ,    0.   ,    0.   ,    0.   ,    1.   ,    0.   ],
               [   0.   ,    0.   ,    0.   ,    0.   ,    0.   ,    1.   ],
               [  18.337,  -75.864,    6.395,    0.   ,    0.   ,    0.   ],
               [ -22.175,  230.549,  -49.01 ,    0.   ,    0.   ,    0.   ],
               [   4.353, -175.393,   95.29 ,    0.   ,    0.   ,    0.   ]])

B = np.array([[ 0.   ,  0.   ,  0.   ],
               [ 0.   ,  0.   ,  0.   ],
               [ 0.   ,  0.   ,  0.   ],
               [ 0.292, -0.785,  0.558],

from sympy import symbols, simplify, lambdify, solve
from sympy.physics.mechanics import dynamicsymbols, ReferenceFrame, Point, inertia, RigidBody, KanesMethod
from pydy.codegen.code import generate_ode_function
from numpy import array, linspace, deg2rad, rad2deg, ones, concatenate, sin, cos, pi, zeros, dot, eye
from numpy.linalg import inv
from scipy.integrate import odeint
from scipy.linalg import solve_continuous_are
import matplotlib.animation as animation
import matplotlib.pyplot as plt
from matplotlib.pyplot import plot, xlabel, ylabel, legend, rcParams
import numpy as np
from body_model_setup import theta1, theta2, theta3,theta4, omega1, omega2, omega3,omega4, l_ankle_torque, l_hip_torque, waist_torque, r_hip_torque, coordinates, speeds, kane, mass_matrix, forcing_vector, specified, parameter_dict, constants, numerical_constants, l_leg, r_leg, crotch, body
from sympy.physics.vector import init_vprinting, vlatex
from math import fabs
init_vprinting()
import pickle


# Control
# =======

inputK = open('bm_LQR_K_useful.pkl','rb')
inputa1 = open('bm_angle_one_useful.pkl','rb')
inputa2 = open('bm_angle_two_useful.pkl','rb')
inputa3 = open('bm_angle_three_useful.pkl','rb')
inputa4 = open('bm_angle_four_useful.pkl','rb')

K = pickle.load(inputK)
a1 = pickle.load(inputa1)
a1 = np.asarray(a1, dtype = float)
a2 = pickle.load(inputa2)

#%% 
from __future__ import print_function, division
from sympy import symbols, simplify, Matrix
from sympy import trigsimp
from sympy.physics.mechanics import dynamicsymbols, ReferenceFrame, Point, Particle, inertia, RigidBody, KanesMethod
from numpy import deg2rad, rad2deg, array, zeros, linspace
from sympy.physics.vector import init_vprinting, vlatex
import numpy as np

from sympy.physics.vector import init_vprinting
init_vprinting(use_latex='mathjax', pretty_print=False)
#%% 
# define all reference frames of all individually moving links of the robot
print("Defining reference frames")
inertial_frame = ReferenceFrame('I')
lower_leg_left_frame = ReferenceFrame('R_1')
upper_leg_left_frame = ReferenceFrame('R_2')
hip_frame = ReferenceFrame('R_3')
upper_leg_right_frame = ReferenceFrame('R_4')
lower_leg_right_frame = ReferenceFrame('R_5')
#%% Angles
# everything is symbolic, so create all angles of your robot
# NOTE: the angle phi is the angle between your robot and the inertial frame
theta0, theta1, theta2, theta3, phi = dynamicsymbols('theta0, theta1, theta2, theta3, phi')

lower_leg_left_frame.orient(inertial_frame, 'Axis', (phi, inertial_frame.z))
simplify(lower_leg_left_frame.dcm(inertial_frame))

upper_leg_left_frame.orient(lower_leg_left_frame, 'Axis', (theta0, -lower_leg_left_frame.z))
simplify(upper_leg_left_frame.dcm(inertial_frame))