% MPC for cooler
% v0: not working yet. Tamal, Cristina and Adriana worked on it 15.Oct.17
% Need to check: 
% 1. Mismatch between "vectorized model" and original model because the steady
% state is not the same.
% 2. Implementation of complementary constraint. 
% v1: Using original model.

clc; clear; close all;
%addpath('C:\Documentos\Casadi\casadi-matlabR2014b-v3.2.3')
addpath('M:\Documents\phd\Software\casadi-matlabR2014b-v3.3.0')
import casadi.*

F_c_max = 10;        % maximum cold fluid mass flow, kg/s
F_h_max = 3;         % maximum hot fluid mass flow, kg/s 
u10=0.5;             % initial opening for cold valve, [-]
u20= 1;              % initial opening for hot valve, [-]
rho_c = 980 ;        % cold fluid density, kg/m3
rho_h = 890;         % hot fluid density,  kg/m3
U = 150*1e-3;        % heat transfer coefficient,   kW/m2/K
A = 90/10;           % heat transfer area, m2
cp_c = 1.5;          % specific heat for cold fluid, kJ/kg/K
cp_h = 1.2;          % specific heat for hot fluid,  kJ/kg/K
Tinh = 70+273.15; % parameter
Tin_c0 = 20 + 273.15; % cold fluid inlet temperature
b = 0.5;      % volume fraction,taken as equal for both shell and tube sides
Aspec = 100;   % specific area to volume, taken from mathisen.phd
V_c = A *b/Aspec; % volume cold side, from mathisen.phd
V_h = A *b/Aspec; % volume hot side, from mathisen.phd

%% casadi definitions

xc = MX.sym('xc',10,1); % x(1)
xh = MX.sym('xh',10,1);
x = [xc;xh]; % Temperatures
u = MX.sym('u',2,1); % control input;  valve opening
F_c = u(1)*F_c_max; % cold side mass flow
F_h = u(2)*F_h_max; % hot side mass flow
Tinc = MX.sym('Tinc',1,1); % disturbance
d = Tinc;

%% model derivatives
xc_bar = [xc(2:10);Tinc]; %inlet to each control volume
xh_bar = [Tinh;xh(1:9)];
xc_dot = F_c/rho_c/V_c*(xc_bar-xc) + (xh-xc)*U*A/rho_c/V_c/cp_c; % cold side
xh_dot = F_h/rho_h/V_h*(xh_bar-xh) - (xh-xc)*U*A/rho_h/V_h/cp_h; % hot side

%% defining casadi function for the model
%f= Function('f', {x,u,d}, {HX_ODE(x,u,d)});
f = Function('f', {x,u,d}, {[xc_dot;xh_dot]});

dae= struct('x',x,'p',[u;d],'ode', f(x,u,d));
op = struct('t0',0,'tf',10);
F = integrator('F','idas',dae,op);

%% Initial values

Tc0 = [314.1411  310.3426  307.0824  304.2841  301.8823  299.8208  298.0514  296.5326  295.2291  294.1102]';
Th0 = [335.2388  328.4495  322.6229  317.6224  313.3307  309.6473  306.4859  303.7724  301.4434  299.4442]';
x0 = [Tc0;Th0]; 
u0 = [u10;u20;Tin_c0]; %u, d

% two alternatives for finding steady state
x01 = fsolve(@(x) HX_ODE(x,u0(1:2),u0(3)),x0);

r=rootfinder('r', 'newton', f);
x02 = full(r(x0, u0(1:2),u0(3)));

%% Steady state
Xk = x0;  Uk = u0;
t_final = 299;
x_int = zeros(20,t_final+1);
x_int(:,1) = Xk;
for i =2:t_final+1
Fk = F('x0',Xk,'p',Uk);
Xk = full(Fk.xf);
x_int(:,i) = Xk;
 if i == 30
    Uk = [0.5;1;295];
 end
end
plot(0:t_final,x_int(:,:))



%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%% MPC

% Collocation

% Degree of interpolating polynomial
d = 3;
% Get collocation points
tau_root = collocation_points(d,'legendre');
tau_root = [0 tau_root];
% Coefficients of the collocation equation
C = zeros(d+1,d+1);
% Coefficients of the continuity equation
D = zeros(d+1, 1);
% Coefficients of the quadrature function
B = zeros(d+1, 1);
% Construct polynomial basis
for j=1:d+1
  % Construct Lagrange polynomials to get the polynomial basis at the collocation point
  coeff = 1;
  for r=1:d+1
    if r ~= j
      coeff = conv(coeff, [1, -tau_root(r)]);
      coeff = coeff / (tau_root(j)-tau_root(r));
    end
  end
  % Evaluate the polynomial at the final time to get the coefficients of the continuity equation
  D(j) = polyval(coeff, 1.0);

  % Evaluate the time derivative of the polynomial at all collocation points to get the coefficients of the continuity equation
  pder = polyder(coeff);
  for r=1:d+1
    C(j,r) = polyval(pder, tau_root(r));
  end

  % Evaluate the integral of the polynomial to get the coefficients of the quadrature function
  pint = polyint(coeff);
  B(j) = polyval(pint, 1.0);
end

% Time horizon
T = 300;

% Control discretization
N = 60; % number of control intervals
h = T/N; % sampling time

% Weighting matrices in the optimization objective
Q = diag([1 1]); %There are two objectives: Thsp and Fhsp
R = diag([1 1]); %For the two inputs
P = Q;

% Start with an empty NLP
w={};
w0 = [];
lbw = [];   %input, slack and states
ubw = [];
J = 0;
g={};
lbg = []; %model
ubg = [];

Th_sp = 26.9+273.15;
Fh_sp = 3;


% "Lift" initial conditions
X0 = MX.sym('X0', 20);
%w = {w{:}, X0};
%lbw = [lbw; x_0];
%ubw = [ubw; x_0];
%w0 = [w0; x_0];
%u_bar = 5;

% Formulate the NLP
Xk = X0;
for k=0:N-1
    % New NLP variable for the control
    Uk = MX.sym(['U_' num2str(k)],2,1);
    w = {w{:}, Uk}; %first uk that goes into the loop
    lbw = [lbw; 0 ; 0]; %input goes from 0 to 1
    ubw = [ubw; 1 ; 1];
    w0 = [w0;  [0.5;1]]; %input initial values
    
    Slack_k = MX.sym(['Slack_' num2str(k)],1,1);
    w = {w{:}, Slack_k};
    lbw = [lbw; -inf];
    ubw = [ubw; inf];
    w0 = [w0;  0];
    
    
    % adding the complementarity constraint for active set changes
       g = {g{:}, (Uk(2)-1)*(Uk(1)-1) - Slack_k}; % slope matching constraint
       lbg = [lbg; 0];
       ubg = [ubg; 0];    

       J = J + 0*Slack_k^2; % weight on the slack variable is 1
    % State at collocation points
    Xkj = {};
    for j=1:d
        Xkj{j} = MX.sym(['X_' num2str(k) '_' num2str(j)], 20,1);
        w = {w{:}, Xkj{j}};
        lbw = [lbw; 290*ones(20,1)]; %lower bounds for temperatures
        ubw = [ubw;  340*ones(20,1)]; %upper bounds for temperatures
        w0 = [w0; x02]; %initial values for temperatures
    end
        
    % Loop over collocation points
    Xk_end = D(1)*Xk;
    for j=1:d
       % Expression for the state derivative at the collocation point
       xp = C(1,j+1)*Xk;
       for r=1:d
           xp = xp + C(r+1,j+1)*Xkj{r};
       end
      
       % Append collocation equations
       fj = f(Xkj{j},Uk,Tinc);
       g = {g{:}, h*fj - xp}; % slope matching constraint
       lbg = [lbg; zeros(20,1)];
       ubg = [ubg; zeros(20,1)];
       
       % Add contribution to the end state
       Xk_end = Xk_end + D(j+1)*Xkj{j};
 
    end    
   
    % New NLP variable for state at end of interval
    Xk = MX.sym(['X_' num2str(k+1)], 20,1);
    w = {w{:}, Xk};
    lbw = [lbw; 290*ones(20,1)];
    ubw = [ubw;  340*ones(20,1)];
    w0 = [w0; x02];

    % Add equality constraint
    g = {g{:}, Xk_end-Xk};  % shooting gap constraint
    lbg = [lbg; zeros(20,1)];
    ubg = [ubg; zeros(20,1)];  
    J = J + ([(Xk(20)-Th_sp)/Th_sp;(Uk(2)*F_h_max - Fh_sp)/Fh_sp])'*Q*([(Xk(20)-Th_sp)/Th_sp;(Uk(2)*F_h_max - Fh_sp)/Fh_sp]); % stage costs
end
J = J + ([(Xk(20)-Th_sp)/Th_sp;(Uk(2)*F_h_max - Fh_sp)/Fh_sp])'*P*([(Xk(20)-Th_sp)/Th_sp;(Uk(2)*F_h_max - Fh_sp)/Fh_sp]); % terminal cost


% Create an NLP solver
prob = struct('f', J, 'x', vertcat(w{:}), 'g', vertcat(g{:}),'p',X0);
solver = nlpsol('solver', 'ipopt', prob);

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% Solve the NLP
t_f =500;
x_0 = x02;
n_mpc = t_f/h;
x_vec = [x_0';zeros(n_mpc-1,20)];
u_vec = zeros(n_mpc,2);
w0=repmat([0.5;1;0;repmat(x_0,4,1)],N,1);
for i = 1:n_mpc
sol = solver('p',x_0,'x0', w0, 'lbx', lbw, 'ubx', ubw,...
            'lbg', lbg, 'ubg', ubg); %open loop soln
w_opt = full(sol.x);
w0 = w_opt;
u = w_opt(1:2);

Fk = F('x0',x_0,'p',[u;Tin_c]); % this is the plant, change disturbance here.
x_plant = full(Fk.xf); % closed loop solution
u_vec(i,:) = u';
x_vec(i+1,:) = x_plant'; %closed loop solution 
x_0 = x_plant;
%x1 = w_opt(8);
%x2 = w_opt(9);
end

% Plot the solution
% x1_opt = w_opt(1:3+2*d:end);
% x2_opt = w_opt(2:3+2*d:end);
% u_opt = w_opt(3:3+2*d:end);

% u_opt = w_opt(1:3+2*d:end);
% x1_opt = [x_0(1);w_opt(8:3+2*d:end)];
% x2_opt = [x_0(2);w_opt(9:3+2*d:end)];

%x_plant = ode45(@(x) ode(x,u),[0 h],x_0);
tgrid = 0:h:t_f;
%tgrid = linspace(0, T, N+1);
 clf;
 figure(1);
 hold on
 plot(tgrid, x_vec(:,1), '--')
 plot(tgrid, x_vec(:,2), '-')
 stairs(tgrid, [u_vec; nan], '-.')
 xlabel('t')
 legend('x1','x2','u')

% % Inspect Jacobian sparsity
% g = vertcat(g{:});
% w = vertcat(w{:});
% Jac = jacobian(g, w);
% figure(2)
% spy(sparse(DM.ones(Jac.sparsity())))
% 
% % Inspect Hessian of the Lagrangian sparsity
% lam = MX.sym('lam', g.sparsity());
% L = J + dot(lam, g);
% Hess = hessian(L, w);
% figure(3)
% spy(sparse(DM.ones(Hess.sparsity())))