import sympy as sm
import sympy.physics.mechanics as me
me.init_vprinting()

m, g, kt, kl, l = sm.symbols('m, g, k_t, k_l, l')
q1, q2, q3 = me.dynamicsymbols('q1, q2, q3')
u1, u2, u3 = me.dynamicsymbols('u1, u2, u3')

N = me.ReferenceFrame('N')
A = me.ReferenceFrame('A')
B = me.ReferenceFrame('B')

A.orient_axis(N, q1, N.z)
B.orient_axis(A, q2, A.x)

A.set_ang_vel(N, u1*N.z)
B.set_ang_vel(A, u2*A.x)

O = me.Point('O')
Ao = me.Point('A_O')
Bo = me.Point('B_O')
Q = me.Point('Q')

Ao.set_pos(O, l/2*A.x)
Bo.set_pos(O, l*A.x)
Q.set_pos(Bo, q3*B.y)

O.set_vel(N, 0)
Ao.v2pt_theory(O, N, A)
Bo.v2pt_theory(O, N, A)
Q.set_vel(B, u3*B.y)
Q.v1pt_theory(Bo, N, B)

qd_repl = {q1.diff(): u1, q2.diff(): u2, q3.diff(): u3}
qd_repl

Q.set_acc(N, Q.acc(N).xreplace(qd_repl))
Q.acc(N)

v_Ao_1 = Ao.vel(N).diff(u1, N)
v_Ao_2 = Ao.vel(N).diff(u2, N)
v_Ao_3 = Ao.vel(N).diff(u3, N)

v_Ao_1, v_Ao_2, v_Ao_3

v_Bo_1 = Bo.vel(N).diff(u1, N)
v_Bo_2 = Bo.vel(N).diff(u2, N)
v_Bo_3 = Bo.vel(N).diff(u3, N)

v_Bo_1, v_Bo_2, v_Bo_3

v_Q_1 = Q.vel(N).diff(u1, N)
v_Q_2 = Q.vel(N).diff(u2, N)
v_Q_3 = Q.vel(N).diff(u3, N)

v_Q_1, v_Q_2, v_Q_3

w_A_1 = A.ang_vel_in(N).diff(u1, N)
w_A_2 = A.ang_vel_in(N).diff(u2, N)
w_A_3 = A.ang_vel_in(N).diff(u3, N)

w_A_1, w_A_2, w_A_3

w_B_1 = B.ang_vel_in(N).diff(u1, N)
w_B_2 = B.ang_vel_in(N).diff(u2, N)
w_B_3 = B.ang_vel_in(N).diff(u3, N)

w_B_1, w_B_2, w_B_3

I = m*l**2/12
I_A_Ao = I*me.outer(A.y, A.y) + I*me.outer(A.z, A.z)
I_B_Bo = I*me.outer(B.x, B.x) + I*me.outer(B.z, B.z)

I_A_Ao, I_B_Bo

F = me.dynamicsymbols('F')

R_Ao = m*g*N.x
R_Bo = m*g*N.x - F*B.y
R_Q = m/4*g*N.x + F*B.y

R_Ao, R_Bo, R_Q

T_A = -kt*q1*N.z + kt*q2*B.x
T_B = -kt*q2*B.x

T_A, T_B

Rs_Ao = -m*Ao.acc(N)
Rs_Bo = -m*Bo.acc(N)
Rs_Q = -m/4*Q.acc(N)

Ts_A = -(me.dot(A.ang_acc_in(N), I_A_Ao) + 
         me.dot(me.cross(A.ang_vel_in(N), I_A_Ao), A.ang_vel_in(N)))
Ts_A

Ts_B = -(me.dot(B.ang_acc_in(N), I_B_Bo) + 
         me.dot(me.cross(B.ang_vel_in(N), I_B_Bo), B.ang_vel_in(N)))
Ts_B

Fr = sm.Matrix([
    v_Ao_1.dot(R_Ao) + v_Bo_1.dot(R_Bo) + v_Q_1.dot(R_Q) +
    w_A_1.dot(T_A) + w_B_1.dot(T_B),
    v_Ao_2.dot(R_Ao) + v_Bo_2.dot(R_Bo) + v_Q_2.dot(R_Q) +
    w_A_2.dot(T_A) + w_B_2.dot(T_B),
    v_Ao_3.dot(R_Ao) + v_Bo_3.dot(R_Bo) + v_Q_3.dot(R_Q) +
    w_A_3.dot(T_A) + w_B_3.dot(T_B),
])

Frs = sm.Matrix([
    v_Ao_1.dot(Rs_Ao) + v_Bo_1.dot(Rs_Bo) + v_Q_1.dot(Rs_Q) +
    w_A_1.dot(Ts_A) + w_B_1.dot(Ts_B),
    v_Ao_2.dot(Rs_Ao) + v_Bo_2.dot(Rs_Bo) + v_Q_2.dot(Rs_Q) +
    w_A_2.dot(Ts_A) + w_B_2.dot(Ts_B),
    v_Ao_3.dot(Rs_Ao) + v_Bo_3.dot(Rs_Bo) + v_Q_3.dot(Rs_Q) +
    w_A_3.dot(Ts_A) + w_B_3.dot(Ts_B),
])

kanes_eqs = Fr + Frs
kanes_eqs

u = sm.Matrix([u1, u2, u3])
t = me.dynamicsymbols._t
ud = u.diff(t)

Md = Frs.jacobian(ud)
Md

ud_zero_repl = {u1.diff(): 0, u2.diff(): 0, u3.diff(): 0}

gd = Frs.xreplace(ud_zero_repl) + Fr
gd

me.find_dynamicsymbols(Md)

me.find_dynamicsymbols(gd)

Md.free_symbols | gd.free_symbols

q = sm.Matrix([q1, q2, q3])
q

u

p = sm.Matrix([g, kt, l, m])
p

# specified inputs
r = sm.Matrix([F])
r

eval_dyn = sm.lambdify((q, u, p, r), (Md, gd))

import numpy as np

q_vals = np.array([
    np.deg2rad(25.0),  # q1, rad
    np.deg2rad(5.0),  # q2, rad
    0.1,  # q3, m
])

u_vals = np.array([
    0.1,  # u1, rad/s
    2.2,  # u2, rad/s
    0.3,  # u3, m/s
])

p_vals = np.array([
    9.81,  # g, m/s**2
    0.01,  # kt, Nm/rad
    0.6,  # l, m
    1.0,  # m, kg
])

r_vals = np.array([100.0])  # Newtons

Md_vals, gd_vals = eval_dyn(q_vals, u_vals, p_vals, r_vals)

Md_vals

gd_vals

np.linalg.solve?

ud_vals = np.linalg.solve(-Md_vals, np.squeeze(gd_vals))
ud_vals

qd_vals = u_vals
qd_vals

xd_vals = np.hstack((qd_vals, ud_vals))
xd_vals

def eval_rhs(t, x, p, r):
    q = x[:3]
    u = x[3:]
    Md, gd = eval_dyn(q, u, p, r)
    ud = np.linalg.solve(-Md, np.squeeze(gd))
    qd = u
    return np.hstack((qd, ud))

eval_rhs(2.1, np.hstack((q_vals, u_vals)), p_vals, r_vals)

def eval_rhs(t, x, p):
    
    q = x[:3]
    u = x[3:]
    #r = np.array([1.0*np.sin(2*np.pi*t)])
    r = np.array([-2.0*q[2]])
    Md, gd = eval_dyn(q, u, p, r)
    ud = np.linalg.solve(-Md, np.squeeze(gd))
    qd = u
    return np.hstack((qd, ud))

eval_rhs(2.1, np.hstack((q_vals, u_vals)), p_vals)

from scipy.integrate import solve_ivp

x0 = np.hstack((q_vals, u_vals))
t0 = 0.0
tf = 4.0

sol = solve_ivp(eval_rhs, (t0, tf), x0, args=(p_vals, ))

sol.t

sol.y

import matplotlib.pyplot as plt

sol.t.shape

sol.y.shape

%matplotlib notebook

plt.plot(sol.t, np.transpose(sol.y))

plt.xlabel('Time [s]')

plt.legend(['q1', 'q2', 'q3', 'u1', 'u2', 'u3'])

def plot_results(ts, xs):
    """Returns the array of axes of a 4 panel plot of the state trajectory
    versus time.

    Parameters
    ==========
    ts : array_like, shape(m,)
       Values of time.
    xs : array_like, shape(m, 6)
       Values of the state trajectories corresponding to ``ts`` in order
       [q1, q2, q3, u1, u2, u3].

    Returns
    =======
    axes : ndarray, shape(4,)
       Matplotlib axes for each panel.

    """

    fig, axes = plt.subplots(4, 1, sharex=True)

    fig.set_size_inches((10.0, 6.0))

    axes[0].plot(ts, np.rad2deg(xs[:, :2]), '.-')
    axes[1].plot(ts, xs[:, 2], '.-')
    axes[2].plot(ts, np.rad2deg(xs[:, 3:5]), '.-')
    axes[3].plot(ts, xs[:, 5], '.-')

    axes[0].legend([me.vlatex(q[0], mode='inline'),
                    me.vlatex(q[1], mode='inline')])
    axes[1].legend([me.vlatex(q[2], mode='inline')])
    axes[2].legend([me.vlatex(u[0], mode='inline'),
                    me.vlatex(u[1], mode='inline')])
    axes[3].legend([me.vlatex(u[2], mode='inline')])

    axes[0].set_ylabel('Angle [deg]')
    axes[1].set_ylabel('Distance [m]')
    axes[2].set_ylabel('Angular Rate [deg/s]')
    axes[3].set_ylabel('Speed [m/s]')

    axes[3].set_xlabel('Time [s]')

    fig.tight_layout()

    return axes

plot_results(sol.t, np.transpose(sol.y))

x0 = np.hstack((q_vals, u_vals))
t0 = 0.0
tf = 4.0

ts = np.linspace(t0, tf, num=100)

sol = solve_ivp(eval_rhs, (t0, tf), x0, args=(p_vals, ), t_eval=ts)

plot_results(sol.t, np.transpose(sol.y))

M = me.ReferenceFrame('M')
M.orient_axis(N, sm.pi/2, N.z)

Bl = me.Point('B_l')
Br = me.Point('B_r')
Bl.set_pos(Bo, -l/2*B.y)
Br.set_pos(Bo, l/2*B.y)

coordinates = O.pos_from(O).to_matrix(M)
for point in [Bo, Q, Bl, Br]:
    coordinates = coordinates.row_join(point.pos_from(O).to_matrix(M))

eval_point_coords = sm.lambdify((q, p), coordinates)
eval_point_coords(q_vals, p_vals)

# initial configuration of the points
x, y, z = eval_point_coords(q_vals, p_vals)

# create a figure
fig = plt.figure()
fig.set_size_inches((10.0, 10.0))

# setup the subplots
ax_top = fig.add_subplot(2, 2, 1)
ax_3d = fig.add_subplot(2, 2, 2, projection='3d')
ax_front = fig.add_subplot(2, 2, 3)
ax_right = fig.add_subplot(2, 2, 4)

# common line and marker properties for each panel
line_prop = {
    'color': 'black',
    'marker': 'o',
    'markerfacecolor': 'blue',
    'markersize': 10,
}

# top view
lines_top, = ax_top.plot(x, z, **line_prop)
ax_top.set_xlim((-0.5, 0.5))
ax_top.set_ylim((0.5, -0.5))
ax_top.set_title('Top View')
ax_top.set_xlabel('x')
ax_top.set_ylabel('z')
ax_top.set_aspect('equal')

# 3d view
lines_3d, = ax_3d.plot(x, z, y, **line_prop)
ax_3d.set_xlim((-0.5, 0.5))
ax_3d.set_ylim((0.5, -0.5))
ax_3d.set_zlim((-0.8, 0.2))
ax_3d.set_xlabel('x')
ax_3d.set_ylabel('z')
ax_3d.set_zlabel('y')

# front view
lines_front, = ax_front.plot(x, y, **line_prop)
ax_front.set_xlim((-0.5, 0.5))
ax_front.set_ylim((-0.8, 0.2))
ax_front.set_title('Front View')
ax_front.set_xlabel('x')
ax_front.set_ylabel('y')
ax_front.set_aspect('equal')

# right view
lines_right, = ax_right.plot(z, y, **line_prop)
ax_right.set_xlim((0.5, -0.5))
ax_right.set_ylim((-0.8, 0.2))
ax_right.set_title('Right View')
ax_right.set_xlabel('z')
ax_right.set_ylabel('y')
ax_right.set_aspect('equal')

fig.tight_layout()

from matplotlib.animation import FuncAnimation

xs = np.transpose(sol.y)
def animate(i):
    x, y, z = eval_point_coords(xs[i, :3], p_vals)
    lines_top.set_data(x, z)
    lines_3d.set_data_3d(x, z, y)
    lines_front.set_data(x, y)
    lines_right.set_data(z, y)

animate(20)

ani = FuncAnimation(fig, animate, len(sol.t))

from IPython.display import HTML

HTML(ani.to_html5_video())

HTML(ani.to_jshtml(fps=30))