clear all
clc

% Addpath to casadi here
import casadi.*

%% Optimizer Model
[par,diff,x_var,p_var,d_var,L,~,f,w_opt_SS] = CSTR_model(0.0);
xf = par.dx0;
dist = [1;0.15];
par.N       = 20;

%% Actual plant Model
par.E2_plant = par.E2;
F_integrator = CSTR_plant(par);

par.ubx(3) = 425;

%% Build the NLP problem with scenario Tree
p_case = [0.0;0.1;0.2]; % different scenarios
nR = 2;                 % Length of robust horizon
[solver,init] = BuildScenMPC(par,f,p_case,nR);

%% simulation

par.nIter = 60;
h = waitbar(0,'NMPC Simulation in Progress...');

x_hat =par.dx0;
u_in = par.u0;
d_hat = dist;

dist1 = [1;p_case(1)];
dist2 = [1;p_case(2)];
dist3 = [1;p_case(3)];

x1_opt_all = [];
x2_opt_all = [];
x3_opt_all = [];
u_opt_all = [];

for sim_k = 1:par.nIter
    waitbar(sim_k/par.nIter)
    
    if sim_k>=30
        dist = [1;0.0];
    end
    % solve NMPC
    nlp_par = vertcat(dist1,u_in(1),x_hat,dist2,u_in(1),x_hat,dist3,u_in(1),x_hat);
    
    tic
    sol         = solver('x0',init.w0,'p',nlp_par,'lbx',init.lbw,'ubx',init.ubw,'lbg',init.lbg,'ubg',init.ubg);
    sim.sol_t(sim_k)    = toc;
    
    flag = solver.stats();
    if (strcmp(flag.return_status,'Solve_Succeeded')~=1)
        disp(['time step ' num2str(sim_k) ': Solver did not converge !!!'])
        %     else
        %         disp(['Time step ' num2str(sim_k) ': Optimal Solution Found'])
    end
    
    % Extract solution
    w_opt = full(sol.x);
    T_in  = w_opt(4);
    init.w0 = w_opt;
    
    n_w_i = 3 + par.N*(13);
    T_in_opt(:,1) = w_opt(4:13:n_w_i);
    x1_opt(:,1) = w_opt([1,5:13:n_w_i]);
    x2_opt(:,1) = w_opt([2,6:13:n_w_i]);
    x3_opt(:,1) = w_opt([3,7:13:n_w_i]);
    
    T_in_opt(:,2) = w_opt(n_w_i+4:13:2*n_w_i);
    x1_opt(:,2) = w_opt([n_w_i+1,n_w_i+5:13:2*n_w_i]);
    x2_opt(:,2) = w_opt([n_w_i+2,n_w_i+6:13:2*n_w_i]);
    x3_opt(:,2) = w_opt([n_w_i+3,n_w_i+7:13:2*n_w_i]);
    
    T_in_opt(:,3) = w_opt(2*n_w_i+4:13:3*n_w_i);
    x1_opt(:,3) = w_opt([2*n_w_i+1,2*n_w_i+5:13:3*n_w_i]);
    x2_opt(:,3) = w_opt([2*n_w_i+2,2*n_w_i+6:13:3*n_w_i]);
    x3_opt(:,3) = w_opt([2*n_w_i+3,2*n_w_i+7:13:3*n_w_i]);
    
    % send first control input to plant
    u_in = [T_in;dist];
    
    % Plant Simulator using integrator
    Fk = F_integrator('x0',xf,'p',u_in);
    xf = full(Fk.xf);
    sim.J(sim_k) = full(Fk.qf);
    sim.y(:,sim_k) = xf;
    sim.u(sim_k) = u_in(1);
    sim.d(:,sim_k) = dist;
    
    ymeas = [sim.y([1:3],sim_k)] ;
    uEKF = [u_in(1)];
    
    x_hat = xf;
    
    
    x1_opt_all = [x1_opt_all,x1_opt];
    x2_opt_all = [x2_opt_all,x2_opt];
    x3_opt_all = [x3_opt_all,x3_opt];
    u_opt_all = [u_opt_all,T_in_opt];
end

sim.x1_opt = x1_opt_all;
sim.x2_opt = x2_opt_all;
sim.x3_opt = x3_opt_all;
sim.u_opt = u_opt_all;

sim.t = 0:10:par.nIter;
save('sim_Central2_nR2.mat','sim')

close(h);
%%
figure(124)

subplot(321)
hold all
plot(sim.y(1:2,:)')
% plot(est.x(1:2,:)','--')
ylabel('concentration')
legend('C_A','C_B')
grid on
box on

subplot(322)
hold all
% plot(sim.Ju)
% ylabel('Gradient')
grid on
box on

subplot(323)
hold all
% plot(sim.Ideal.T_opt,'--')
stairs(sim.u)
grid on
box on
ylabel('Temp In')

subplot(324)
hold all
plot(sim.d')
% plot(est.d(1:2,:)','--')
ylabel('d')
legend('C_{A,in}','C_{B,in}')
grid on
box on

subplot(325)
hold all
plot(sim.y(3,:))
% plot(est.x(3,:)','--')
grid on
box on
ylabel('Temp Out')

subplot(326)
hold all
plot(-sim.J)
% plot(-1*(-(pCB.*sim.y(2,:) - (pTi*sim.y(3,:)).^2)),'--')
ylabel('Cost')
grid on

% figure();
% stem(xEstErr)

