%% RAPTOR tutorial on trajectory optimization with user defined sub functions
% This tutorial will explore more advanced features of RAPTOR and its use
% for trajectory optimization.

% It is still in progress of being improved and optimized, and may not
% work in this version of RAPTOR
disp('this tutorial is under construction')
return
%% Simulate nominal ramp-up
%
% First define the environment, load the params structure and modify some
% default values to suit our needs.

% Load and modify default parameters
close all hidden;
clear; run('../RAPTOR_path.m');

[config] = RAPTOR_config;

% modify some parameters
config.grid.rhogrid = linspace(0,1,11); % spatial grid
config.grid.tgrid = [0:0.001:0.1]; % time grid
config.debug.iterplot = 0; % no plot, no disp this time (so we can see the optimizer output)
config.debug.iterdisp = 0; % 
config.numerics.restol = 1e-8; % relax solution tolerance to be a bit faster

% set ECH/ECCD parameters
config.echcd.params.active = true;
config.echcd.params.rdep = [0.4]; % off axis
config.echcd.params.wdep = [0.4];
config.echcd.params.cd_eff = [2]; % off-axis current drive
config.echcd.params.uindices = uint32([2]); % indices in input vectors

config.nbhcd = nbhcd('none'); % no nbhcd

% generate model, params, init, geometry g, kinetic profiles v
[model,params,init,g,v,U0] = build_RAPTOR_model(config);
%%
% Define a "nominal" ramp-up trajectory with constant ramp-rate and
% sudden switch-on of current drive actuators and simulate this nominal
% case.

% Ip 
U0(1,:) = 250e3*ones(size(params.tgrid)); % input Ip trace: ramp from 80kA to 250kA
U0(1,params.tgrid<=0.025) = linspace(80e3,250e3,sum(params.tgrid<=0.025));

% Off-axis power
U0(2,:) = 2e6*ones(size(params.tgrid)); % off-axis power
U0(2,params.tgrid<0.05) = 1;

%% 
% Define initial conditions and run simulation

% Initial conditions
init.Ip0 = U0(1,1);
x0 = RAPTOR_initial_conditions(model,init,g(:,1),v(:,1));

% Run simulation of nominal case
[simres0] = RAPTOR_predictive(x0,g,v,U0,model,params);
out0 = RAPTOR_out(simres0,model,params);

%% Prepare optimization problem
% We need to first define the discrete set of parameters which determine
% the actuator trajectories, via control vector parametrization.

% tknots: 5 linearly spaced knot points.
% must be a vector.
tknots = linspace(0,1,5)'*params.tgrid(end);
params.cvp.tknots = tknots;
% Parameters for optimization: Ipl and Paux
% Store information about actuator trajectories U which should be
% optimized: true or false.
params.cvp.uind = [true true]; % U(1,:)==Ipl - true, U(2,:)== P_aux - true;
% MA, MW for human scale for optimization parameters.
params.cvp.pscale = [1e6 1e6];
% piecewise linear Ip, piecewise constant Paux
params.cvp.psporder = [1 0];

% initial condition for parameters
p0 = [zeros(1,numel(tknots)),zeros(1,numel(tknots)-1)]';

% generate basis functions in time for these settings.
[~,dU_dp,~,nnp] = RAPTOR_cvp(p0,params);
% nnp is the number of parameters for each input trajectory
nnp
% so U(1,:) is parametrized by 5 numbers, and U(2,;) by 4.

% Now fit parameters p to initial U0 by least squares i.e.
% find p such that U0=dU_dp*p
UUP = [squeeze(dU_dp(1,:,:)),squeeze(dU_dp(2,:,:))]'; % fit matrix
p0 = UUP \ reshape(U0',numel(U0),1);
p0 = max(p0,0); % (to avoid roundoff problems)

% checking the result:
p0
% generate U with these p0;
U(1,:) = squeeze(dU_dp(1,:,:))'*p0;
U(2,:) = squeeze(dU_dp(2,:,:))'*p0;

% check whether they match
figure(1); 
clf;subplot(211); plot(params.tgrid,U(1,:),'b-',params.tgrid,U0(1,:),'r--'); %compare
    title('I_p'); xlabel('time [s]'); ylabel('[A]')
    subplot(212); plot(params.tgrid,U(2,:),'b-',params.tgrid,U0(2,:),'r--'); %compare
    title('P_{aux}'); xlabel('time [s]'); ylabel('W')
    
%% Preparation of the p vector for optimizaiton procedure.
% p vector can contain parameters fixed (for example, initial and final
% value of optimization paramerts) and varying during optimization.
% The goal is to construct a parameter pred (i.e p reduced), which contains
% only varying parameters, and give it as input argument to optimization
% solver, i.e. fmincon.
% In this example numel(p) = 7. Fixed parameters - Ip and P_aux at t=tknots(1) and
% t=tknots(end), varying(i.e. optimizing) parameters - Ip at t(2),t(3),
% t(4) and P_aux at t(2), t(3), since P_aux is piecewise constant
% parameter.

% Define interval where Ip is fixed. First element of the array is always
% t_1, last element - t_end according to RAPTOR time grid, it is necessary
% fot the following interpolation to tknots grid.
% Ip is fixed only at t_1 and t_end.
t_Ipfixed = [params.tgrid(1) params.tgrid(1)+1e-6 params.tgrid(end)-1e-6 params.tgrid(end)];
Ipfixed = [1 0 0 1];
% In the same way define interval where Paux is fixed.
t_Pauxfixed = [params.tgrid(1) params.tgrid(1)+1e-6 params.tgrid(end)-1e-6 params.tgrid(end)];
Pauxfixed = [1 0 0 1];

%%
% Interpolation to tknots grid. This parameter shows if Ip is fixed/free
% at each tknot.
tknots_Ipfixed = interp1(t_Ipfixed, Ipfixed,tknots)';
% Number of parameters in pred corresponding to Ip.
nnpred(1) = sum(tknots_Ipfixed==0);

% Interpolation to tknots grid for Paux.
% Piecewise constant Paux -> necessary to reduce tknots grid by one
% element, but not the last one.
tknotsP = [tknots(1:end-2); tknots(end)];
tknots_Pauxfixed = interp1(t_Pauxfixed,Pauxfixed,tknotsP)';
% Number of parameters in pred corresponding to Paux.
nnpred(2) = sum(tknots_Pauxfixed==0);
% 1 - fixed parameter, 0 - varying parameter
p_ind = [tknots_Ipfixed tknots_Pauxfixed]';

% Contruct transition matrix between p (which contains fixed/free elements)
% and pred (which contains free elements).
% pred' = p'*matrix_ppred
np = numel(p0);
matrix_ppred = bsxfun(@times,~p_ind,diag(ones(1,np)));
matrix_ppred = matrix_ppred(:,any(matrix_ppred));
% Construct pred
pred = (p0'*matrix_ppred)';
% Disp matrix_ppred, nnpred, pred
matrix_ppred
nnpred
pred
% Store in params.cvp
params.cvp.transp_matrix = matrix_ppred;

% otherwise if it's OK to use for optimization full p with fixed elements put
% params.cvp.transp_matrix = [] and define below lbou and ubou for fixed
% elements equal to its values. 

%% Set up linear constraints and parameter bounds

% flag to evolve forward sensitivity equation, required for optimization
% problem.
params.numerics.calc_sens = 1; 
params.debug.checkfwsens = true;
% set constraints on parameter vector p
% Linear equality constraints Aeq*p = beq
% Linear inequality constraints Aineq*p <= bineq
Aineq = []; bineq = []; Aeq = []; beq = [];

% lower and upper bounds lbou <= p <= ubou (element-wise)
lbou = zeros(numel(pred),1); ubou = zeros(numel(pred),1); % initialize
lbou(1:nnpred(1)) = 0.05; % 50kA minimum current
lbou(nnpred(1)+1:end) = 1e-6; % 0 power minimum
ubou(1:nnpred(1)) = 0.5; % 500kA maximum limit
ubou(nnpred(1)+1:end) = 2; % 2MW maximum Paux

%%% If full p with fixed elements is used for optimization 
% % Some parameters are fixed: ensure this by specifying equal upper and
% % lower bounds. Could also do this by equality constraints directly.
% pfixmask = false(sum(nnp),1); % mask which is "true" for fixed parameters.
% pfixmask(1) = true; % fix init Ip
% pfixmask(nnp(1)) = true; % fix final Ip
% pfixmask(nnp(1)+1) = true; % fix init power
% pfixmask(sum(nnp)) = true; % fix final power
% 
% % impose fixed elements of p
% lbou(pfixmask) = p0(pfixmask); % impose fixed p0 elements through bounds
% ubou(pfixmask) = p0(pfixmask); % impose fixed p0 elements through bounds

%% Cost function/constraint parameters
% set up optimization function parameters by calling the function
% |OL_cost_constraint| without input arguments
%
OL_params = OL_cost_constraint_with_subfunctions;
disp(OL_params)
% This parameter structure contains "cost" and "constraint" fields with
% separate parameters to choose cost function and nonlinear constraints.

%% 
% We also pass the RAPTOR model, parameter, and initial conditions to the
% OL_params structure, so that the optimization routines know what RAPTOR
% model to simulate.

OL_params.RAPTOR_model  = model;
OL_params.RAPTOR_params = params;

OL_params.x0 = x0;
OL_params.g0 = g;
OL_params.v0 = v;
OL_params.u0 = U0;
OL_params.p0 = p0;

OL_params.check = 0;

%% Defining the SQP solver options
% As a final step, we need to configure the optimization solver as well

% SQP solver options

% A custom plot function handle
h_plotfun = @(p,optimValues,state) OL_plot_function(p,optimValues,state,model,params,OL_params);

% These are the plot functions to use during the optimization run.
PlotFcns= {h_plotfun,@optimplotfval,@optimplotconstrviolation};

% SQP solver options: pass the plot functions, and some other parameters.
opt = optimset('PlotFcns',PlotFcns,...
                'disp','iter',...
                'GradObj','on',...
                'GradConstr','on',...
                'algorithm','sqp',...
                'maxiter',30,...
                'Tolfun',1e-4,...
                'TolX',1e-4);
            
%% Cost function definition
% The cost function is configured through the |OL_params.cost| field
% The general form reads
%
% $J = \nu_{\iota}J_\iota + \nu_{ss}J_{ss} + \nu_{\psi_{oh}}J_{\psi_{oh}} + \nu_{U_{pl}}J_{U_{pl}}$
% 
% where 
%
% * $J_{\iota} = ||\iota(t_f) -- \iota_{target}||_{W_{\iota}}^2$ (proximity to target
% $\iota$ profile)
% * $J_{ss} = ||\frac{\partial U_{pl}}{\partial{\rho}}(t_f)||_{W_{ss}}^2$ Stationarity
% factor.
% * $J_{\psi_{oh}} = \Delta \Psi_{OH}^2$ Ohmic flux swing
% * $J_{U_{pl}} = ||U_{pl}(t_f) -- U_{pl,ref}||_{U_{pl}}^2$ Distance from target loop
% voltage profile.
%
% and each norm is weighted by a matrix $W$.
% In this case we only desire a stationary profile at the end of the simulation, 
% i.e. flat loop voltage profile.
%% User defined cost function
% Cost function parameters
% Target zero final loop voltage (stationary state)
costfun_params.nu_ss1   = 1;
costfun_params.W_ss1 = 1;
costfun_params.iterpointsplot = 0;
costfun_params.it_cost = Inf;
% Cost function
OL_params.cost_function = @(p,simres) cost_function_tutopt(p,simres,model,params,costfun_params);

%% Constraint function definition
% The implemented constraints are on the rotational transform profile or on
% the (edge) loop voltage. They are simply imposed through
% |OL_params.constraint|
% disp(OL_params.constraint);
%
% We choose to only require q>1 at all times
% constraint function parameters
%% User defined constraint function
% Constraint function parameters
constraintfun_params.iotalimit = 1;
constraintfun_params.Li_max = 1.2;
constraintfun_params.Li_min = 0.65;
% constraint relaxation parameter
constraintfun_params.epsilon_iota = 1e-6;    
constraintfun_params.epsilon_Li3edge = 1e-6;
constraintfun_params.constrplot = 0;
constraintfun_params.constrdisp = 0;
% Constraint function
OL_params.constraint_function = @(p,simres) constraint_function_tutopt(p,simres,model,params,constraintfun_params);

%% User defined function for U, g, and v calculation

% During optimizaiton procedure for each new p proposed by optimization
% algorithm parameters U or/and geometry g or/and v have to be changed
% according to a new p. 
OL_params.prep_ugv = @(p) prep_ugv_tutopt(p,OL_params,model,params);

%% Run SQP optimization algorithm

% Define the cost function handle, i.e. calling OL_cost_constraint with
% 'cost' as second argument.
h_costfun = @(p) OL_cost_constraint_with_subfunctions(p,'cost',OL_params);

% Define the constraint function handle, i.e. calling OL_cost_constraints
% with 'constraint' as second argument.
h_constraints = @(p) OL_cost_constraint_with_subfunctions(p,'constraint',OL_params);
% h_constraints = []; % no constraints - uncomment this to have
% unconstrained optimization.

% Launch optimization problem
disp('please wait, optimizing...');
[poptimalred,fval] = fmincon(h_costfun,pred,Aineq,bineq,Aeq,beq,lbou,ubou,h_constraints,opt);

%% compute trajectories corresponding to optimal p
% Construct full p vector with fixed parameters
poptimal = p0 + ((poptimalred' - p0'*matrix_ppred)*matrix_ppred')';
% Construct trajectories for optimized parameters.
[U,dU_dp,ddU_dt_dp] = RAPTOR_cvp(poptimal,params);
% Run RAPTOR with optimized trajectories.
simres_opt = RAPTOR_predictive(x0,g,v,U,model,params);
% Store results.
out_opt = RAPTOR_out(simres_opt,model,params);

%% Compare nominal and optimized trajectories.
disp('optimized case')
RAPTOR_plot_GUI({out0,out_opt});

%% compare loop voltages
figure(100);
subplot(221);
hp(1) = plot(out0.time,out0.U(1,:)/1e3,'k'); hold on;
hp(2) = plot(out_opt.time,out_opt.U(1,:)/1e3,'r');
legend(hp,{'nominal','optimal'},'location','southeast');
title('Ip [kA]'); xlabel('time [s]')

subplot(222);
plot(out0.time,out0.U(2,:)/1e6,'k');hold on;
plot(out_opt.time,out_opt.U(2,:)/1e6,'r');
title('P_{ECCD} [MW]'); xlabel('time [s]')

subplot(223);
plot(out0.rho,out0.upl(:,end),'k');hold on;
plot(out_opt.rho,out_opt.upl(:,end),'r');
ylabel('[V]');xlabel('\rho')
title('Final loop voltage profile')

subplot(224);
plot(out0.rho,out0.q(:,end),'k');hold on;
plot(out_opt.rho,out_opt.q(:,end),'r');
ylabel('q');xlabel('\rho')
title('Final q profile')
