import numpy as np
from scipy import integrate
import math
import matplotlib.pyplot as plt
g = 9.81
m = 1
M = 10
l = 10


def u(t):
    return 10


def equation(y0, t):
    r, theta, q, w = y0
    return [q, w, r * w ** 2 - g * math.sin(theta),
             (u(t) - 2 * m * r * q * w - m * g * r * math.cos(theta)) / (2 * M * l ** 2 + m * r ** 2)]


y0 = [1, 0, 0, 0]
t = np.linspace(0, 10, 10000)
sol = integrate.odeint(equation, y0, t)
plt.plot(t, sol[:, 0], 'b', label='r(t)')
plt.plot(t, sol[:, 1], 'g', label='theta(t)')
plt.plot(t, sol[:, 2], 'red', label='r_pr(t)')
plt.plot(t, sol[:, 3], 'yellow', label='theta_pr(t)')
plt.legend(loc='best')
plt.xlabel('t')
plt.grid()
plt.show()