from __future__ import division
import numpy as np
from numpy import sin, cos, pi
from scipy.integrate import odeint
import matplotlib.pyplot as plt
from matplotlib import rcParams

rcParams['font.size'] = 14
rcParams['legend.handlelength'] = 0

def pend(x, t, grav=9.81, length=1, amplitude=0.05, omega=10):
    theta, dtheta = x
    dxdt = [dtheta,
            sin(theta)/length*(grav - amplitude*omega**2 * sin(omega*t))]
    return dxdt

plt.figure(figsize=(14, 5))
plt.subplot(121)
time = np.linspace(0, 1.5, 101)
sol = odeint(pend, [0.1, 0], time, args=(9.81, 1, 0.05, 10))
theta, _ = sol.T
plt.plot(time, theta*180/pi, color="#e41a1c", lw=2)
plt.xlim(0, 1.5)
plt.ylim(0, 100)
plt.xlabel(r"$t$ (s)")
plt.ylabel(r"$\theta$ (deg)")
plt.legend([r"$\omega=10$"], numpoints=1, framealpha=0)

plt.subplot(122)
time = np.linspace(0, 3, 1001)
sol = odeint(pend, [0.05, 0], time, args=(9.81, 1, 0.05, 200))
theta, omega = sol.T
plt.plot(time, theta*180/pi, color="#e41a1c", lw=2)
plt.xlim(0, 2.5)
plt.xlabel(r"$t$ (s)")
plt.ylabel(r"$\theta$ (deg)")
plt.legend([r"$\omega=200$"], numpoints=1, framealpha=0)

plt.savefig("inverted_pendulum_oscillatory_base.svg", bbox_inches="tight")
plt.show()