from gekko import GEKKO
import numpy as np
import matplotlib.pyplot as plt

#number of points in time discretization
n = 91

#Initialize Model
m = GEKKO(remote=False)

#define time discretization
m.time = np.linspace(0,90,n)

#make array of drag coefficients, changing at time 60
drag = [(0.2 if t<=60 else 10) for t in m.time]

#define constants
g = m.Const(value=9.81)
mass = m.Const(value=80)

#define drag parameter
d = m.Param(value=drag)

#initialize variables
x,y,vx,vy,v,Fx,Fy = [m.Var(value=0) for i in range(7)]

#initial conditions
y.value = 5000
vx.value = 50

#Equations
# momentum balance
m.Equation(Fx == -d * vx**2)
m.Equation(Fy == -mass*g + d*vy**2)
#F = ma
m.Equation(Fx/mass == vx.dt())
m.Equation(Fy/mass == vy.dt())
#vel = dxdt
m.Equation(vx == x.dt())
m.Equation(vy == y.dt())
#total velocity
m.Equation(v == (vx**2 + vy**2)**.5)

#Set global options
m.options.IMODE = 4 #dynamic simulation

#Solve simulation
m.solve()

#%% Plot results
plt.figure(1)
plt.plot(x.value,y.value,'r--',label='Path')
plt.xlabel('x')
plt.ylabel('y')
plt.legend(); plt.grid()

plt.figure(2)
plt.subplot(2,1,1)
plt.plot(m.time,x.value,label='x')
plt.plot(m.time,y.value,label='y')
plt.xlabel('time')
plt.legend(); plt.grid()

plt.subplot(2,1,2)
plt.plot(m.time,vx.value,label='vx')
plt.plot(m.time,vy.value,label='vy')
plt.xlabel('time')
plt.legend(); plt.grid()

plt.show()