# -*- coding: utf-8 -*-
"""
Ordinary Differential Equations
legge oraria e phase portrait del pendolo
"""

import numpy as np
import scipy as sp
import matplotlib.pyplot as plt
from scipy.integrate import ode    

def func(t, y, w0, damping):
    """w0^2=g/l
    y[0]=theta;   
    y[1] = dtheta/dt
    dy[0]/dt = y[1]
    dy[1]/dt = - w0**2 * sin(y[0])
    """
    return [y[1], - w0**2 * np.sin(y[0])] # array or list
    
w0 = 1.0            # 'natural' frequency sqrt(g/L)
T0 = 2*np.pi / w0   # period for small (harmonic) oscillations

t0 = 0.0            # initial time
theta0 = 0.0        # assume initial position at the bottom
# initial dtheta/dt=v0; 
# libration for v0 < 2*w0,
# rotation otherwise

dt = 0.01 * T0      # 'reasonable' time step

r = ode(func)       # GRAB THE ODE OBJECT  dy/dt = func(t,y)
r.set_integrator('vode', rtol=1.e-8)
"""
# legge oraria theta(t)
v0 = 3                     
r.set_f_params(w0, 0.0)    # set natural freq and damping
r.set_initial_value([theta0, v0], t0)
time = np.array([t0])
theta = np.array([theta0])
v = np.array([v0])
E = np.array([0.5*v0**2 + w0**2*(1.-np.cos(theta0))])   # initial energy

while r.successful() and r.t<20*T0 and r.y[0]<20.1*np.pi:
    r.integrate(r.t+dt)
    time = np.hstack((time, r.t))
    theta = np.hstack((theta, r.y[0]))
    v = np.hstack((v, r.y[1]))
    E = np.hstack((E, 0.5*r.y[1]**2 +  w0**2*(1.-np.cos(r.y[0]))))

plt.figure()
plt.subplot(311)
plt.plot(time, theta)
plt.ylabel('theta')
plt.subplot(312)
plt.plot(time, v)
plt.ylabel('v')
plt.subplot(313)
plt.plot(time, E)
plt.ylabel('energy')
plt.xlabel('time')

"""
# phase portrait

v0 = np.array([0.1, 0.5, 1.0, 1.5, 2.0, 2.5, 3.0])
plt.figure()
for ii in range(len(v0)):
    r.set_f_params(w0, 0.0)
    r.set_initial_value([theta0, v0[ii]], t0)
    time = np.array([t0])
    theta = np.array([theta0])
    v = np.array([v0[ii]])
    while r.successful() and r.t<4*T0 and r.y[0]<2.1*np.pi:
        r.integrate(r.t+dt)
        time = np.hstack((time, r.t))
        theta = np.hstack((theta, r.y[0]))
        v = np.hstack((v, r.y[1]))
    plt.plot(theta, v, color='b')


plt.show()

"""
Energia E/mL^2 = 0.5(dtheta/dt)^2 +w0^2(1-cos(theta))
deviazione dal valore iniziale al variare
"""