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

m = GEKKO(remote=False)
n = 101; m.time = np.linspace(0,1,n)

ksi = 0.05236
tf = m.FV(value=5, lb=1, ub=10)
tf.STATUS = 1
m.Minimize(tf/n)

# x0 = angle of attack
# x1 = pitch angle
# x2 = pitch rate
x0,x1,x2 = m.Array(m.Var,3,value=0,lb=-1,ub=1)
x0.value = 0.4655

# Control signal - aileron angle
w = m.MV(value=0, lb=0, ub=1); w.STATUS=1

m.Equation(x0.dt()/tf == -0.877*x0 + x2 - 0.088*x0*x2
           + 0.47*x0**2 - 0.019*x1**2 -x0**2*x2 + 3.846*x0**3
           -(0.215*ksi-0.28*x0**2*ksi -0.47*x0*ksi**2 - 0.63*ksi**3)*w
           -(-0.215*ksi + 0.28*x0**2 -0.47*x0*ksi**2 + 0.63*ksi**3)*(1-w))
m.Equation(x1.dt()/tf==x2)
m.Equation(x2.dt()/tf == -4.208*x0 - 0.396*x2 - 0.47*x0**2 - 3.564*x0**3
           +20.967*ksi - 6.265*x0**2*ksi + 46*x0*ksi**2 -61.4*ksi**3
           -(20.967*ksi - 6.265*x0**2*ksi - 61.4*ksi**3)*2*w)

m.fix_final(x0,0)
m.fix_final(x1,0)
m.fix_final(x2,0)

m.options.IMODE=6
m.options.SOLVER=1
m.options.MAX_ITER=1000

m.options.COLDSTART=0
m.solve()

tm = tf.value[0]*m.time
plt.figure(figsize=(6,4))
plt.subplot(2,1,1)
plt.plot(tm,w,'r:',lw=2,label='Aileron Angle')
plt.grid(); plt.legend()
plt.subplot(2,1,2)
plt.plot(tm,x0,'r-',lw=2,label='AOA')
plt.plot(tm,x1,'b:',lw=2,label='Pitch')
plt.plot(tm,x2,'k--',lw=2,label='Pitch Rate')
plt.grid(); plt.legend()
plt.savefig('aircraft.png',dpi=300);
plt.tight_layout(); plt.show()