function [D_full, w_fin] = NLP_beamlet_optimizer
% This function performs the beamlet optimization
% Inputs: none. Still needs to find "Geometry" and "beamlets" from Wiscplan
% Outputs: full dose image dose: D_full, optimal beamlet weights: w_fin
%
% Made by Peter Ferjancic 1. May 2018
% Last updated: 1. May 2018
% Inspired by Ana Barrigan's REGGUI optimization procedures

% To-do:
% - Add robusness aspect (+take worst case scenario, see REGGUI)


%% Program starts here
fprintf('starting NLP optimization process: ')

% -- LOAD GEOMETRY AND BEAMLETS --
patient_dir = 'C:\010-work\003_localGit\WiscPlan_v2\data\PatientData';
load([patient_dir '\matlab_files\Geometry.mat'])
beamlet_batch_filename = [patient_dir '\beamlet_batch_files\' 'beamletbatch0.bin'];
beamlet_cell_array = read_ryan_beamlets(beamlet_batch_filename, 'ryan');
fprintf('.')

% -- SET INITIAL PARAMETERS --
numVox  = numel(Geometry.data);
numBeam = size(beamlet_cell_array,2);

% ROI_idx=Geometry.ROIS{1, 2}.ind;
fprintf('.')

%% create input for optimization
% -- BEAMLET DOSE DELIVERY --
beamlets = sparse(numVox, numBeam);
fprintf('.')
wbar1 = waitbar(0, 'Creating beamlet array');
for beam_i=1:numBeam
    % for each beam define how much dose it delivers on each voxel
    idx=beamlet_cell_array{1, beam_i}.non_zero_indices;
    beamlets(idx, beam_i) = 1000*beamlet_cell_array{1, beam_i}.non_zero_values;
    waitbar(beam_i/numBeam)
end
close(wbar1)
fprintf('.')

% -- OPTIMIZATION TARGETS --
optGoal = make_ROI_goals(Geometry, beamlets);

% -- INITIALIZE BEAMLET WEIGHTS --
w0 = ones(numBeam,1);
w0 = mean(optGoal{1}.target ./ (optGoal{1}.beamlets_pruned*w0)) * w0;

% -- MAKE IT ROBUST --
RO_params=0;
optGoal = make_robust_optGoal(optGoal, RO_params);


% -- CALLBACK OPTIMIZATION FUNCTION --
fun = @(x) get_penalty(x, optGoal);

% -- OPTIMIZATION PARAMETERS --
% define optimization parameters
A = [];
b = [];
Aeq = [];
beq = [];
lb = zeros(1, numBeam);
ub = [];
nonlcon = [];

% define opt limits, and make it fmincon progress
options = optimoptions('fmincon');
options.MaxFunctionEvaluations = 30000;
options.Display = 'iter';
options.PlotFcn = 'optimplotfval';
fprintf('-')

%% Run the optimization
tic
w_fin = fmincon(fun,w0,A,b,Aeq,beq,lb,ub,nonlcon,options);
t=toc;
disp(['Optimization run time = ',num2str(t)]);

%% evaluate the results
D_full = reshape(beamlets * w_fin, size(Geometry.data));

%% save outputs
NLP_result.dose = D_full;
NLP_result.weights = w_fin;
save('C:\010-work\003_localGit\WiscPlan_v2\data\PatientData\matlab_files\NLP_result.mat', 'NLP_result');

% plot_DVH(D_full, Geometry, 2, 1)
optGoal_idx=[1,3];
plot_DVH(D_full, optGoal, optGoal_idx)

end

% orthoslice(D_full, [0,70])

%% support functions
% ---- PENALTY FUNCTION 1 ----
function penalty = get_penalty(x, optGoal)
    % this function gets called by the optimizer. It checks the penalty for
    % all the robust implementation and returns the worst result.
    
    NumScenarios = optGoal{1}.NbrRandScenarios * optGoal{1}.NbrSystSetUpScenarios * optGoal{1}.NbrRangeScenarios;
    fobj = zeros(NumScenarios,1);  
    sc_i = 1;
    
    for nrs_i = 1:optGoal{1}.NbrRandScenarios 
        for sss_i = 1:optGoal{1}.NbrSystSetUpScenarios % syst. setup scenarios = ss
            for rrs_i = 1:optGoal{1}.NbrRangeScenarios % range scenario = rs
                fobj(sc_i)=eval_f(x, optGoal, nrs_i, sss_i, rrs_i);
                sc_i = sc_i + 1;
            end
        end
    end
    % take the worst case penalty of evaluated scenarios
    penalty=max(fobj);
end

% ---- Evaluate penalty for single scenario ----
function penalty = eval_f(x, optGoal, nrs_i, sss_i, rrs_i)
    penalty = 0;
    % for each condition
    for goal_i = 1:numel(optGoal)
        switch optGoal{goal_i}.function
            case 'min'
                % penalize if achieved dose is lower than target dose
                d_penalty = 1.0e0 * sum(max(0, ...
                    (optGoal{goal_i}.nrs{nrs_i}.sss{sss_i}.rrs{rrs_i}.target) -...
                    (optGoal{goal_i}.nrs{nrs_i}.sss{sss_i}.rrs{rrs_i}.beamlets_pruned * x)));
            case 'max'
                % penalize if achieved dose is higher than target dose
                d_penalty = 1.0e0 * sum(max(0, ...
                    (optGoal{goal_i}.nrs{nrs_i}.sss{sss_i}.rrs{rrs_i}.beamlets_pruned * x)-...
                    (optGoal{goal_i}.nrs{nrs_i}.sss{sss_i}.rrs{rrs_i}.target)));
            case 'LeastSquare'
                % penalize with sum of squares any deviation from target
                % dose
                d_penalty = 1.0e-1* sum(((...
                    optGoal{goal_i}.nrs{nrs_i}.sss{sss_i}.rrs{rrs_i}.beamlets_pruned * x) - ...
                    optGoal{goal_i}.nrs{nrs_i}.sss{sss_i}.rrs{rrs_i}.target).^2);
        end
        penalty = penalty + d_penalty * optGoal{goal_i}.opt_weight;
    end
end


% ---- DVH PLOT FUNCTION ----
function plot_DVH(dose, optGoal, optGoal_idx)
    % this function plots the DVHs of the given dose
    nfrac=1;
    
    figure
    hold on
    for roi_idx=optGoal_idx
        % plot the histogram
        
        [dvh dosebins] = get_dvh_data(dose,optGoal{roi_idx}.ROI_idx,nfrac);
        plot(dosebins, dvh,'Color', optGoal{roi_idx}.dvh_col,'LineStyle', '-','DisplayName', optGoal{roi_idx}.ROI_name);
    end
    hold off
    
end
function [dvh dosebins] = get_dvh_data(dose,roi_idx,nfrac);
    % this function calculates the data for the DVH
    dosevec = dose(roi_idx);
    
    
    [pdf dosebins] = hist(dosevec, 999);
    % clip negative bins
    pdf = pdf(dosebins >= 0);
    dosebins = dosebins(dosebins >= 0);
    dvh = fliplr(cumsum(fliplr(pdf))) / numel(dosevec) * 100;
    % append the last bin
    dosebins = [dosebins dosebins(end)+0.1];
    dvh = [dvh 0];
    
end

% ---- MAKE ROI GOALS ----
function optGoal = make_ROI_goals(Geometry, beamlets)
    optGoal={};
    
    % -- START DEFINITION OF GOAL --
    goal_1.name = 'Target_min';
    goal_1.ROI_name = Geometry.ROIS{1, 2}.name;
    ROI_idx = Geometry.ROIS{1, 2}.ind;
    goal_1.ROI_idx = ROI_idx;
    goal_1.D_final = 60;
    goal_1.function = 'min';
    goal_1.beamlets_pruned = beamlets(ROI_idx, :);
    goal_1.target   = ones(numel(ROI_idx), 1) * goal_1.D_final;
    goal_1.opt_weight = 1 / numel(ROI_idx); % normalize to volume of target area
    goal_1.dvh_col = [0.9, 0.2, 0.2]; % color of the final DVH plot
    % assign target
    optGoal{end+1}=goal_1;
    % -- END DEFINITION OF GOAL --
    
    % -- START DEFINITION OF GOAL --
    goal_2.name = 'Target_max';
    goal_2.ROI_name = Geometry.ROIS{1, 2}.name;
    ROI_idx = Geometry.ROIS{1, 2}.ind;
    goal_2.ROI_idx = ROI_idx;
    goal_2.D_final = 65;
    goal_2.function = 'max';
    goal_2.beamlets_pruned = beamlets(ROI_idx, :);
    goal_2.target   = ones(numel(ROI_idx), 1) * goal_2.D_final;
    goal_2.opt_weight = 1 / numel(ROI_idx); % normalize to volume of target area
    goal_2.dvh_col = [0.9, 0.2, 0.2]; % color of the final DVH plot
    % assign target
    optGoal{end+1}=goal_2;
    % -- END DEFINITION OF GOAL --
    
    % -- START DEFINITION OF GOAL --
    goal_3.name = 'ROI 1_max';
    goal_3.ROI_name = Geometry.ROIS{1, 3}.name;
    ROI_idx = Geometry.ROIS{1, 3}.ind;
    goal_3.ROI_idx = ROI_idx;
    goal_3.D_final = 0;
    goal_3.function = 'LeastSquare';
    goal_3.beamlets_pruned = beamlets(ROI_idx, :);
    goal_3.target   = ones(numel(ROI_idx), 1) * goal_3.D_final;
    goal_3.opt_weight = 1 / numel(ROI_idx); % normalize to volume of target area
    goal_3.dvh_col = [0.2, 0.9, 0.2]; % color of the final DVH plot
    % assign target
    optGoal{end+1}=goal_3;
    % -- END DEFINITION OF GOAL --
    
end
% ---- MAKE ROI ROBUST ----
function optGoal = make_robust_optGoal(optGoal, RO_params);
    % take regular optimal goal and translate it into several robust cases
    
    % nrs - random scenarios
    % sss - system setup scenarios
    % rrs - random range scenarios
    nrs_scene_list={[0,0,0]};
    sss_scene_list={[0,0,0], [5,5,5], [-5, -5, -5]};
    rrs_scene_list={[0,0,0]};
    
    for i = 1:numel(optGoal)
        optGoal{i}.NbrRandScenarios     =numel(nrs_scene_list);
        optGoal{i}.NbrSystSetUpScenarios=numel(sss_scene_list);
        optGoal{i}.NbrRangeScenarios    =numel(rrs_scene_list);
    end
    

    for goal_i = 1:numel(optGoal)
        out_01.beamlets_pruned  = optGoal{goal_i}.beamlets_pruned;
        out_01.target           = optGoal{goal_i}.target;
        
        for nrs_i = 1:optGoal{goal_i}.NbrRandScenarios          % num. of random scenarios
            
            out_02.beamlets_pruned  = out_01.beamlets_pruned;
            out_02.target           = out_01.target;
            
            for sss_i = 1:optGoal{goal_i}.NbrSystSetUpScenarios   % syst. setup scenarios = sss
                
                out_03 = get_RO_sss(out_02);
                
                for rrs_i = 1:optGoal{goal_i}.NbrRangeScenarios   % range scenario = rrs
                    
                    out_final.beamlets_pruned = out_03.beamlets_pruned;
                    out_final.target = out_03.target;
                    
                    optGoal{goal_i}.nrs{nrs_i}.sss{sss_i}.rrs{rrs_i}.beamlets_pruned   = optGoal{goal_i}.beamlets_pruned;
                    optGoal{goal_i}.nrs{nrs_i}.sss{sss_i}.rrs{rrs_i}.target     = optGoal{goal_i}.target;
                end
            end
        end
    end
end

function out_03 = get_RO_sss(out_02);
    % this function returns a systemic perturbation of the original data
    
    %% THIS PART DOES NOT WORK WELL! NEEDS TO BE FIXED SO THAT IT ACTUALLY PERMUTES DATA
%     wbar1 = waitbar(0, 'Creating SSS beamlet array');
%     for beam_i=1:numBeam
%         % for each beam define how much dose it delivers on each voxel
%         idx=beamlet_cell_array{1, beam_i}.non_zero_indices;
%         beamlets(idx, beam_i) = 1000*beamlet_cell_array{1, beam_i}.non_zero_values;
%         waitbar(beam_i/numBeam)
%     end
%     close(wbar1)

    out_03= out_02;
end