% this is a simple example for a constrained NLP solver from matlabs
% optimization toolbox

function [D_full, w_fin] = NLP_beamlet_optimizer
fprintf('starting NLP optimization process')

% load the data
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
targetROI_ind = 1;
numVox  = numel(Geometry.data);
numBeam = size(beamlet_cell_array,2);

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

%% create input for optimization
% ideal target values
target = sparse(zeros(numVox,1));
target(ROI_idx) = 60; % 60 units of dose to the target, 0 otherwise

% create the map on how beamlets deliver the dose
beamlets = sparse(zeros(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('.')

% prune the beamlets matrix only to areas of interest
pruneIdx= find(target~=0);
beamlets_pruned= beamlets(pruneIdx, :);
target_pruned = full(target(pruneIdx));

% initial beamlet weights
w0 = ones(numBeam,1);
w0 = mean(target_pruned./ (beamlets_pruned*w0)) * w0;

% optimization function - least squares
% fun = @(x) sum(((beamlets_pruned * x) - target_pruned).^2);
fun = @(x) eval_f(x, beamlets_pruned, target_pruned);


% all weights*-1 should be < 0 (b)
A = [];
b = [];
Aeq = [];
beq = [];
lb = zeros(1, numBeam);
ub = [];
nonlcon = [];

% define optimizer options
options = optimoptions('fmincon');
options.MaxFunctionEvaluations = 100000;
options.Display = 'iter';

fprintf('--')
%% Run the optimization
tic
x = 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 * x, size(Geometry.data));



w_fin=x;

%% save outputs

NLP_result.dose = D_full;
NLP_result.weights = w_fin;
% orthoslice(NLP_result.dose, [0,70])

save('C:\010-work\003_localGit\WiscPlan_v2\data\PatientData\matlab_files\NLP_result.mat', 'NLP_result');
end

function penalty = eval_f(x, beamlets, target)
    % sum squares
%     penalty = sum(((beamlets * x) - target).^2);

    % at least>target
    penalty = 0;
    penalty = penalty + sum(max(0, target-(beamlets * x)));
    
    penalty = penalty + sum(max(0, (beamlets * x)-(target+5)));
end