# Using GEKKO to solve Discrete-time Optimal Control Problems stated in 
# DMISER OCTManual
#
# Discrete-time Problem 1 (OCTManual page 3.5)
#
# we are trying to introduce the use of GEKKO to native DMISER user
# this file is prepared by Dr. Joseph Lee
#
# Dr. H. W. J. Lee
# Department of Applied Mathematics
# The Hong Kong Polytechnic University

# if gekko is not installed, uncomment (delete #) in the following line
#! pip install gekko
# only need to do it once, and put it (#) back as comment once gekko is
# installed, e.g for anaconda python, you only need to do it once to
# install gekko
#
# however, if you are using Google CoLab to run python,
# then, you need to install it every time, if so, keep uncomment if you are
# using Google CoLab

# The DMISER problem 1
import numpy as np
from gekko import GEKKO
import matplotlib.pyplot as plt
# Initialize gekko model
m = GEKKO()

M=100 # time from i=0 to i=99

u =  [m.Var(0.04,lb=0,ub=0.04) for i in range(M)]
x1 = [m.Var(1,lb=-1000,ub=1000) for i in range(M)] 
#note, initial value 1 must also be put here
x2 = [m.Var(0,lb=-1000,ub=1000) for i in range(M)] 
#note, initial value 0 must also be put here
x3 = [m.Var(0,lb=-1000,ub=1000) for i in range(M)] 
#note, initial value 0 must also be put here

# initial conditions must be put here as well
m.Equation(x1[0]==1.0)
m.Equation(x2[0]==0.0)
m.Equation(x3[0]==0.0)

# Equations (different for each phase)
for i in range(M-1):
  m.Equation(x1[i+1]==x1[i]-u[i])
  m.Equation(x2[i+1]==x2[i]+x3[i])
  m.Equation(x3[i+1]==x3[i]+(2*u[i]-0.05*(x3[i])**2*m.exp(0.01*x2[i]))/x1[i]-0.01)

# terminal constraint
m.Equation(x1[M-1]==0.2)

# Objective
# Maximize final x2 while keeping third phase = -1
m.Obj(-x2[M-1])
# Solver options
m.options.IMODE = 3
# Solve
m.solve(disp=False)






# Plot
plt.figure(figsize=(15,8))
plt.subplot(3,1,1)
plt.plot(x1,'ro', label='$x1$')
plt.plot(x3,'go', label='$x3$')
plt.legend()
plt.subplot(3,1,2)
plt.plot(x2,'bo', label='$x2$')
plt.legend()
plt.subplot(3,1,3)
plt.plot(u,'yo', label='$u$')
plt.legend()
plt.show()