%% Master project, continuation of my Specialization project
% (simulation of three tanks in series)
% Student: D. Alexander L. Skaug, Main supervisor: Sigurd Skogestad,
%                                 co-supervisor: Cristina Zotica 
clear; clc; close all

% MPC controller

addpath('C:\Users\Alexander\Dropbox\Alexander\NTNU\Prosess-systemteknikk\MPC\CasADi')
import casadi.*

% Setting the system
T = 600; % Simulation time, minutes -- CHANGE SIMULATION TIME HERE --
N_interval = 6; % Number of intervals per minute
N = T*N_interval ; % Number of sampling intervals
dt = T/N; % Sampling time
Horizon = 60; % Horizon for the MPC, in minutes -- CHANGE HORIZON HERE --
Tpc = Horizon*N_interval;
% h = dt; % Scaling for M-matrix for collocation

% Nominal point and setpoint for cost-function 
h10 = 1;   h1_SP = 1; % [m]
h20 = 1.2; h2_SP = 1.2;  
h30 = 0.7; h3_SP = 0.7;
x0 = [h10;h20;h30]; x_SP = [h1_SP; h2_SP; h3_SP];

z0_0 = 0.4; % Valve openings [-]
z1_0 = 0.5;
z2_0 = 0.7;
z3_0 = 0.6;
Z0 = [z0_0; z1_0; z2_0; z3_0];

% Choose between linear or nonlinear modelling of flow
 % 1: Linear; flow dependent on valve position only
 % 0: Nonliner; flow dependent on valve position and liquid height in a tank
Linear = 1; 
 
% For plotting
H1 = []; 
H2 = [];
H3 = [];
U_inputs = [];


%% Orthogonal collocation, Gauss-Radau with 3 collocation points
first = 0.155051; second = 0.644949; third = 1;
M_1 = [first  1/2*first^2  1/3*first^3;
       second 1/2*second^2 1/3*second^3;
       third  1/2*third^2  1/3*third^3];
M_2 = [1 first first^2; 1 second second^2; 1 third third^2];
M = M_1*M_2^(-1);

% Parameters
rho = 1000; % Density for fluid, kg/m^3
gravity = 9.81; % Gravitational constant, m/s^2
A1 = 5; % Bottom area of tanks, m^2
A2 = 5; % m^2
A3 = 5; % m^2
Cv_0 = 1.7889*10^(-3); % Valve constants
Cv_1 = 3.2308*10^(-3);
Cv_2 = 2.1067*10^(-3);
Cv_3 = 3.2180*10^(-3);
DeltaP_0 = 50000; % [Pa], constant pressure drop over z0
Delta_U_max = 1; % [-], (hard) input use constraint, value between 0-1
                 % for change in valve positions. Value of 1 will mean that
                 % a valve *can* change over the whole region, but how much
                 % will depend on the cost in the opbjective function J.

w={};
w0 = [];
lbw = [];
ubw = [];
J = 0;
g={};
lbg = [];
ubg = [];

%% Weights in cost function
        % Weight for level setpoints
        Q = 30; %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
        
        % Weight for input use
        R = 20;
        
        % Weight for valve position
          W_zi_setpoint = 0.01; %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%      
        % Choose which valve is the TPM
        Z_man = MX.sym('Z_man', 4);
        w = {w{:}, Z_man}; %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
        lbw = [lbw; 1; 0; 0; 0]; % z0,z1,z2,z3
        ubw = [ubw; 1; 0; 0; 0];
        w0  = [w0;  1; 0; 0; 0];
    
        
        zi_setpoint = MX.sym('zi_setpoint', 4);
        w = {w{:}, zi_setpoint};
        lbw = [lbw; Z0]; % z0,z1,z2,z3
        ubw = [ubw; Z0];
        w0  = [w0;  Z0];

%% Initial condition (nominal point) for MPC
Uk = MX.sym(['U_' num2str(0)], 4);
    w = {w{:}, Uk}; 
    lbw = [lbw; Z0(1); Z0(2); Z0(3); Z0(4)];
    ubw = [ubw; Z0(1); Z0(2); Z0(3); Z0(4)]; 
    w0  = [w0;  Z0(1); Z0(2); Z0(3); Z0(4)];
    
    x1 = MX.sym(['x1' num2str(0)], 3);
    w = [w, {x1}];                        
    lbw = [lbw; h10;h10;h10];
    ubw = [ubw; h10;h10;h10];
    w0  = [w0;  h10;h10;h10];
    
    x2 = MX.sym(['x2' num2str(0)], 3); 
    w = [w, {x2}];
    lbw = [lbw; h20;h20;h20];
    ubw = [ubw; h20;h20;h20];
    w0  = [w0;  h20;h20;h20];
    
    x3 = MX.sym(['x3' num2str(0)], 3); 
    w = [w, {x3}];
    lbw = [lbw; h30;h30;h30];
    ubw = [ubw; h30;h30;h30];
    w0  = [w0;  h30;h30;h30];

%% Creating the MPC problem formulation with cost function
for k=1:Tpc-1
    Uk = MX.sym(['U_' num2str(k)], 4);
    w = {w{:}, Uk};
    lbw = [lbw; 0;0;0;0];
    ubw = [ubw; 1;1;1;1]; 
    w0 = [w0; z0_0; z1_0; z2_0; z3_0];
    
    x1 = MX.sym(['x1_' num2str(k)], 3);
    w = [w, {x1}];                        
    lbw = [lbw; 0;0;0];
    ubw = [ubw; 2;2;2];
    w0 = [w0; h10;h10;h10];
    
    x2 = MX.sym(['x2_' num2str(k)], 3); 
    w = [w, {x2}];
    lbw = [lbw; 0;0;0];
    ubw = [ubw; 2.4;2.4;2.4];
    w0 = [w0; h20;h20;h20];
    
    x3 = MX.sym(['x3_' num2str(k)], 3); 
    w = [w, {x3}];
    lbw = [lbw; 0;0;0];
    ubw = [ubw; 1.4;1.4;1.4];
    w0 = [w0; h30;h30;h30];
    
    %% Orthogonal collocation
    switch Linear
        case 0 % Nonlinear modeling of flow
    % Derivatives of the level in the first tank
    x1_1_dot = 1/A1*( Cv_0*Uk(1)*sqrt(DeltaP_0) - Cv_1*Uk(2)*sqrt(rho*gravity*x1(1)));
    x1_2_dot = 1/A1*( Cv_0*Uk(1)*sqrt(DeltaP_0) - Cv_1*Uk(2)*sqrt(rho*gravity*x1(2)));
    x1_3_dot = 1/A1*( Cv_0*Uk(1)*sqrt(DeltaP_0) - Cv_1*Uk(2)*sqrt(rho*gravity*x1(3)));
    
    % Derivatives of the level in the second tank
    x2_1_dot = 1/A2*( Cv_1*Uk(2)*sqrt(rho*gravity*x1(1)) - Cv_2*Uk(3)*sqrt(rho*gravity*x2(1)));
    x2_2_dot = 1/A2*( Cv_1*Uk(2)*sqrt(rho*gravity*x1(2)) - Cv_2*Uk(3)*sqrt(rho*gravity*x2(2)));
    x2_3_dot = 1/A2*( Cv_1*Uk(2)*sqrt(rho*gravity*x1(3)) - Cv_2*Uk(3)*sqrt(rho*gravity*x2(3)));
    
    % Derivatives of the level in the third tank
    x3_1_dot = 1/A3*( Cv_2*Uk(3)*sqrt(rho*gravity*x2(1)) - Cv_3*Uk(4)*sqrt(rho*gravity*x3(1)));
    x3_2_dot = 1/A3*( Cv_2*Uk(3)*sqrt(rho*gravity*x2(2)) - Cv_3*Uk(4)*sqrt(rho*gravity*x3(2)));
    x3_3_dot = 1/A3*( Cv_2*Uk(3)*sqrt(rho*gravity*x2(3)) - Cv_3*Uk(4)*sqrt(rho*gravity*x3(3)));
    
    % Inequality constraint
    M1_1 = -x1(1) + w{1,(k-1)*4+4}(3) + M(1,1)*x1_1_dot + M(1,2)*x1_2_dot + M(1,3)*x1_3_dot;
    M1_2 = -x1(2) + w{1,(k-1)*4+4}(3) + M(2,1)*x1_1_dot + M(2,2)*x1_2_dot + M(2,3)*x1_3_dot;
    M1_3 = -x1(3) + w{1,(k-1)*4+4}(3) + M(3,1)*x1_1_dot + M(3,2)*x1_2_dot + M(3,3)*x1_3_dot;
    
    M2_1 = -x2(1) + w{1,(k-1)*4+5}(3) + M(1,1)*x2_1_dot + M(1,2)*x2_2_dot + M(1,3)*x2_3_dot;
    M2_2 = -x2(2) + w{1,(k-1)*4+5}(3) + M(2,1)*x2_1_dot + M(2,2)*x2_2_dot + M(2,3)*x2_3_dot;
    M2_3 = -x2(3) + w{1,(k-1)*4+5}(3) + M(3,1)*x2_1_dot + M(3,2)*x2_2_dot + M(3,3)*x2_3_dot;
    
    M3_1 = -x3(1) + w{1,(k-1)*4+6}(3) + M(1,1)*x3_1_dot + M(1,2)*x3_2_dot + M(1,3)*x3_3_dot;
    M3_2 = -x3(2) + w{1,(k-1)*4+6}(3) + M(2,1)*x3_1_dot + M(2,2)*x3_2_dot + M(2,3)*x3_3_dot;
    M3_3 = -x3(3) + w{1,(k-1)*4+6}(3) + M(3,1)*x3_1_dot + M(3,2)*x3_2_dot + M(3,3)*x3_3_dot;
    
    
        case 1 % Linear modeling of flow
    % Changing the modelling of the flow from being dependent on the liquid
    % height in a tank to only being a function of valve position z:
      
    x1_3_dot = 1/A1*( Cv_0*Uk(1)*sqrt(DeltaP_0) - Cv_1*Uk(2)*sqrt(rho*gravity*h10));
    
    x2_3_dot = 1/A2*( Cv_1*Uk(2)*sqrt(rho*gravity*h10) - Cv_2*Uk(3)*sqrt(rho*gravity*h20));
    
    x3_3_dot = 1/A3*( Cv_2*Uk(3)*sqrt(rho*gravity*h20) - Cv_3*Uk(4)*sqrt(rho*gravity*h30));
    
    M1_3 = -x1(3) + w{1,(k-1)*4+4}(3) + x1_3_dot*dt;
    
    M2_3 = -x2(3) + w{1,(k-1)*4+5}(3) + x2_3_dot*dt;
   
    M3_3 = -x3(3) + w{1,(k-1)*4+6}(3) + x3_3_dot*dt;
    
    M1_1 = 0;
    M1_2 = 0;
    M2_1 = 0;
    M2_2 = 0;
    M3_1 = 0;
    M3_2 = 0;
    end
    %%

      D_constr = (Uk(1:4) - w{1,3+(k-1)*4}(1:4)); % Input use constraint 
      % Cost function J
      J = J + Q*((x1(3)-x_SP(1))^2+(x2(3)-x_SP(2))^2+(x3(3)-x_SP(3))^2) + ...
          D_constr(1)^2*R + D_constr(2)^2*R + D_constr(3)^2*R + ...
          (D_constr(4))^2*R + ... 
          W_zi_setpoint*(Uk(1)-w{1,2}(1)).^2*w{1,1}(1) + ...
          W_zi_setpoint*(Uk(2)-w{1,2}(2)).^2*w{1,1}(2) + ...
          W_zi_setpoint*(Uk(3)-w{1,2}(3)).^2*w{1,1}(3) + ...
          W_zi_setpoint*(Uk(4)-w{1,2}(4)).^2*w{1,1}(4);
    g = {g{:}, M1_1, M1_2, M1_3, M2_1, M2_2, M2_3, M3_1, M3_2, M3_3, D_constr};
    lbg = [lbg; 0; 0; 0; 0; 0; 0; 0; 0; 0; -Delta_U_max; -Delta_U_max; -Delta_U_max; -Delta_U_max];
    ubg = [ubg; 0; 0; 0; 0; 0; 0; 0; 0; 0;  Delta_U_max;  Delta_U_max;  Delta_U_max;  Delta_U_max];  
end
% Creating a solver
prob = struct('f', J, 'x', vertcat(w{:}), 'g', vertcat(g{:}));
solver = nlpsol('solver', 'ipopt', prob);

%% Waitbar for the simulation (for practicality)
f = waitbar(0,'Simulating','CreateCancelBtn','setappdata(gcbf,''canceling'',1)');
setappdata(f,'canceling',0);


%% Simulation
for k=1:N 
%% Nominal point
if k == 1
index = 0; % Indexes: 0 = z0, 1 = z1, 2 = z2, 3 = z3
valve_pos = z0_0; % Desired valve position
end
%% Disturbance
% Nominal point: z0_0 = 0.4, z1_0 = 0.5, z2_0 = 0.7, z3_0 = 0.6;
t_disturbance_1 = 20; % The time at which a disturbance occur, in minutes
 if k == t_disturbance_1 * N_interval % Timestep at which the disturbance occur
   index = 0; % Indexes: 0 = z0, 1 = z1, 2 = z2, 3 = z3
   valve_pos = 0.25; % Desired valve position
 end
 t_disturbance_2 = 80; % The time at which a disturbance occur, in minutes
 if k == t_disturbance_2 * N_interval % Timestep at which the disturbance occur
   index = 0; % Indexes: 0 = z0, 1 = z1, 2 = z2, 3 = z3
   valve_pos = 0.5; % Desired valve position
 end
 
%% Implementing the desired valve position/TPM in the cost function
   switch index
       case 0 % Choosing z0 as desired TPM
           w0(1:4)  = [1; 0; 0; 0]; % Only z0 is chosen
           ubw(1:4) = w0(1:4);
           lbw(1:4) = w0(1:4);
           w0(5)  = valve_pos; % Desired valve position of z0
           ubw(5) = w0(5);
           lbw(5) = w0(5);
       case 1 % Choosing z1 as desired TPM
           w0(1:4)  = [0; 1; 0; 0]; % Only z1 is chosen
           ubw(1:4) = w0(1:4);
           lbw(1:4) = w0(1:4);
           w0(6)  = valve_pos; % Desired valve position of z1
           ubw(6) = w0(6);
           lbw(6) = w0(6);
       case 2 % Choosing z2 as desired TPM
           w0(1:4)  = [0; 0; 1; 0]; % Only z2 is chosen
           ubw(1:4) = w0(1:4);
           lbw(1:4) = w0(1:4);
           w0(7)  = valve_pos; % Desired valve position of z2
           ubw(7) = w0(7);
           lbw(7) = w0(7);
       case 3 % Choosing z3 as desired TPM
           w0(1:4)  = [0; 0; 0; 1]; % Only z3 is chosen
           ubw(1:4) = w0(1:4);
           lbw(1:4) = w0(1:4);
           w0(8)  = valve_pos; % Desired valve position of z3
           ubw(8) = w0(8);
           lbw(8) = w0(8);
   end
   
%% Calculating optimal input
sol = solver('x0', w0, 'lbx', lbw, 'ubx', ubw,'lbg', lbg, 'ubg', ubg);
w_opt = full(sol.x);
Cost = full(sol.f);
% pause(0.1) % To see exitflag = OPTIMAL SOLUTION FOUND
% Solution found from ipopt
u_input = w_opt(22:25);
U_inputs = [U_inputs u_input];
lbw(9:12) = u_input; % Updating initial condition in the MPC
ubw(9:12) = u_input;
w0(9:12) =  u_input;
        
%% Saving first solution of MPC to check against "real" system (Plant)
 if k==1
 flows = [];
 states = [];
     for i=1:Tpc
      flows = [flows w_opt((9+(i-1)*13):(12+(i-1)*13))];
      states = [states w_opt(15+(i-1)*13:3:21+(i-1)*13)];
     end
 end
%% Integrating the model
% Simulating the plant behavior during dt
 if k==1
       x = Plant_three_tanks(x0, u_input(1:4), dt); %  qi_setpoint);
       H1 = [H1 x(1)];
       H2 = [H2 x(2)];
       H3 = [H3 x(3)];
       % Updating initial condition of states in MPC
       lbw(13:15) = x(1); ubw(13:15) = x(1);
       lbw(16:18) = x(2); ubw(16:18) = x(2);
       lbw(19:21) = x(3); ubw(19:21) = x(3);
 else
       x = Plant_three_tanks(x, u_input(1:4), dt); % qi_setpoint);
       H1 = [H1 x(1)];
       H2 = [H2 x(2)];
       H3 = [H3 x(3)];
       % Updating initial condition of states in MPC
       lbw(13:15) = x(1); ubw(13:15) = x(1);
       lbw(16:18) = x(2); ubw(16:18) = x(2);
       lbw(19:21) = x(3); ubw(19:21) = x(3);
       
 end
 % Stops the code if "cancel" is pressed on waitbar
    if getappdata(f,'canceling')
        break
    end
 waitbar(k/N,f);
end
delete(f) % Deletes the waitbar
%% Plotting
H1 = [x0(1), H1];
H2 = [x0(2), H2];
H3 = [x0(3), H3];
U_inputs = [Z0 U_inputs];
figure
t = linspace(0, T, N+1);

% Normalizing levels to 0-100 [%]
norm_h1 = 100/2.0;
norm_h2 = 100/2.4;
norm_h3 = 100/1.4;

H1 = H1*norm_h1;
H2 = H2*norm_h2;
H3 = H3*norm_h3;

% h1_setpoint = h1_setpoint*norm_h1;
% h2_setpoint = h2_setpoint*norm_h2;
% h3_setpoint = h3_setpoint*norm_h3;
% h1_setpoint = h1_setpoint*norm_h1;
% h2_setpoint = h2_setpoint*norm_h2;
% h3_setpoint = h3_setpoint*norm_h3;

% subplot(2,1,1) %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
hold on
plot(t, H1(1:end),'r','LineWidth', 2)
plot(t, H2(1:end),'g','LineWidth', 2)
plot(t, H3(1:end),'b','LineWidth', 2)
ylim([0 100])
legend({'h1','h2','h3'},'Location','eastoutside','FontSize',18); grid
% title('Tank levels')
ylabel('Level [%]'); xlabel('Time [min]');
figure
% subplot(2,1,2) %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
hold on
stairs(t, U_inputs(1,1:end),'', 'LineWidth', 2)
stairs(t, U_inputs(2,1:end),'', 'LineWidth', 2)
stairs(t, U_inputs(3,1:end),'', 'LineWidth', 2)
stairs(t, U_inputs(4,1:end),'', 'LineWidth', 2)
% title('Valve openings')
ylabel('Valve [-]'); xlabel('Time [min]');
ylim([0.0 1])
legend({'z0','z1','z2','z3'},'Location','eastoutside','FontSize',18); grid

%% Plotting first solution from MPC
% t = linspace(0, T, Tpc);
% subplot(2,2,2); hold on
% plot(t, states(1,1:end),'r', 'LineWidth', 2)
% plot(t, states(2,1:end),'g', 'LineWidth', 2)
% plot(t, states(3,1:end),'b', 'LineWidth', 2)
% % ylim([3.8 4.2])
% legend({'h1','h2','h3'},'Location','east','FontSize',12); grid
% title('First solution from MPC')
% subplot(2,2,4); hold on
% t = linspace(0, T, Tpc);
% stairs(t, flows(1,1:end), 'LineWidth', 2)
% stairs(t, flows(2,1:end), 'LineWidth', 2)
% stairs(t, flows(3,1:end), 'LineWidth', 2)
% stairs(t, flows(4,1:end), 'LineWidth', 2)
% % ylim([4000 7000])
% legend({'z0','z1','z2','z3'},'Location','southeast','FontSize',12); grid

%% Saving the solution for plotting later
% save('MPC_sol','H1','H2','H3','t','U_inputs')
% save('MPC_sol_new','H1','H2','H3','t','U_inputs')
% save('MPC_sol_1219','H1','H2','H3','t','U_inputs')
% plotting

