'''Programma per studiare il moto di un pendolo doppio'''


import numpy as np
import scipy as sp
import matplotlib.pyplot as plt

class pendolo(object):
    def __init__(self,m1,m2,l1,l2):
        self.g=9.81
        self.m1=m1
        self.m2=m2
        self.l1=l1
        self.l2=l2

    def coordinate(self,y):
        x=np.zeros((4),dtype=np.float64)
        x[0]=self.l1*np.sin(y[0])
        x[1]=-self.l1*np.cos(y[0])
        x[2]=self.l1*np.sin(y[0])+self.l2*np.sin(y[1])
        x[3]=-self.l1*np.cos(y[0])-self.l2*np.cos(y[1])
        return x

    def velocita(self,y):
        v=np.zeros((4),dtype=np.float64)
        v[0]=self.l1*np.cos(y[0])*y[2]
        v[1]=self.l1*np.sin(y[0])*y[2]
        v[2]=self.l1*np.cos(y[0])*y[2]+self.l2*np.cos(y[1])*y[3]
        v[3]=self.l1*np.sin(y[0])*y[2]+self.l2*np.sin(y[1])*y[3]
        return v

    def energia_cinetica_cart(self,v):
        kin=0.5*self.m1*(v[0]**2+v[1]**2)+0.5*self.m2*(v[2]**2+v[3]**2)
        return kin
    
    def energia_potenziale_cart(self,x):
        pot=self.m1*self.g*x[1]+self.m2*self.g*x[3]
        return pot
    
    def forza_generalizzata(self,y):
        f=np.zeros((4),dtype=np.float64)
        f[0]=y[2]
        f[1]=y[3]
        
        return f



def leap_frog(y0,y1,f,dt):
    y2=y0+2.0*f(y1)*dt
    return y2

def eulero_esplicito(y0,f,dt):
    y2=y0+f(y0)*dt
    return y2

def runge_kutta(y0,f,dt):
    Y1=y0
    Y2=y0+f(Y1)*dt/2.0
    Y3=y0+f(Y2)*dt/2.0
    Y4=y0+f(Y3)*dt
    y1=y0+(f(Y1)+2.0*f(Y2)+2.0*f(Y3)+f(Y4))*dt/6.0
    return y1








def main():

    m1=float(input("Massa 1 (Kg):\t"))
    m2=float(input("Massa 2 (Kg):\t"))
    l1=float(input("Lunghezza 1 (m):\t"))
    l2=float(input("Lunghezza 2 (m):\t"))

    ang1=float(input("Angolo 1 (Gradi):\t"))*np.pi/180
    ang2=float(input("Angolo 2 (Gradi):\t"))*np.pi/180

    vang1=float(input("Velocità angolare 1 (Gradi):\t"))*np.pi/180
    vang2=float(input("Velocità angolare 2 (Gradi):\t"))*np.pi/180

    pend=pendolo(m1,m2,l1,l2)

    
    iprop=input("Propagatore: 1-Eulero esplicito 2-Leap-Frog 3-Runge Kutta: ")
    iprop=int(iprop)

    y0=np.zeros((4),dtype=np.float64)
    y0[0]=ang1
    y0[1]=ang2
    y0[2]=vang1
    y0[3]=vang2

    dt=float(input("Passo temporale (s):\t"))
    n=int(input("Numero di passi temporali:\t"))

    for it in range(n):
        if(iprop==1 or (iprop==2 and it==0)):
            y1=eulero_esplicito(y0,pend.forza_generalizzata,dt)
        elif(iprop==2):
            y1=leap_frog(ym1,y0,pend.forza_generalizzata,dt)
        elif(iprop==3):
            y1=runge_kutta(y0,pend.forza_generalizzata,dt)

        ym1=y0
        y0=y1

        cart_x=pend.coordinate(ym1)
        cart_v=pend.velocita(ym1)
        ene_kin=pend.energia_cinetica_cart(cart_v)
        ene_pot=pend.energia_potenziale_cart(cart_x)
        print("Passo:\t"+str(it)+"\tEnergia (J):\t"+str(ene_kin+ene_pot))   
    
   

if __name__ == "__main__":
    main()