from sympy import *
from sympy.solvers.solveset import linsolve
from sympy import Matrix, S
from sympy import symbols
from scipy.integrate import odeint
import numpy as np

init_printing()

n = 3

# arrays for all the constants in the system
l = []
m = []
k = []
g=9.81

# fill in the constant values
for i in range(n):
    l.append(1)
    m.append(1)
    k.append(0.1)

    

# arrays for all the symbols in the system
phi = []       #angles
dphi = []      #first derivatives
ddphi = []     #second derivatives
diff2 = []     #list of tuples that map expressions phi.diff(t).diff(t) to symbols ddphi
diff1 = []     #see above for first derivatives
x = []         #x positions of pendulums masses
y = []         #y positions
Fr=[]          #Friction terms
t = symbols("t")


# Lagrangian, an empty expression
L = Add(0, 0)

for i in range(n):
    temp_phi = symbols("phi_" + str(i))
    temp_phi = temp_phi(t)
    phi.append(temp_phi)

    temp_dphi = symbols("dphi_" + str(i))
    dphi.append(temp_dphi)

    temp_ddphi = symbols("ddphi_" + str(i))
    ddphi.append(temp_ddphi)

    diff1.append((temp_phi.diff(t), temp_dphi))
    diff2.append((temp_phi.diff(t).diff(t), temp_ddphi))

    temp_x = symbols("x_" + str(i))
    temp_y = symbols("y_" + str(i))

    temp_x = l[i] * sin(phi[-1])
    temp_y = -l[i] * cos(phi[-1])

    if i > 0:
        temp_x = temp_x + x[-1]
        temp_y = temp_y + y[-1]

    x.append(temp_x)
    y.append(temp_y)

    temp_v = temp_x.diff(t) ** 2 + temp_y.diff(t) ** 2

    Fr.append(-temp_dphi*k[i])

    # adding the kinetic energy of this pendulum to L
    L += 1 / 2 * temp_v * m[i]

    # adding the potential energy of this pendulum to L
    L -= m[i] * g * temp_y

    

eqs = []

for i in range(n):
    a = L.diff(phi[i].diff(t))
    da = a.diff(t)
    b = L.diff(phi[i]) +Fr[i]
    eq = da - b
    eq=eq.subs(diff2).subs(diff1).simplify()
    eqs.append(eq)

# this can be solved for the second order derivatives
import time
start=time.time()
res = linsolve(eqs, ddphi)
end=time.time()
print(end-start)

res = [res for res in res]
res = res[0]

lambdified_solns = [lambdify(phi + dphi, soln) for soln in res]


# state
#  phi_1
#  phi_1'
#  phi_2
#  phi_2'

# phi
# dphi
def derivs(state, t):
    d=np.empty(state.shape)
    d[:-n] = state[n:]
    for i in range(n):
        d[i+n]=lambdified_solns[i](*state)
    return d


# solve ode
start_phi=[]
for i in range(n):
    start_phi.append(3)

start_dphi=[]
for i in range(n):
    start_dphi.append(0.5)

steps = 5000
t_end = 300
res_ode = odeint(derivs, start_phi + start_dphi, np.linspace(0, t_end, steps))