
% Clearing of the workspace and closing of all windows
clc; clear; close all;

% Definition of the current path folder
FileName = mfilename('fullpath');
[directory,~,~] = fileparts(FileName);
[dirmain,~,~] = fileparts(directory);

% Import of CASADI
addpath([dirmain '\Util1. Casadi3.1.0'])
addpath([dirmain '\Util2. Functions'])
addpath(genpath([dirmain,'\Util3. Splinter 3.0']));
setup
import casadi.*

% Load of the model
[func,dae,d,lb,ub,n,Cineq,ToutFun] = AmmonReac_NewDist(1);

% Definition of the number of measurements in the QP approach
n.y_used = 1;

% Definition of decision variables
dec.plot = true;
dec.H    = 1;
dec.tra_opt  = true;
dec.tra_save = false;
dec.val_opt  = false;
dec.val_run  = true;
dec.val_save = false;
dec.surr  = true;
dec.Htype = 'In_Out'; % Norm_2, Norm_1, Norm_inf, Max_Gain, In_Out, In
dec.QP    = true;

% Definition of nominal disturbance variables
in_nom = zeros(6,1);
in_nom(1) = 252/3.6;        % Feed flow (kg/s)
in_nom(2) = 200;            % Pressure [bar]
in_nom(3) = 250;            % Temperature [C]
in_nom(4) = 0.08;           % Mass fraction ammonia [-]
in_nom(5) = 3;              % Hydrogen/nitrogen molar ratio [-]
in_nom(6) = 1;              % Catalyst activity 

% Definition of the options for the grid
Var.lbmub = [-10.5 -15 -15 -0.01 -0.2  -0.05;...
               0     0   0  0     0     0;...
              10.5  15  15  0.01  0.2   0.05];
Var.traind = 1:5;
opt.Np     = 3;
opt.Ns     = 3;
opt.val    = 'n';
opt.type   = 'CG';

% Definition of the options for the validation grid
Var.valind   = 1:5;
optval.Np    = 500;
optval.val   = 'y';
optval.type  = 'LHS';

% Definition of the save and load name
if dec.QP
    savename_H = ['H_dist12345_QP_n' num2str(n.y_used) '.mat'];
    loadname_H = ['H_dist12345_QP_n' num2str(n.y_used) '.mat'];
else
    savename_H = ['H_dist12345_' dec.Htype '.mat'];
    loadname_H = ['H_dist12345_' dec.Htype '.mat'];
end

% Load of the initial guesses and the H matrices
load([dirmain '\Data\Init_In.mat'])
if dec.Htype == 'In'
    H{1} = eye(n.u);
    ind.y_red = ind.y_red(1:2:end);
    ind.y_used = ind.y_used(1:2:end);
else
    load(['Local Data\' loadname_H])
end

% Definition of the options for casadi
opts = struct('warn_initial_bounds',false, ...
              'print_time',false, ...
              'ipopt',struct('print_level',1) ...
              );

%% Part a) Root finding for the modified system

% Combination of the variables
w_root = [dae.x;dae.z];
w0     = [x0;z0];

% Combination of the equality constraints and substitution of inlet values
g = vertcat(dae.ode,dae.alg);
g_root = substitute(g,d,zeros(n.d,1));
g_root = substitute(g_root,dae.p(1:n.u),p0(1:n.u));

% Definition and solving of the system
g_fun = Function('g_fun',{w_root},{g_root});
cost = Function('cost',{[w_root;dae.p]},{dae.quad});
G = rootfinder('G','newton',g_fun);
[root.comb] = G('i0',w0);

% Calculation of the number of CSTRs in total
N = size(dae.x,1);

if ~exist('ind','var')
    % Definition of indices used in the solution
    ind.T = [3*N+1 1:N/3       3*N+2 N/3+1:2*N/3   3*N+3  2*N/3+1:N]; % 3*N+4 
    ind.w = [3*N+5 N+1:4*N/3   3*N+6 4*N/3+1:5*N/3 3*N+7  5*N/3+1:2*N];
    ind.R = [3*N+8 2*N+1:7*N/3 3*N+9 7*N/3+1:8*N/3 3*N+10 8*N/3+1:3*N];

    ind.x = 1:N;
    ind.y = [3*N+4 3*N+1 N/3 3*N+2 N*2/3 3*N+3 N];
    ind.z = (N+1):(3*N+10);

    ind.u = (3*N+11):(3*N+13);

    ind.y_used = ind.y;
    ind.y_red = [34 1 11 12 22 23 33];
end

% Reassignment of variables
root.x = full(root.comb.o0(ind.x));
root.z = full(root.comb.o0(ind.z));

root.T = full(root.comb.o0(ind.T));
root.w = full(root.comb.o0(ind.w));

root.f = -full(cost([root.comb.o0;p0;zeros(6,1)]));
x0 = root.x;
z0 = root.z;

% Print of the found results
fprintf('The initialization is performed.\n\n')

%% Part b) Calculation of the training set

% Load of the optimal values
load([dirmain '\Data\Init_Opt.mat'])
y0 = [y0(2); x0(1:10); y0(4); x0(11:20); y0(6); x0(21:30); y0(1)]; 

% Combination of the variables
w_opt = [dae.x; dae.z; dae.p(1:n.u)];
lbw   = [lb.x;  lb.z;  lb.p(1:n.u)];
ubw   = [ub.x;  ub.z;  ub.p(1:n.u)];
w0    = [x0;    z0;    p0(1:n.u)];

% Combination of the equality and inequality constraints
g_opt = [g; Cineq];
lbg  = [zeros(size(g_opt,1)-1,1); -1];
ubg  = [zeros(size(g_opt))];

% Combination of the equality and inequality constraints
g_SOC = [g; H{dec.H}*w_opt(ind.y_used)-H{dec.H}*y0(ind.y_red)];
% g_SOC = [g; w_opt(ind.u)-p0(1:3)];

% Definition of the varying matrix
in_tra = GridDef(Var.lbmub,Var.traind,opt);
n.tra = length(in_tra.out);

% Predefinition of the solution vector
tra.x = zeros(n.x,n.tra);
tra.z = zeros(n.z,n.tra);

tra.T = zeros(N+3,n.tra);
tra.w = zeros(N+3,n.tra);
tra.R = zeros(N+3,n.tra);

tra.p = zeros(n.p,n.tra);
tra.u = zeros(n.u,n.tra);

tra.f    = zeros(1,n.tra);
tra.Tout = zeros(1,n.tra);

time.in.tra = tic;
for k = 1:size(in_tra.out,1)
    if dec.tra_opt
        % Substitution of the disturbance values
        g_sol = substitute(g_opt,d,in_tra.out(k,:));
        J_sol = substitute(dae.quad,d,in_tra.out(k,:));

        % Solving of the non-linear problem
        nlp = struct('x',w_opt,'f',J_sol,'g',g_sol);
        time.in.trait = tic;
        solver = nlpsol('solver','ipopt',nlp,opts);
        tra.comb = solver('x0',w0,'lbx',lbw,'ubx',ubw,'lbg',lbg,'ubg',ubg);
        tra.conv{k,1} = solver.stats.return_status;
        time.trait(k) = toc(time.in.trait);
        
        % Reassignment of the solution vector
        tra.a(:,k) = full(tra.comb.x);
        tra.f(k)   = -full(tra.comb.f);
        
        % Redefinition of the new initial values
        if strcmp(tra.conv{k,1},'Solve_Succeeded')
            w0 = full(tra.comb.x);
            tra.flag(k) = true;
        else
            w0 = [full(root.comb.o0);p0];
            tra.flag(k) = false;
            fprintf('Iteration %4i in the training space failed to converge.\n',k)
        end
    else
        % Substitution of the disturbance values
        g_sol = substitute(g_SOC,d,in_tra.out(k,:));

        % Solving of the non-linear systems of equations
        g_fun = Function('g_fun',{w_opt},{g_sol});
        time.in.trait = tic;
        G = rootfinder('G','newton',g_fun);
        tra.comb = G('i0',w0);
        time.trait(k) = toc(time.in.trait);
        
        % Reassignment of the solution vector
        tra.a(:,k) = full(tra.comb.o0);
        tra.f(k)    = -full(cost([tra.a(:,k); in_tra.out(k,:)']));
    end
    
    % Reassignment of variables
    tra.x(:,k) = tra.a(ind.x,k);
    tra.y(:,k) = tra.a(ind.y_used,k);
    tra.z(:,k) = tra.a(ind.z,k);
    
    tra.T(:,k) = tra.a(ind.T,k);
    tra.w(:,k) = tra.a(ind.w,k);
    tra.R(:,k) = tra.a(ind.R,k);
    
    tra.u(:,k) = tra.a(ind.u,k);
    tra.p(1:n.u,k) = tra.u(:,k);
    
    tra.Tout(k) = full(ToutFun([tra.a(:,k);in_tra.out(k,:)']));
    
    % Redefinition of the new initial values
%     w0 = tra.a(:,k);

    % Print of the progession
    if mod(k,50) == 0
        fprintf('%4i iterations out of %4i in the training space are completed.\n',k,n.tra)
    end
end
time.tra = toc(time.in.tra);

% Redefinition of the new initial values
% w0 = [full(root.comb.o0);p0(1:3)];

% Save of the results
if dec.tra_save
    save(['Local Data\' savename],'SS_opt')
end

% Print of the found results
fprintf('The training space is sampled.\n\n')

plot(tra.u')
[~,asd] = min(tra.u(1,:));
in_tra.grid(asd,:)
sum(tra.u < 0,2)

%% Part c) Sampling of the validation space

if dec.val_run
% Definition of the varying matrix
in_val = GridDef(Var.lbmub,Var.valind,optval);
n.val = length(in_val.out);

% Predefinition of the solution vector
val.x = zeros(n.x,n.val);
val.z = zeros(n.z,n.val);

val.T = zeros(N+3,n.val);
val.w = zeros(N+3,n.val);
val.R = zeros(N+3,n.val);

val.p = zeros(n.p,n.val);
val.u = zeros(n.u,n.val);

val.f    = zeros(1,n.val);
val.Tout = zeros(1,n.val);

time.in.val = tic;
for k = 1:size(in_val.out,1)
    if dec.val_opt
        % Substitution of the disturbance values
        g_sol = substitute(g_opt,d,in_val.out(k,:));
        J_sol = substitute(dae.quad,d,in_val.out(k,:));

        % Solving of the non-linear problem
        nlp = struct('x',w_opt,'f',J_sol,'g',g_sol);
        time.in.valit = tic;
        solver = nlpsol('solver','ipopt',nlp,opts);
        val.comb = solver('x0',w0,'lbx',lbw,'ubx',ubw,'lbg',lbg,'ubg',ubg);
        val.conv{k,1} = solver.stats.return_status;
        time.valit(k) = toc(time.in.valit);
        
        % Reassignment of the solution vector
        val.a(:,k) = full(val.comb.x);
        val.f(k)   = -full(val.comb.f);
        
        % Redefinition of the new initial values
        if strcmp(val.conv{k,1},'Solve_Succeeded')
            w0 = full(val.comb.x);
            val.flag(k) = true;
        else
            w0 = [full(root.comb.o0);p0(1:n.u)];
            val.flag(k) = false;
            fprintf('Iteration %4i in the validation space failed to converge.\n',k)
        end
    else
        % Substitution of the disturbance values
        g_sol = substitute(g_SOC,d,in_val.out(k,:));

        % Solving of the non-linear systems of equations
        g_fun = Function('g_fun',{w_opt},{g_sol});
        time.in.valit = tic;
        G = rootfinder('G','newton',g_fun);
        val.comb = G('i0',w0);
        time.valit(k) = toc(time.in.valit);
        
        % Reassignment of the solution vector
        val.a(:,k) = full(val.comb.o0);
        val.f(k)    = -full(cost([val.a(:,k); in_val.out(k,:)']));
        val.flag(k) = true;
    end
    
    % Reassignment of variables
    val.x(:,k) = val.a(ind.x,k);
    val.y(:,k) = val.a(ind.y_used,k);
    val.z(:,k) = val.a(ind.z,k);
    
    val.T(:,k) = val.a(ind.T,k);
    val.w(:,k) = val.a(ind.w,k);
    val.R(:,k) = val.a(ind.R,k);
    
    val.u(:,k) = val.a(ind.u,k);
    val.p(1:n.u,k) = val.u(:,k);
    
    val.Tout(k) = full(ToutFun([val.a(:,k);in_val.out(k,:)']));
    % Redefinition of the new initial values
%     if nom_H.f(i) ~= 0
%         w0 = full(nom_H.comb.o0);
%     else
%         w0 = [full(root.comb.o0);p0(1:n.u)];
%     end

    % Print of the progession
    if mod(k,50) == 0
        fprintf('%4i iterations out of %4i in the validation space are completed.\n',k,n.val)
    end
    val = rmfield(val,'comb');
    if dec.val_save
        save('Local Data\OptSave_RND5000.mat','val','in_val')
    end
end
time.val = toc(time.in.val);
else
    load('Local Data\OptSave_RND5000.mat')
    n.val = length(in_val.out);
end

% Print of the found results
fprintf('The validation space is sampled.\n\n')

%% Part d) Surrogate model generation

if dec.surr
    time.in.fit = tic;
    surr.EoR  = BSplineBuilder(in_tra.grid, tra.f',      min(opt.Np-1,4)).build;
    surr.Tout = BSplineBuilder(in_tra.grid, tra.Tout',   min(opt.Np-1,4)).build;
    surr.u1   = BSplineBuilder(in_tra.grid, tra.u(1,:)', min(opt.Np-1,4)).build;
    surr.u2   = BSplineBuilder(in_tra.grid, tra.u(2,:)', min(opt.Np-1,4)).build;
    surr.u3   = BSplineBuilder(in_tra.grid, tra.u(3,:)', min(opt.Np-1,4)).build;
    time.fit = toc(time.in.fit);

    surr.f = surr.EoR.eval(in_val.out(val.flag,Var.traind));
    err.EoR  = (surr.f-val.f(val.flag)')./val.f(val.flag)'*100;
    sae.EoR  = sort(abs(err.EoR));
    [mae.EoR,ind.max.EoR] = max(abs(err.EoR));
    [meae.EoR]            = mean(abs(err.EoR));


    surr.T = surr.Tout.eval(in_val.out(val.flag,Var.traind));
    err.Tout = (surr.T-val.Tout(val.flag)')./val.Tout(val.flag)'*100;
    sae.Tout = sort(abs(err.Tout));
    [mae.T,ind.max.T] = max(abs(err.Tout));
    [meae.T]          = mean(abs(err.Tout));


    surr.u(:,1) = surr.u1.eval(in_val.out(val.flag,Var.traind));
    surr.u(:,2) = surr.u2.eval(in_val.out(val.flag,Var.traind));
    surr.u(:,3) = surr.u3.eval(in_val.out(val.flag,Var.traind));
    err.u = (surr.u-val.u(:,val.flag)')./val.u(:,val.flag)'*100;
    [mae.u,ind.max.u] = max(abs(err.u),[],1);
    [meae.u]          = mean(abs(err.u));
    % 
    % ind.y_red
    % H{1}
    % mean(time.trait)
    fprintf('The maximum absolute relative error for the extent of reaction is %2.3f %%.\n',mae.EoR)
    fprintf('The maximum absolute relative error for the outlet temperature is %2.3f %%.\n',mae.T)

    save(['Local Data\Surr_' loadname_H],'err','mae','meae') 
end

%% Part e) Plotting of the results for the different cases

if numel(Var.valind) == 2
    figure(1)
    f_plot = reshape(val.f,[optval.Np optval.Np]);
    A_plot = reshape(in_val.grid(:,1),[optval.Np optval.Np]);
    B_plot = reshape(in_val.grid(:,2),[optval.Np optval.Np]);
    surfc(A_plot,B_plot,f_plot)
end

if numel(Var.traind) == 2
    f_plot = reshape(tra.f,       [opt.Np opt.Np]);
    T_plot = reshape(tra.Tout,    [opt.Np opt.Np]);
    A_plot = reshape(in_tra.grid(:,1),[opt.Np opt.Np]);
    B_plot = reshape(in_tra.grid(:,2),[opt.Np opt.Np]);
    figure(3)
    surfc(A_plot,B_plot,f_plot)
    figure(4)
    surfc(A_plot,B_plot,T_plot)
end

mean(abs((surr.T-val.Tout(val.flag)')))
max(abs((surr.T-val.Tout(val.flag)')))

%% [~,ind.EoR] = sort

% for k = 2:n.val
%     if abs(sae.EoR(k)-mean(sae.EoR)) > abs(sae.EoR(k-1)-mean(sae.EoR))
%         ind.mea.EoR = k-1;
%         break
%     end
% end
% 
% figure(1)
% loglog(sae.EoR,'k')
% hold on
% loglog([0 512 512]',[median(abs(err.EoR))*ones(1,2) 0]','r')
% loglog([0 ind.mea.EoR ind.mea.EoR],[mean(abs(err.EoR))*ones(1,2) 0],'b')
% % xlim([0 2000])
% 
% for k = 2:n.val
%     if abs(sae.Tout(k)-mean(sae.Tout)) > abs(sae.Tout(k-1)-mean(sae.Tout))
%         ind.mea.Tout = k-1;
%         break
%     end
% end
% 
% figure(2)
% loglog(sae.Tout,'k')
% hold on
% loglog([0 512 512],[median(abs(err.Tout))*ones(1,2) 0],'r')
% loglog([0 ind.mea.Tout ind.mea.Tout],[mean(abs(err.Tout))*ones(1,2) 0],'b')
% % xlim([0 2000])