%typeset_mode True
import numpy as np
import scipy
import matplotlib.pyplot as plt
import numpy.fft as fft
import matplotlib
import matplotlib.pyplot as plt1
import matplotlib.pyplot as plt2
from scipy import integrate, optimize, interpolate
from scipy.special import gamma
%var angle1, x, x_dot,x1dot, x2dot, epsilon, A, Aa, alpha, alpha_dot, m1, m2, k1,g, l1, Ek, Ep, Lagran1, Lagran2
#х-путь, alpha - угол отклонения шара, m1-масса тележки, m2-масса шара, k1-жесткость пружины,g-9.81, l1-длинна нити, за которую подвешен шар
Ek = m2*((x_dot^2)+(2*x_dot*l1*cos(alpha)*alpha_dot)+(l1^2)*(alpha_dot^2)*((cos(alpha)**2)-(sin(alpha)**2)))/2
Ep = m2*g*(l1-l1*cos(alpha))+(k1*(x^2)/2)
L = Ek - Ep
Ldiff = L.diff(x_dot )
Ldiff2 = L.diff(alpha_dot)
t=var('t')
xdot_tmp = var ('X_Dot')
alphadot_tmp = var ('Alpha_Dot')
alpha_tmp = function ('Alpha')(t)# For substitution alpha
x_tmp = function ('X')(t) #For substitution x
angle = var ('theta')
Ldiff3 = Ldiff.subs([x_dot==x_tmp.diff(t), x==x_tmp ,alpha==alpha_tmp,alpha_dot==alpha_tmp.diff(t)]).diff(t)
Ldiff4 = Ldiff2.subs([x_dot==x_tmp.diff(t), x==x_tmp ,alpha==alpha_tmp,alpha_dot==alpha_tmp.diff(t)]).diff(t)
Lagran1 = Ldiff3-L.diff(x)-A*sin(angle*t)
Lagran2= Ldiff4-L.diff(alpha)
print ('Lagran1')
Lagran1
print ('Lagran2')
Lagran2
omega =var ('omega') #For subs second derivative alpha. It is pendulum acceleration
Aa =var('AA') #For subs second derivative x. It is trolley acceleration
eq1=Lagran1.subs([x_tmp.diff(t)==x_dot,x_tmp==x,alpha_tmp.diff(t)==alpha_dot,alpha_tmp==alpha, alpha_tmp.diff(t).diff(t)==omega, x_tmp.diff(t).diff(t)==Aa])
print ('eq1')
eq1
eq2=Lagran2.subs([x_tmp.diff(t)==x_dot,x_tmp==x,alpha_tmp.diff(t)==alpha_dot,alpha_tmp==alpha, alpha_tmp.diff(t).diff(t)==omega, x_tmp.diff(t).diff(t)==Aa])
print ('eq2')
eq2
sol_list=solve([ eq1,eq2  ],[Aa,omega ])
a111=sol_list[0][0].rhs()
a112=sol_list[0][1].rhs()
a111=a111.subs(m1==100,m2==10,l1==5,g==9.81,k1==2,A==0,angle==30)
print ('a111')
a111
a112=a112.subs(m1==100,m2==10,l1==5,g==9.81,k1==2,A==0,angle==30)
print ('a112')
a112
def fun_ode(y,t_num):
    x_num       =y[0]
    v_num       =y[1]
    alpha_num   =y[2]
    alphadot_num=y[3]
    xdot1       =v_num
    xdot2       =alphadot_num #alpha_dot
    v1dot=a111.subs(x==x_num,alpha_dot==alphadot_num,alpha==alpha_num,t==t_num)
    v2dot=a112.subs(x==x_num,alpha_dot==alphadot_num,alpha==alpha_num,t==t_num)
    return[xdot1,v1dot,xdot2,v2dot] # xdot1 -скорость тележки, v1dot-ускорение тележки, xdot2-скорость шара, v2dot-ускорение шара. 
t_span = np.arange(0,20.0, 0.0001)
x_0 =0.0001
v_0 =0.0001
x1_0=0.0001
v1_0=0.0001
y_0 = [x_0,v_0,x1_0,v1_0]
dt = t_span[1]-t_span[0]; dt
y = integrate.odeint(fun_ode, y_0, t_span,  rtol=1e-12, atol=1e-12)
bp_x = t_span
bp_y = y [:,0]
plt1.figure(figsize=(15,5))
probes_no = len(y)
rms_spec = abs(fft.fft(bp_y))
rms_spec_shift = fft.fftshift(rms_spec)
f_span=fft.fftshift(fft.fftfreq(probes_no, d=dt))
plt1.plot(f_span,rms_spec_shift, linewidth=2, linestyle="-", color="blue")
plt2.plot(f_span,rms_spec.imag,color="red")
plt1.grid()
plt1.xlim([-5,5])
plt1.show()
plt2.figure(figsize=(15,5))
plt2.plot(bp_x,bp_y)
plt2.grid()
plt2.show()

 v1dot=a111.subs(x==x_num,alpha_dot==alphadot_num,alpha==alpha_num,t==t_num)