import numpy as np
import matplotlib.pyplot as plt
from sympy import *
import sympy.physics.mechanics as me
from pydy.system import System
from pydy.viz import Cylinder, Plane, VisualizationFrame, Scene
from sympy import sin, cos, symbols, solve
from pydy.codegen.ode_function_generators import generate_ode_function
from scipy.integrate import odeint
from IPython.display import SVG
SVG(filename='openrov2.svg')

%matplotlib nbagg

me.init_vprinting(use_latex='mathjax')

"""
Explanation: Setup
End of explanation
"""


# Inertial Reference Frame
N = me.ReferenceFrame('N')

# Define a world coordinate origin
O = me.Point('O')
O.set_vel(N, 0)

rot = list(me.dynamicsymbols('r0:3'))
#rot = list(symbols('r0:3'))
drot = list(me.dynamicsymbols('dr0:3'))
x = list(me.dynamicsymbols('v0:3')) # Coordinates of robot in World Frame
dx = list(me.dynamicsymbols('dv0:3'))
kin_diff=Matrix(x+rot).diff()-Matrix(dx+drot)
kin_diff

#xxx = me.dynamicsymbols('xxx')
#from sympy import *
#eval(sm.srepr(kin_diff))
#import pickle
#pickle.dump(xxx,open('/tmp/forcing_vector.pkl','wb'))

# Constants for the Robot Body
Wx = symbols('W_x')                                     # 2*w is the width between thrusters
Wh = symbols('W_h')                                    
T1 = symbols('T_1')                                     # Distance between thruster base and center of mass
T2 = symbols('T_2')                                     
Bh = symbols('B_h')                                     
Bw = symbols('B_w')                                     

m_b = symbols('m_b')                                 # Mass of the body
v_b = symbols('v_b') # Volume of the body
mu = symbols('\mu') #drag
mu_r = symbols('\mu_r') #rotational drag
g = symbols('g')
I = list(symbols('Ixx, Iyy, Izz'))             # Moments of inertia of body


# Robot Reference Frame
Rz=N.orientnew('R_z', 'Axis', (rot[2], N.z))
Rz.set_ang_vel(N,drot[2]*N.z)

Rx=Rz.orientnew('R_x', 'Axis', (rot[0], Rz.x))
Rx.set_ang_vel(Rz,drot[0]*Rz.x)

R=Rx.orientnew('R', 'Axis', (rot[1], Rx.y))
R.set_ang_vel(Rx,drot[1]*Rx.y)

#### adding dumping Torqe for each rotation
T_z=(R,-drot[2]*N.z*mu_r) #rotaional dumping Torqe
T_x=(R,-drot[0]*Rz.x*mu_r) #rotaional dumping Torqe
T_y=(R,-drot[1]*Rx.y*mu_r) #rotaional dumping Torqe

# Center of mass of body
COM = O.locatenew('COM', x[0]*N.x + x[1]*N.y + x[2]*N.z)

# Set the velocity of COM
COM.set_vel(N, dx[0]*N.x + dx[1]*N.y + dx[2]*N.z)

#center of bouyency
COB = COM.locatenew('COB', R.x*Bw+R.z*Bh)
COB.v2pt_theory(COM, N, R);
R.ang_vel_in(N)

# Calculate inertia of body
Ib = me.inertia(R, *I)
# Create a rigid body object for body
Body = me.RigidBody('Body', COM, R, m_b, (Ib, COM))


# Points of thrusters
L1 = COM.locatenew('L_1', -R.x*T1-Wx*R.y)
L2 = COM.locatenew('L_2', -R.x*T1+Wx*R.y)
L3 = COM.locatenew('L_3', -R.x*T2-Wh*R.z)

# Set the velocity of points 
L1.v2pt_theory(COM, N, R)
L2.v2pt_theory(COM, N, R)
L3.v2pt_theory(COM, N, R);


"""
Explanation: Define Variables
constants:

$m_A$: Mass of the Rov.
$T_1$: Distance from COM to back thruster axis.
$T_2$: Distance from COM to upper thruster axis.
$W_x$: distance from center axis to back thruster axis.
$W_h$: distance from center axis to upper thruster.
$B_h$: the z messure of COB (center of bouyency) distance from COG (center of gravity) 
$B_h$: the x messure of COB distance from COG.
End of explanation
"""


#dCw=Cw.diff()
v=N.x*dx[0]+N.y*dx[1]+N.z*dx[2]
Fd=-v.normalize()*v.magnitude()**2*mu
Fd

#rotational drags
#Fr=-R.ang_vel_in(R_static)*mu
#Fr=(N,-drot[2]**2*N.z*mu),-drot[0]**2*Rz.x-drot[1]**2*Rx.y
#Fr=Fr*mu
#Fr

#thrust forces symbols
F1, F2, F3 = symbols('f_1, f_2, f_3') 
Fg = -N.z *m_b * g
Fb = N.z * v_b * 1e3 *g #whight of 1m^3 water in kg (MKS units)

kane = me.KanesMethod(N, q_ind=x+rot, u_ind=dx+drot, kd_eqs=kin_diff)

bodies = (Body,)
loads = (
    (L1, F1 * R.x),
    (L2, F2 * R.x),
    (L3, F3 * R.z), 
    (COM, Fg ), 
    (COB, Fb ), 
    (COM, Fd ),
    T_x,
    T_y,
    T_z
    )

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

mu_r

mass_matrix = trigsimp(kane.mass_matrix_full)
mass_matrix

forcing_vector = trigsimp(kane.forcing_full)

#open('/tmp/mass_matrix.srepr','wb').write(mass_matrix)


coordinates = tuple(x+rot)
coordinates

speeds = tuple(dx+drot)
speeds

specified = (F1, F2, F3)
constants = [Wx,Wh,T1,T2,Bh,Bw,m_b,v_b,mu,mu_r,g]+I

if 1:
    open('./forcing_vector.srepr','wb').write(srepr((\
                                                    forcing_vector,
                                                   coordinates,
                                                    mass_matrix,
                                                   speeds,
                                                    constants,
                                                   specified,
                                                   )).encode())


if 1:
    (forcing_vector,\
    coordinates,
    mass_matrix,
    speeds,
    constants,
    specified)=eval(open('./forcing_vector.srepr','rb').read())


right_hand_side = generate_ode_function(forcing_vector, coordinates,
                                        speeds, constants,
                                        mass_matrix=mass_matrix,
                                        specifieds=specified)

help(right_hand_side)

x0 = np.zeros(12)
#MKS units
#constants = [Wx,Wh,T1,T2,Bh,Bw,m_b,v_b,mu,g]+I
numerical_constants = np.array([
                            0.1,  # Wx [m]
                            0.15,  # Wh [m]
                            0.1,  # T1 [m]
                            0.05,  # T2 [m]
                            0.08,  # Bh [m]
                            0.01,  # Bw [m]
                            1.0 *3,  # m_b [kg]
                            0.001 *3 ,  # v_b [M^3]
                            5.9,  # mu
                            0.2,  # mu_r
                            9.8,  # g MKS
                            0.5,  # Ixx [kg*m^2]
                            0.5,  # Iyy [kg*m^2]
                            0.5,  # Izz [kg*m^2]
                               ]
                            ) 

#args = {'constants': numerical_constants,
numerical_specified=[0.8,0.5,0]
frames_per_sec = 60.0
final_time = 60.0

t = np.linspace(0.0, final_time, int(final_time * frames_per_sec))

right_hand_side(x0, 0.0, numerical_specified, numerical_constants)

def controller(x, t):
    if t<20 or t>35:
        #return [0.8,0.5,0]
        return [0.0,0.0,0]
    else:
        return [-0.55*30,-0.5*30,0]
#def controller(x, t):
#    return [0.0,0.0,0]
y = odeint(right_hand_side, x0, t, args=(controller, numerical_constants))

#y = odeint(right_hand_side, x0, t, args=(numerical_specified, numerical_constants))
y.shape

def plot():
    plt.figure()
    #plt.plot(sys.times, np.rad2deg(x[:, :3]))
    plt.subplot(2,3,1)
    #plt.plot(t, np.rad2deg(y[:, 0]))
    plt.plot(t, y[:, :3])
    plt.legend([latex(s, mode='inline') for s in coordinates[:3]])
    plt.subplot(2,3,2)
    plt.plot(t, np.rad2deg(y[:, 3:6]))
    plt.legend([latex(s, mode='inline') for s in coordinates[3:6]])
    plt.subplot(2,3,3)
    plt.title('XY')
    plt.plot(y[:,0],y[:,1])
    plt.axis('equal')
    plt.subplot(2,3,6)
    plt.title('Z')
    plt.plot(y[:,2])
    plt.axis('equal')
    plt.subplot(2,3,4)
    plt.plot(t, np.rad2deg(y[:, 9:12]))
    plt.legend([latex(s, mode='inline') for s in coordinates[9:12]])
plot()

"""
Explanation: Calculating hydrodynamic drag
under sphire assumption and ignoring inertia forces
$F_{D}\,=\,{\tfrac {1}{2}}\,\rho \,v^{2}\,C_{D}\,A$
https://en.wikipedia.org/wiki/Drag_(physics)
we define $\mu$ as:
$\mu=\,{\tfrac {1}{2}}\,\rho \,C_{D}\,A$
then:
$F_{D}\,=\mu \,v^{2}$
End of explanation
"""


x0 = np.zeros(12)
xx=x0
y=[]
for ct in t:
    x_dot=right_hand_side(xx, ct, controller, numerical_constants)
    y.append(xx)
    xx=xx+x_dot*1/frames_per_sec
y=np.array(y)
y.shape

plot()

"""
Explanation: another method
End of explanation
"""


