from numpy import sin, cos
import numpy as np
import matplotlib.pyplot as plt
import scipy.integrate as integrate
from ipywidgets import *

G = 9.8  # acceleration due to gravity, in m/s^2

def derivs(state, t,params):
    L1, L2, M1, M2 = params

    dydx = np.zeros_like(state)
    dydx[0] = state[1]

    del_ = state[2] - state[0]
    den1 = (M1 + M2)*L1 - M2*L1*cos(del_)*cos(del_)
    dydx[1] = (M2*L1*state[1]*state[1]*sin(del_)*cos(del_) +
               M2*G*sin(state[2])*cos(del_) +
               M2*L2*state[3]*state[3]*sin(del_) -
               (M1 + M2)*G*sin(state[0]))/den1

    dydx[2] = state[3]

    den2 = (L2/L1)*den1
    dydx[3] = (-M2*L2*state[3]*state[3]*sin(del_)*cos(del_) +
               (M1 + M2)*G*sin(state[0])*cos(del_) -
               (M1 + M2)*L1*state[1]*state[1]*sin(del_) -
               (M1 + M2)*G*sin(state[2]))/den2

    return dydx

def example(th1,w1,th2,w2,l1,l2,m1,m2, Tmax):
    pars = (l1, l2, m1, m2)
    ders = lambda state, t: derivs(state, t, pars)
    dt = 0.001
    t = np.arange(0.0, Tmax, dt)
    state = np.radians([th1, w1, th2, w2])
    y = integrate.odeint(ders, state, t)
    x1 = l1*sin(y[:, 0])
    y1 = -l1*cos(y[:, 0])
    x2 = l2*sin(y[:, 2]) + x1
    y2 = -l2*cos(y[:, 2]) + y1
    plt.plot(x2,y2, label = 'bottom pendulum')
    plt.legend()
    plt.scatter([0.0],[0.0],c='orange',s=36.0)
    plt.scatter([x2[0]],[y2[0]],c='green',s=36.0)
    plt.scatter([x2[-1]],[y2[-1]],c='red',s=36.0)

interact(example,
         th1 = BoundedFloatText(value=120.0, min = - 180.0, max = 180.0, description='Angle 1 (degrees):'), 
         w1 = FloatText(value=0.0, description='Angular velocity 1:'),
         th2 = BoundedFloatText(value=-10.0, min = - 180.0, max = 180.0, description='Angle 2 (degrees):'),
         w2 = FloatText(value=0.0, description='Angular velocity 2:'),
         l1 = BoundedFloatText(value=1.0, min = 0.01, max = 5.0, description='L1:'),
         l2 = BoundedFloatText(value=1.0, min = 0.01, max = 5.0, description='L2:'),
         m1 = BoundedFloatText(value=1.0, min = 0.01, max = 10.0, description='M1:'),
         m2 = BoundedFloatText(value=1.0, min = 0.01, max = 10.0, description='M2:'),
         Tmax = BoundedFloatText(value=20.0, min = 0.01, max = 1000.0, description='Max time:'),
         __manual=True)