import numpy as np
from APMonitor.apm import *  # pip install APMonitor

##################################################################
## Step #1 - Solve model with p = 1
##################################################################

## Step #1a - write data.csv
n = 31
time = np.linspace(0,3,n)
p = np.ones(31)
x = 2 * np.ones(31)

fid = open('data.csv','w')

## write time row
fid.write('time, ')
for i in range(n-1):
    fid.write(str(time[i]) +  ', ')
fid.write(str(time[n-1]) + '\n')

## write 'p' row (input parameter)
fid.write('p, ')
for i in range(n-1):
    fid.write(str(p[i]) +  ', ')
fid.write(str(p[n-1]) + '\n')

## write 'x' row (state variable initialization)
fid.write('x, ')
# imode: https://apmonitor.com/wiki/index.php/Main/Modes
# for imode=4-6, include all initialization values
# for imode=7-9, include only the initial condition for variables
imode = 7
if ((imode>=4) and (imode<=6)):
    for i in range(n-1):
        fid.write(str(x[i]) + ', ')
    fid.write(str(x[n-1]) + '\n')
else:
    fid.write(str(x[0]) + ', ')
    for i in range(1,n-1):
        fid.write('-, ')
    fid.write('-\n')

# close file
fid.close()

## Step 1b - Load and solve model
s = 'https://byu.apmonitor.com'
a = 'model_init'

apm(s,a,'clear all')
apm_load(s,a,'model.apm')
csv_load(s,a,'data.csv')

apm_option(s,a,'apm.time_shift',1)
apm_option(s,a,'apm.imode',imode)
output1 = apm(s,a,'solve')

## Step 1c - Retrieve results with solution.csv
solution1 = apm_sol(s,a)

##################################################################
## Step 2 - Solve again with prior solution for initialization and
##          p as a step from 0 to 2
##################################################################

## Change to imode = 4 and change p trajectory
p[0:5] = 0.0
p[5:n] = 2.0

## Step 2a - Write new row at the end of solution.csv
fname = 'solution_' + a + '.csv'
fid = open(fname,'a')  # append to file
fid.write('p, ')
for i in range(n-1):
    fid.write(str(p[i]) +  ', ')
fid.write(str(p[n-1]) + '\n')
# close file
fid.close()

## Step 2b - Reload csv file for initialization
apm(s,a,'clear csv')
csv_load(s,a,fname)

## Step 2c - Solve again but with new inputs
imode = 4
apm_option(s,a,'apm.time_shift',0)
apm_option(s,a,'apm.imode',imode)
output2 = apm(s,a,'solve')
print(output2)

## Step 2d - Retrieve results with solution.csv
solution2 = apm_sol(s,a)

##################################################################
## Step 3 - Create plots
##################################################################
import matplotlib.pyplot as plt

plt.figure(1)
plt.subplot(2,1,1)
plt.plot(solution1['time'],solution1['p'],'k-',lw=2)
plt.plot(solution2['time'],solution2['p'],'b--',lw=2)
plt.legend([r'$p_1$',r'$p_2$'])
plt.subplot(2,1,2)
plt.plot(solution1['time'],solution1['x'],'r--',lw=2)
plt.plot(solution2['time'],solution2['x'],'g:',lw=2)
plt.legend([r'$x_1$',r'$x_2$'])
plt.xlabel('time')
plt.show()