import numpy as np
import sympy as smp
from sympy import init_printing
from scipy.integrate import odeint
import matplotlib.pyplot as plt
from matplotlib import animation
from mpl_toolkits.mplot3d import Axes3D
from matplotlib.animation import PillowWriter
init_printing()
t, g = smp.symbols('t g')
m1, m2 = smp.symbols('m1 m2')
L_c1, L_c2 = smp.symbols('L_{c1}, L_{c2}')
L_1 = smp.symbols('L_1')
the1, the2 = smp.symbols(r'\theta_1, \theta_2', cls=smp.Function)
tau_2 = smp.symbols('tau_2')
the1 = the1(t)
the2 = the2(t)
the1_d = smp.diff(the1, t)
the2_d = smp.diff(the2, t)
the1_dd = smp.diff(the1_d, t)
the2_dd = smp.diff(the2_d, t)
x1 = L_c1*smp.sin(the1)
z1 = L_c1*smp.cos(the1)
x2 = L_1*smp.sin(the1)-L_c2*smp.sin(the1+the2)
z2 = L_1*smp.cos(the1)-L_c2*smp.cos(the1+the2)
# Kinetic
T1 = 1/2 * m1 * (smp.diff(x1, t)**2 + smp.diff(z1, t)**2)
T2 = 1/2 * m2 * (smp.diff(x2, t)**2 + smp.diff(z2, t)**2)
T = T1+T2
# Potential
V1 = m1*g*z1
V2 = m2*g*z2
V = V1 + V2
# Lagrangian
L = T-V
L
LE1 = smp.diff(smp.diff(L, the1_d), t).simplify() - smp.diff(L, the1)
LE2 = smp.diff(smp.diff(L, the2_d), t).simplify() - smp.diff(L, the2)
eq1 = smp.Eq(LE1,0)
eq2 = smp.Eq(LE2,tau_2)
eq1.simplify()
eq2.simplify()