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')

Ao.set_pos(O, l/2*A.x)
Bo.set_pos(O, l*A.x)

O.set_vel(N, 0)
Ao.v2pt_theory(O, N, A)
Bo.v2pt_theory(O, N, A)

Ao.vel(N), Bo.vel(N), A.ang_vel_in(N), B.ang_vel_in(N)

Q = me.Point('Q')

Q.set_pos(Bo, q3*B.y)

Q.set_vel(B, u3*B.y)

Q.v1pt_theory(Bo, N, B)

Ao.acc(N)

Bo.acc(N)

Q.acc(N)

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)

A.ang_acc_in(N)

B.ang_acc_in(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

me.find_dynamicsymbols(Ts_A, reference_frame=N)

me.find_dynamicsymbols(Ts_B, reference_frame=N)

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),
])
Fr

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),
])
Frs

kanes_eqs = Fr + Frs
kanes_eqs

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

Md = Frs.jacobian(ud)
Md

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

gd = Frs.xreplace(ud_zero_repl) + Fr
gd

res = sm.cse(Md.LUsolve(-gd))

res[0]

res[1]