import numpy as np
import matplotlib.pyplot as pplt

def euler(f, y0, T):
    Y = [y0]
    for i in range(1,len(T)):
        dt = T[i]-T[i-1]
        Y.append(Y[-1]+dt*f(Y[-1],dt))
        if(Y[-1][0]<0):
            return Y
    return Y
  
# Y = [z,z',y,y']
T = np.linspace(0,10,10000,dtype=np.float)

g = 9.81
alpha = 1
me = 63
qe = 1
B = 0.1

def equadiff(Y,dt):
    z,zp,y,yp = Y
    
    nzpp = - alpha/me*zp*zp - qe*B/me*yp - g
    nypp = - alpha/me*yp*yp - qe*B/me*zp
    
    nzp = zp + dt*nzpp
    nyp = yp + dt*nypp
    
    return np.array([nzp,nzpp,nyp,nypp],dtype=np.float)
    

Y = euler(equadiff,np.array([1000,0,0,0],dtype=np.float),T)
Tt = T[:len(Y)]
Yy = [y[0] for y in Y]
Zz = [y[2] for y in Y]
#pplt.plot(Tt,Yy)
#pplt.plot(Tt,Zz)

pplt.plot(Zz,Yy)
pplt.show()