%% RAPTOR tutorial on trajectory optimization
% This tutorial will explore more advanced features of RAPTOR and its use
% for trajectory optimization.

%% 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 = RAPTORmodule('echcd_gaussian');
config.echcd.Configuration.n_units = 1;
config.echcd.Configuration.n_rdep_external = 0;

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

% generate model, params, init, geometry g, kinetic profiles v
[model,params,init,g,v,U] = build_RAPTOR_model(config);
%%

params.echcd.active = true;
params.echcd.rdep = [0.4]; % off axis
params.echcd.wdep = [0.4];
params.echcd.cd_eff = [2]; % off-axis current drive

% 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) = 0;

%% 
% 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;
params.cvp.pscale = [1e6 1e6]; % MA, MW for human scale.
params.cvp.psporder = [1 0]; % piecewise linear Ip, piecewise constant Paux
params.cvp.uind = [1 1];
params.cvp.transp_matrix = [];

% 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
disp(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')
    
%% Set up linear constraints and parameter bounds

% flag to evolve forward sensitivity equation, required for optimization
% problem.
params.numerics.calc_sens = 1; 
% 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(p0),1); ubou = zeros(numel(p0),1); % initialize
lbou(1:nnp(1)) = 0.05; % 50kA minimum current
lbou(nnp(1)+1:end) = 1e-6; % 0 power minimum
ubou(1:nnp(1)) = 0.5; % 500kA maximum limit
ubou(nnp(1)+1:end) = 2; % 2MW maximum Paux

%%
% 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;
disp(OL_params)
% This parameter structure contains "cost" and "constraint" fields with
% separate parameters to choose cost function and nonlinear constraints.

%% 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$.
disp(OL_params.cost)
%% 
% In this case we only desire a stationary profile at the end of the simulation, 
% i.e. flat loop voltage profile.

% cost function parameters
OL_params.cost.nu_iota  = 0;
OL_params.cost.nu_ss1   = 1; % target zero final loop voltage (stationary state)
OL_params.cost.nu_psioh = 0;
OL_params.cost.nu_upl   = 0; 

%% Constraint 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
OL_params.constraint.iotalimit = 1;
OL_params.constraint.Uedgelimit = [];
OL_params.constraint.constrdisp = 0;

%% 
% 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.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);

%% 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(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(p,'constraint',OL_params);
% h_constraints = []; % no constraints - uncomment this to have
% unconstrained optimization.

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

% compute trajectories corresponding to optimal p
[U,dU_dp,ddU_dt_dp] = RAPTOR_cvp(poptimal,params);
simres_opt = RAPTOR_predictive(x0,g,v,U,model,params);
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
clf;
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')
