# -*- coding: utf-8 -*-
"""
Creazione di una class con metodi per
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])]

class PendulumExercise(object):
    
    def __init__(self, w0=1.0, damping=0.0, t0=0.0, theta0=0.0):
        self.t0 = t0
        self.theta0 = theta0
        self.w0 = w0
        r = ode(func)
        r.set_f_params(w0, damping)
        self.r = r
    
    def evolve(self, v0, iplot=False):
        r = self.r
        r.set_initial_value([self.theta0, v0], self.t0)
        T0 = 2*np.pi / self.w0
        dt = 0.01 * T0  #<<<<<<<<<<
        time = np.array([self.t0])
        theta = np.array([self.theta0])
        v = np.array([v0])
        while r.successful() and r.t<2*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]))
        self.time = time
        if iplot:
            plt.figure()
            h1 = plt.subplot(211)
            plt.plot(time, theta)
            plt.subplot(212, sharex=h1)
            plt.plot(time, v)
        else:
            return theta, v
            
    def portrait(self, v0):
        plt.figure()
        for ii in range(len(v0)):
            theta, v = self.evolve(v0[ii])
            plt.plot(theta, v, color='b')


v0 = np.array([0.1, 0.5, 1.0, 1.5, 2.0, 2.5, 3.0])
p = PendulumExercise()
p.portrait(v0)

plt.figure()
theta, v = p.evolve(3.0)
plt.plot(p.time, theta)
theta, v = p.evolve(5.0)
plt.plot(p.time, theta)

plt.show()