import numpy as np
from gekko import GEKKO
import matplotlib.pyplot as plt
import json
from IPython import display
from IPython import get_ipython

# detect if IPython is running
ipython = get_ipython() is not None

m = GEKKO(remote=False)
t = np.linspace(0,15,31)
m.time = t # 0.5 cycle time
u = m.MV(1,name='u')
x = m.CV(name='x')
y = m.CV(name='y')
m.Equations([5*x.dt() ==-x+1.5*u,\
             15*y.dt()==-y+3.0*u])
m.options.IMODE = 6
m.options.NODES = 3

# MV tuning
u.DCOST   = 0.1
u.DMAX    = 1.5
u.STATUS  = 1
u.FSTATUS = 0
u.UPPER   = 10
u.LOWER   = 0

# CV tuning
x.STATUS  = 1
x.FSTATUS = 0
x.SPHI    = 10
x.SPLO    = 9
x.TAU     = 2
x.TR_INIT = 1
x.TR_OPEN = 5
x.WSPHI   = 100
x.WSPLO   = 100

y.STATUS  = 1
y.FSTATUS = 0
y.SPHI    = 7
y.SPLO    = 2
y.TAU     = 0
y.TR_INIT = 0
y.TR_OPEN = 1
y.WSPHI   = 200
y.WSPLO   = 50

# Make an MP4 animation?
make_mp4 = True
if make_mp4:
    import imageio  # required to make animation
    import os
    try:
        os.mkdir('./figures')
    except:
        pass

# plot solution
plt.figure(figsize=(8,5))
plt.ion()
plt.show()

for i in range(40):
    m.solve(disp=False)

    # update time
    tm = m.time+i*0.5

    # x is more important at t=10
    if tm[0]>=10:
        x.WSPLO = 500
        x.WSPHI = 500

    # get trajectory information
    with open(m.path+'//results.json') as f:
        z = json.load(f)

    if ipython:
        display.clear_output(wait=True)
    else:
        plt.clf()

    plt.figure(1,figsize=(8,5))
    plt.subplot(3,1,1)
    plt.plot(tm,u.value,'k-',lw=2,label='u')
    plt.ylim([-1,11])
    plt.legend(loc=1)
    plt.ylabel('MV')

    plt.subplot(3,1,2)
    plt.plot(tm,x.value,'b-',lw=2,label='x')
    plt.plot(tm,z['x.tr_hi'],'k--',label='x traj')
    plt.plot(tm,z['x.tr_lo'],'k--')
    plt.legend(loc=1)

    plt.subplot(3,1,3)
    plt.plot(tm,y.value,'r-',lw=2,label='y')
    plt.plot(tm,z['y.tr_hi'],'k--',label='y traj')
    plt.plot(tm,z['y.tr_lo'],'k--')
    plt.legend(loc=1)
    plt.draw()
    if make_mp4:
        filename='./figures/plot_'+str(i+10000)+'.png'
        plt.savefig(filename)
    plt.pause(0.01)

# generate mp4 from png figures
if make_mp4:
    images = []
    for i in range(40):
        filename='./figures/plot_'+str(i+10000)+'.png'
        images.append(imageio.imread(filename))
    try:
        imageio.mimsave('results.mp4', images)
    except:
        print('Error: Install mp4 generator with "pip install imageio-ffmpeg"')