import numpy as np
import sys
from casadi import *

# Add do_mpc to path. This is not necessary if it was installed via pip
sys.path.append('../../../')

# Import do_mpc package:
import do_mpc

import matplotlib.pyplot as plt

"""
Explanation: Continuous stirred tank reactor (CSTR)
In this Jupyter Notebook we illustrate the example CSTR.
Open an interactive online Jupyter Notebook with this content on Binder:

The example consists of the three modules template_model.py, which describes the system model, template_mpc.py, which defines the settings for the control and template_simulator.py, which sets the parameters for the simulator.
The modules are used in main.py for the closed-loop execution of the controller.
The file post_processing.py is used for the visualization of the closed-loop control run.
One exemplary result will be presented at the end of this tutorial as a gif.
In the following the different parts are presented. But first, we start by importing basic modules and do-mpc.
End of explanation
"""


model_type = 'continuous' # either 'discrete' or 'continuous'
model = do_mpc.model.Model(model_type)

"""
Explanation: Model
In the following we will present the configuration, setup and connection between these blocks, starting with the model.
The considered model of the CSTR is continuous and has 4 states and 2 control inputs.
The model is initiated by:
End of explanation
"""


# States struct (optimization variables):
C_a = model.set_variable(var_type='_x', var_name='C_a', shape=(1,1))
C_b = model.set_variable(var_type='_x', var_name='C_b', shape=(1,1))
T_R = model.set_variable(var_type='_x', var_name='T_R', shape=(1,1))
T_K = model.set_variable(var_type='_x', var_name='T_K', shape=(1,1))

"""
Explanation: States and control inputs
The four states are concentration of reactant A ($C_{\text{A}}$), the concentration of reactant B ($C_{\text{B}}$), the temperature inside the reactor ($T_{\text{R}}$) and the temperature of the cooling jacket ($T_{\text{K}}$):
End of explanation
"""


# Input struct (optimization variables):
F = model.set_variable(var_type='_u', var_name='F')
Q_dot = model.set_variable(var_type='_u', var_name='Q_dot')

"""
Explanation: The control inputs are the feed $F$ and the heat flow $\dot{Q}$:
End of explanation
"""


# Certain parameters
K0_ab = 1.287e12 # K0 [h^-1]
K0_bc = 1.287e12 # K0 [h^-1]
K0_ad = 9.043e9 # K0 [l/mol.h]
R_gas = 8.3144621e-3 # Universal gas constant
E_A_ab = 9758.3*1.00 #* R_gas# [kj/mol]
E_A_bc = 9758.3*1.00 #* R_gas# [kj/mol]
E_A_ad = 8560.0*1.0 #* R_gas# [kj/mol]
H_R_ab = 4.2 # [kj/mol A]
H_R_bc = -11.0 # [kj/mol B] Exothermic
H_R_ad = -41.85 # [kj/mol A] Exothermic
Rou = 0.9342 # Density [kg/l]
Cp = 3.01 # Specific Heat capacity [kj/Kg.K]
Cp_k = 2.0 # Coolant heat capacity [kj/kg.k]
A_R = 0.215 # Area of reactor wall [m^2]
V_R = 10.01 #0.01 # Volume of reactor [l]
m_k = 5.0 # Coolant mass[kg]
T_in = 130.0 # Temp of inflow [Celsius]
K_w = 4032.0 # [kj/h.m^2.K]
C_A0 = (5.7+4.5)/2.0*1.0 # Concentration of A in input Upper bound 5.7 lower bound 4.5 [mol/l]

# Uncertain parameters:
alpha = model.set_variable(var_type='_p', var_name='alpha')
beta = model.set_variable(var_type='_p', var_name='beta')

"""
Explanation: ODE and parameters
The system model is described by the ordinary differential equation:
\begin{align}
\dot{C}{\text{A}} &= F \cdot (C{\text{A},0} - C_{\text{A}}) - k_1 \cdot C_{\text{A}} - k_3 \cdot C_{\text{A}}^2, \
\dot{C}{\text{B}} &= F \cdot C{\text{B}} + k_1 \cdot C_{\text{A}} - k_2 \cdot C_{\text{B}}, \
\dot{T}{\text{R}} &= \frac{k_1 \cdot C{\text{A}} \cdot H_{\text{R},ab} + k_2 \cdot C_{\text{B}} \cdot  H_{\text{R},bc} + k_3 \cdot C_{\text{A}}^2 \cdot H_{\text{R},ad}} {-\rho \cdot c_p}\
&+ F \cdot (T_{\text{in}} - T_{\text{R}}) + \frac{K_w \cdot A_{\text{R}} \cdot(T_{\text{K}}-T_{\text{R}})}{\rho \cdot c_p \cdot V_{\text{R}}}, \
\dot{T}{\text{K}} &= \frac{\dot{Q} + k_w \cdot A{\text{R}} \cdot T_{\text{dif}}}{m_k \cdot C_{p,k}},
\end{align}
where
\begin{align}
k_1 &= \beta \cdot k_{0,\text{ab}} \cdot \exp\left(\frac{-E_{\text{A},\text{ab}}}{T_{\text{R}}+273.15}\right), \
k_1 &= k_{0,\text{bc}} \cdot \exp \left( \frac{-E_{\text{A},\text{bc}}}{T_{\text{R}}+273.15} \right), \
k_3 &= k_{0,\text{ad}} \cdot \exp \left( \frac{-\alpha \cdot E_{\text{A},\text{ad}}}{T_{\text{R}}+273.15} \right).
\end{align}
The 
parameters $\alpha$ and $\beta$ are uncertain while the rest of the parameters is considered certain:
End of explanation
"""


# Auxiliary terms
K_1 = beta * K0_ab * exp((-E_A_ab)/((T_R+273.15)))
K_2 =  K0_bc * exp((-E_A_bc)/((T_R+273.15)))
K_3 = K0_ad * exp((-alpha*E_A_ad)/((T_R+273.15)))

"""
Explanation: In the next step, we formulate the $k_i$-s:
End of explanation
"""


T_dif = model.set_expression(expr_name='T_dif', expr=T_R-T_K)

"""
Explanation: Additionally, we define an artificial variable of interest, that is not a state of the system, but will be later used for plotting:
End of explanation
"""


model.set_rhs('C_a', F*(C_A0 - C_a) -K_1*C_a - K_3*(C_a**2))
model.set_rhs('C_b', -F*C_b + K_1*C_a - K_2*C_b)
model.set_rhs('T_R', ((K_1*C_a*H_R_ab + K_2*C_b*H_R_bc + K_3*(C_a**2)*H_R_ad)/(-Rou*Cp)) + F*(T_in-T_R) +(((K_w*A_R)*(-T_dif))/(Rou*Cp*V_R)))
model.set_rhs('T_K', (Q_dot + K_w*A_R*(T_dif))/(m_k*Cp_k))

"""
Explanation: WIth the help ot the $k_i$-s and $T_{\text{dif}}$ we can define the ODEs:
End of explanation
"""


# Build the model
model.setup()

"""
Explanation: Finally, the model setup is completed:
End of explanation
"""


mpc = do_mpc.controller.MPC(model)

"""
Explanation: Controller
Next, the model predictive controller is configured.
First, one member of the mpc class is generated with the prediction model defined above:
End of explanation
"""


setup_mpc = {
    'n_horizon': 20,
    'n_robust': 1,
    'open_loop': 0,
    't_step': 0.005,
    'state_discretization': 'collocation',
    'collocation_type': 'radau',
    'collocation_deg': 2,
    'collocation_ni': 2,
    'store_full_solution': True,
    # Use MA27 linear solver in ipopt for faster calculations:
    #'nlpsol_opts': {'ipopt.linear_solver': 'MA27'}
}

mpc.set_param(**setup_mpc)

"""
Explanation: We choose the prediction horizon n_horizon, set the robust horizon n_robust to 1. The time step t_step is set to one second and parameters of the applied discretization scheme orthogonal collocation are as seen below:
End of explanation
"""


mpc.scaling['_x', 'T_R'] = 100
mpc.scaling['_x', 'T_K'] = 100
mpc.scaling['_u', 'Q_dot'] = 2000
mpc.scaling['_u', 'F'] = 100

"""
Explanation: Because the magnitude of the states and inputs is very different, we introduce scaling factors:
End of explanation
"""


_x = model.x
mterm = (_x['C_b'] - 0.6)**2 # terminal cost
lterm = (_x['C_b'] - 0.6)**2 # stage cost

mpc.set_objective(mterm=mterm, lterm=lterm)

mpc.set_rterm(F=0.1, Q_dot = 1e-3) # input penalty

"""
Explanation: Objective
The goal of the CSTR is to obtain a mixture with a concentration of $C_{\text{B,ref}} = 0.6$ mol/l.
Additionally, we add a penalty on input changes for both control inputs, to obtain a smooth control performance.
End of explanation
"""


# lower bounds of the states
mpc.bounds['lower', '_x', 'C_a'] = 0.1
mpc.bounds['lower', '_x', 'C_b'] = 0.1
mpc.bounds['lower', '_x', 'T_R'] = 50
mpc.bounds['lower', '_x', 'T_K'] = 50

# upper bounds of the states
mpc.bounds['upper', '_x', 'C_a'] = 2
mpc.bounds['upper', '_x', 'C_b'] = 2
mpc.bounds['upper', '_x', 'T_K'] = 140

# lower bounds of the inputs
mpc.bounds['lower', '_u', 'F'] = 5
mpc.bounds['lower', '_u', 'Q_dot'] = -8500

# upper bounds of the inputs
mpc.bounds['upper', '_u', 'F'] = 100
mpc.bounds['upper', '_u', 'Q_dot'] = 0.0

"""
Explanation: Constraints
In the next step, the constraints of the control problem are set.
In this case, there are only upper and lower bounds for each state and the input:
End of explanation
"""


mpc.set_nl_cons('T_R', _x['T_R'], ub=140, soft_constraint=True, penalty_term_cons=1e2)

"""
Explanation: If a constraint is not critical, it is possible to implement it as a soft constraint.
This means, that a small violation of the constraint does not render the optimization infeasible.
Instead, a penalty term is added to the objective.
Soft constraints can always be applied, if small violations can be accepted and it might even be necessary to apply MPC in a safe way (by avoiding avoiding numerical instabilities).
In this case, we define the upper bounds of the reactor temperature as a soft constraint by using mpc.set_nl_cons().
End of explanation
"""


alpha_var = np.array([1., 1.05, 0.95])
beta_var = np.array([1., 1.1, 0.9])

mpc.set_uncertainty_values(alpha = alpha_var, beta = beta_var)

"""
Explanation: Uncertain values
The explicit values of the two uncertain parameters $\alpha$ and $\beta$, which are considered in the scenario tree, are given by:
End of explanation
"""


mpc.setup()

"""
Explanation: This means with n_robust=1, that 9 different scenarios are considered.
The setup of the MPC controller is concluded by:
End of explanation
"""


estimator = do_mpc.estimator.StateFeedback(model)

"""
Explanation: Estimator
We assume, that all states can be directly measured (state-feedback):
End of explanation
"""


simulator = do_mpc.simulator.Simulator(model)

"""
Explanation: Simulator
To create a simulator in order to run the MPC in a closed-loop, we create an instance of the do-mpc simulator which is based on the same model:
End of explanation
"""


params_simulator = {
    'integration_tool': 'cvodes',
    'abstol': 1e-10,
    'reltol': 1e-10,
    't_step': 0.005
}

simulator.set_param(**params_simulator)

"""
Explanation: For the simulation, we use the same time step t_step as for the optimizer:
End of explanation
"""


p_num = simulator.get_p_template()
tvp_num = simulator.get_tvp_template()

"""
Explanation: Realizations of uncertain parameters
For the simulatiom, it is necessary to define the numerical realizations of the uncertain parameters in p_num and the time-varying parameters in tvp_num.
First, we get the structure of the uncertain and time-varying parameters:
End of explanation
"""


# function for time-varying parameters
def tvp_fun(t_now):
    return tvp_num

# uncertain parameters
p_num['alpha'] = 1
p_num['beta'] = 1
def p_fun(t_now):
    return p_num

"""
Explanation: We define two functions which are called in each simulation step, which return the current realizations of the parameters, with respect to defined inputs (in this case t_now):
End of explanation
"""


simulator.set_tvp_fun(tvp_fun)
simulator.set_p_fun(p_fun)

"""
Explanation: These two custum functions are used in the simulation via:
End of explanation
"""


simulator.setup()

"""
Explanation: By defining p_fun as above, the function will always return the value 1.0 for both $\alpha$ and $\beta$.
To finish the configuration of the simulator, call:
End of explanation
"""


# Set the initial state of mpc, simulator and estimator:
C_a_0 = 0.8 # This is the initial concentration inside the tank [mol/l]
C_b_0 = 0.5 # This is the controlled variable [mol/l]
T_R_0 = 134.14 #[C]
T_K_0 = 130.0 #[C]
x0 = np.array([C_a_0, C_b_0, T_R_0, T_K_0]).reshape(-1,1)

mpc.x0 = x0
simulator.x0 = x0
estimator.x0 = x0

mpc.set_initial_guess()

"""
Explanation: Closed-loop simulation
For the simulation of the MPC configured for the CSTR, we inspect the file main.py.
We define the initial state of the system and set it for all parts of the closed-loop configuration:
End of explanation
"""


%%capture
for k in range(50):
    u0 = mpc.make_step(x0)
    y_next = simulator.make_step(u0)
    x0 = estimator.make_step(y_next)

"""
Explanation: Now, we simulate the closed-loop for 50 steps (and suppress the output of the cell with the magic command %%capture):
End of explanation
"""


mpc_graphics = do_mpc.graphics.Graphics(mpc.data)

"""
Explanation: Animating the results
To animate the results, we first configure the do-mpc graphics object, which is initiated with the respective data object:
End of explanation
"""


from matplotlib import rcParams
rcParams['axes.grid'] = True
rcParams['font.size'] = 18

"""
Explanation: We quickly configure Matplotlib.
End of explanation
"""


%%capture
fig, ax = plt.subplots(5, sharex=True, figsize=(16,12))
# Configure plot:
mpc_graphics.add_line(var_type='_x', var_name='C_a', axis=ax[0])
mpc_graphics.add_line(var_type='_x', var_name='C_b', axis=ax[0])
mpc_graphics.add_line(var_type='_x', var_name='T_R', axis=ax[1])
mpc_graphics.add_line(var_type='_x', var_name='T_K', axis=ax[1])
mpc_graphics.add_line(var_type='_aux', var_name='T_dif', axis=ax[2])
mpc_graphics.add_line(var_type='_u', var_name='Q_dot', axis=ax[3])
mpc_graphics.add_line(var_type='_u', var_name='F', axis=ax[4])
ax[0].set_ylabel('c [mol/l]')
ax[1].set_ylabel('T [K]')
ax[2].set_ylabel('$\Delta$ T [K]')
ax[3].set_ylabel('Q [kW]')
ax[4].set_ylabel('Flow [l/h]')
ax[4].set_xlabel('time [h]')

"""
Explanation: We then create a figure, configure which lines to plot on which axis and add labels.
End of explanation
"""


# Update properties for all prediction lines:
for line_i in mpc_graphics.pred_lines.full:
    line_i.set_linewidth(2)
# Highlight nominal case:
for line_i in np.sum(mpc_graphics.pred_lines['_x', :, :,0]):
    line_i.set_linewidth(5)
for line_i in np.sum(mpc_graphics.pred_lines['_u', :, :,0]):
    line_i.set_linewidth(5)
for line_i in np.sum(mpc_graphics.pred_lines['_aux', :, :,0]):
    line_i.set_linewidth(5)

# Add labels
label_lines = mpc_graphics.result_lines['_x', 'C_a']+mpc_graphics.result_lines['_x', 'C_b']
ax[0].legend(label_lines, ['C_a', 'C_b'])
label_lines = mpc_graphics.result_lines['_x', 'T_R']+mpc_graphics.result_lines['_x', 'T_K']
ax[1].legend(label_lines, ['T_R', 'T_K'])

fig.align_ylabels()

"""
Explanation: Some "cosmetic" modifications are easily achieved with the structure pred_lines and result_lines.
End of explanation
"""


from matplotlib.animation import FuncAnimation, ImageMagickWriter

"""
Explanation: After importing the necessary package:
End of explanation
"""


def update(t_ind):
    print('Writing frame: {}.'.format(t_ind), end='\r')
    mpc_graphics.plot_results(t_ind=t_ind)
    mpc_graphics.plot_predictions(t_ind=t_ind)
    mpc_graphics.reset_axes()
    lines = mpc_graphics.result_lines.full
    return lines

n_steps = mpc.data['_time'].shape[0]


anim = FuncAnimation(fig, update, frames=n_steps, blit=True)

gif_writer = ImageMagickWriter(fps=5)
anim.save('anim_CSTR.gif', writer=gif_writer)

"""
Explanation: We obtain the animation with:
End of explanation
"""


