#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Mon Mar  8 21:34:49 2021

Gekko implementation of the simple energy storage model found here:
    https://www.sciencedirect.com/science/article/abs/pii/S030626191500402X

Useful link:
    https://apmonitor.com/wiki/index.php/Apps/PeriodicBoundaryConditions    

@author: Nathaniel Gates, John Hedengren
"""

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

m = GEKKO(remote=False)

t = np.linspace(0, 24, 24*3+1)
m.time = t

m.options.SOLVER   = 1
m.options.IMODE    = 6
m.options.NODES    = 3
m.options.CV_TYPE  = 1
m.options.MAX_ITER = 300

p = m.FV()           # production
p.STATUS = 1
s = m.Var(100, lb=0) # storage inventory
store = m.SV()       # store energy rate
vy = m.SV(lb=0)      # store slack variable
recover = m.SV()     # recover energy rate
vx = m.SV(lb=0)      # recover slack variable

eps = 0.7

d = m.MV(-20*np.sin(np.pi*t/12)+100)

m.periodic(s)

m.Equations([p + recover/eps - store >= d,
             p - d == vx - vy,
             store == p - d + vy,
             recover == d - p + vx,
             s.dt() == store - recover/eps,
             store * recover <= 0])
m.Minimize(p)

m.solve(disp=True)

#%% Visualize results
fig, axes = plt.subplots(4, 1, sharex=True)

ax = axes[0]
ax.plot(t, store, 'C3-', label='Store Rate')
ax.plot(t, recover, 'C0-.', label='Recover Rate')

ax = axes[1]
ax.plot(t, d, 'k-', label='Electricity Demand')
ax.plot(t, p, 'C3--', label='Power Production')

ax = axes[2]
ax.plot(t, s, 'C2-', label='Energy Inventory')

ax = axes[3]
ax.plot(t, vx, 'C2-', label='$S_1$')
ax.plot(t, vy, 'C3--', label='$S_2$')
ax.set_xlabel('Time (hr)')

for ax in axes:
    ax.legend(bbox_to_anchor=(1.01, 0.5), \
              loc='center left', frameon=False)
    ax.grid()
    ax.set_xlim(0, 24)
    loc = mtick.MultipleLocator(base=6)
    ax.xaxis.set_major_locator(loc)

plt.tight_layout()
plt.show()