import numpy as np
import copy
import matplotlib.pyplot as plt
import ast

def forza_generalizzata(y,m,g, Catt):
    f=np.zeros(4)
    f[0]=y[2]
    f[1]=y[3]
    vel=np.sqrt(y[2]**2+y[3]**2)
    u_x=y[2]/vel
    u_y=y[3]/vel
    f[2]=-u_x*Catt*vel**2/m
    f[3]=-u_y*Catt*vel**2/m-g
    return f


def main():
    g=9.81
    rho_aria=1.22
    rho_gran=2700.0
    R=0.1
    Cd=0.47


    A=np.pi*R**2
    Catt=0.5*rho_aria*Cd*A
    m=rho_gran*(4/3)*np.pi*R**3


    theta=float(input('Alzo (gradi): '))
    theta*=np.pi/180
    v0=float(input('Velocità iniziale (m/s): '))
    dt=float(input('Time step (s): '))
    
    it=0

    y0=np.zeros(4)
    y1=np.zeros(4)

    y0[2]=np.cos(theta)*v0
    y0[3]=np.sin(theta)*v0

    y0RK2=copy.deepcopy(y0)
    y1RK2=copy.deepcopy(y1)
    ymRK2=copy.deepcopy(y1)

    y0RK4=copy.deepcopy(y0)
    y1RK4=copy.deepcopy(y1)
    Y_1RK4=copy.deepcopy(y1)
    Y_2RK4=copy.deepcopy(y1)
    Y_3RK4=copy.deepcopy(y1)
    Y_4RK4=copy.deepcopy(y1)

    found=False

    it=0
    nn=int(input('Numero passi: '))
    time=0.0

    errEul=[]
    errRK2=[]
    errRK4=[]
    times=[]

    latt=ast.literal_eval(input('Inserisci l\'attrito True/False: '))
 
    if(not latt):
        Catt=0.0
    for it in range(nn):
        #Eulero esplicito
        y1=y0+forza_generalizzata(y0,m,g, Catt)*dt
        y0=y1
        #RK2 o mid-point Explicit Euler 
        ymRK2=y0RK2+forza_generalizzata(y0RK2,m,g, Catt)*dt/2.0
        y1RK2=y0RK2+forza_generalizzata(ymRK2,m,g,Catt)*dt
        y0RK2=y1RK2

        #Runge Kutta
        Y_1RK4=y0RK4
        Y_2RK4=y0RK4+forza_generalizzata(Y_1RK4,m,g,Catt)*dt/2
        Y_3RK4=y0RK4+forza_generalizzata(Y_2RK4,m,g,Catt)*dt/2
        Y_4RK4=y0RK4+forza_generalizzata(Y_3RK4,m,g,Catt)*dt
        y1RK4=y0RK4+(forza_generalizzata(Y_1RK4,m,g,Catt)+2*forza_generalizzata(Y_2RK4,m,g,Catt)+2*forza_generalizzata(Y_3RK4,m,g,Catt)+forza_generalizzata(Y_4RK4,m,g,Catt))*dt/6
        y0RK4=y1RK4

        time+=dt

        x=np.cos(theta)*v0*time
        y=np.cos(theta)*v0*time-(1/2)*g*time**2

        times.append(time)
        if(not latt):
            errEul.append(np.sqrt((x-y1[0])**2+(y-y1[1])**2))
            errRK2.append(np.sqrt((x-y1RK2[0])**2+(y-y1RK2[1])**2))
            errRK4.append(np.sqrt((x-y1RK4[0])**2+(y-y1RK4[1])**2))
        else:
            errEul.append(np.sqrt((y1RK4[0]-y1[0])**2+(y1RK4[1]-y1[1])**2))
            errRK2.append(np.sqrt((y1RK4[0]-y1RK2[0])**2+(y1RK4[1]-y1RK2[1])**2))



    fig, ax = plt.subplots()
    ax.semilogy()

    ax.set_xlabel('tempo (s)')
    ax.set_ylabel('Distanza (m)')
   
    plt.plot(times, errEul, label='Eulero Esplicito')
    plt.plot(times, errRK2, label='Runge-Kutta 2')
    if(not latt): 
        plt.plot(times, errRK4, label='Runge-Kutta 4')

    ax.legend()
    plt.show(block=True)
    










if __name__ == "__main__":
    main()