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


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


import matplotlib.pyplot as plt
import matplotlib.animation as animation

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]
        A=1.0/((self.m1+self.m2*(1.0-np.cos(y[0]-y[1])**2))*self.l1)
        T1=self.m2*self.l2*np.sin(y[0]-y[1])*y[3]**2
        T2=self.g*(self.m1+self.m2)*np.sin(y[0])
        T3=-np.cos(y[0]-y[1])*(-self.m2*self.l1*np.sin(y[0]-y[1])*y[2]**2+self.m2*self.g*np.sin(y[1]))
        f[2]=-A*(T1+T2+T3)
        B=1.0/((self.m2-(self.m2**2/(self.m1+self.m2)*np.cos(y[0]-y[1])**2))*self.l2)
        S1=-(self.m2/(self.m1+self.m2))*np.cos(y[0]-y[1])*(self.m2*self.l2*np.sin(y[0]-y[1])*y[3]**2+self.g*(self.m1+self.m2)*np.sin(y[0]))
        S2=-self.m2*self.l1*np.sin(y[0]-y[1])*y[2]**2
        S3=self.m2*self.g*np.sin(y[1])
        f[3]=-B*(S1+S2+S3)

      
        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 update(frame):
        point.set_data(x_data[:frame + 1], y_data[:frame + 1])
        point.set_3d_properties(z_data[:frame + 1])
        return point



def animazione(dt,t,coord):
    fig, ax = plt.subplots()
    l1x=np.zeros((2),dtype=np.float64)
    l1y=np.zeros((2),dtype=np.float64)
    l1x[0]=0
    l1y[0]=0
    l1x[0]=coord[0,0]
    l1y[0]=coord[1,0]
    line1 = ax.plot(l1x,l1y)[0]
    l2x=np.zeros((2),dtype=np.float64)
    l2y=np.zeros((2),dtype=np.float64)
    l2x[0]=coord[0,0]
    l2y[0]=coord[1,0]
    l2x[0]=coord[2,0]
    l2y[0]=coord[3,0]
    line2 = ax.plot(l2x,l2y)[0]
    ax.set(xlim=[-4, 4], ylim=[-4, 4], xlabel='X [m]', ylabel='Y [m]')
    ax.legend()


    def update(frame):
        l1x[0]=0
        l1y[0]=0
        l1x[1]=coord[0,frame]
        l1y[1]=coord[1,frame]
        l2x[0]=coord[0,frame]
        l2y[0]=coord[1,frame]
        l2x[1]=coord[2,frame]
        l2y[1]=coord[3,frame]
        line1.set_xdata(l1x)
        line1.set_ydata(l1y)
        line2.set_xdata(l2x)
        line2.set_ydata(l2y)
        return (line1)


    ani = animation.FuncAnimation(fig=fig, func=update, frames=len(t), interval=int(dt)*1000)
    plt.show()


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"))
    npassi=int(input("Stampa  ogni passi:\t"))

    t=np.zeros((int(n/npassi)),dtype=np.float64)
    coord=np.zeros((4,int(n/npassi)),dtype=np.float64)
    ip=0
    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)
        if(it%npassi==0):
            t[ip]=it*dt
            coord[0:4,ip]=cart_x[0:4]
            ip+=1
        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))   
    
    animazione(dt*npassi,t,coord)

if __name__ == "__main__":
    main()