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)

N_fcallback1 = 10000;
N_fcallback2 = 50000;
patient = 'phantom';

switch patient
    case 'phantom'
        patient_dir = 'C:\010-work\003_localGit\WiscPlan_v2\data\PatientData';
        blet_in_beam=5; % this is the number of beamlets in a beam. Called "Mxp" in helicalDosecalcSetup
    case 'phantom_HD'
        patient_dir = 'C:\010-work\003_localGit\WiscPlan_v2\data\PD_HD_dicomPhantom';
        blet_in_beam=20; % ???
    case 'doggo'
        patient_dir = 'C:\010-work\003_localGit\WiscPlan_v2\data\PatientData_dog5_2';
        blet_in_beam=5; % ???
    otherwise
        error('invalid case')
end


%% PROGRAM STARTS HERE
% - no tocar lo que hay debajo -
fprintf('starting NLP optimization process... ')

% -- LOAD GEOMETRY AND BEAMLETS --
load([patient_dir '\matlab_files\Geometry.mat'])
% beamlet_batch_filename = [patient_dir '\beamlet_batch_files\' 'beamletbatch0.bin'];
beamlet_batch_filename = [patient_dir '\' 'batch_dose.bin'];
beamlet_cell_array = read_ryan_beamlets(beamlet_batch_filename, 'ryan');
fprintf('\n loaded beamlets...')

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

%% -- BEAMLET DOSE DELIVERY --
beamlets = get_beamlets(beamlet_cell_array, numVox);
% show_joint_beamlets(beamlets, size(Geometry.data), 7:9)
fprintf('\n beamlet array constructed...')
% - merge beamlets into beams -
load([patient_dir '\all_beams.mat'])
beamletOrigin=[0 0 0];
beam_i=0;
beam_i_list=[];
for beamlet_i = 1:numel(all_beams)
    newBeamletOrigin = all_beams{beamlet_i}.ip;
    if any(newBeamletOrigin ~= beamletOrigin)
        beam_i = beam_i+1;
        beamletOrigin = newBeamletOrigin;
    end
    beam_i_list=[beam_i_list, beam_i];
end
beamlets_joined=[];
target_joined=[];
wbar2 = waitbar(0, 'merging beamlets into beams');
numBeam=numel(unique(beam_i_list));
for beam_i = unique(beam_i_list)
    beam_full = sum(beamlets(:,beam_i_list == beam_i), 2); 
    beamlets_joined(:,end+1) = beam_full;
    waitbar(beam_i/numBeam, wbar2)
end
close(wbar2)


%% -- OPTIMIZATION TARGETS --
switch patient
    case 'phantom'
        optGoal = make_ROI_goals(Geometry, beamlets);
        optGoal_beam = make_ROI_goals(Geometry, beamlets_joined);
        N_beamlets_in_beam = 10;
    case 'phantom_HD'
        optGoal = make_ROI_goals(Geometry, beamlets);
        optGoal_beam = make_ROI_goals(Geometry, beamlets_joined);
    case 'doggo'
        optGoal = make_ROI_goals_DOG(Geometry, beamlets);
        optGoal_beam = make_ROI_goals_DOG(Geometry, beamlets_joined);
    otherwise
        error('invalid case')
end
% -- make them robust --
RO_params=0;
optGoal_beam = make_robust_optGoal(optGoal_beam, RO_params, beamlets_joined);
optGoal = make_robust_optGoal(optGoal, RO_params, beamlets);

%% -- INITIALIZE BEAMLET WEIGHTS --
w0_beams = ones(numBeam,1);
w0_beams = mean(optGoal_beam{1}.target ./ (optGoal_beam{1}.beamlets_pruned*w0_beams+0.1)) * w0_beams;
w0_beams = w0_beams + (2*rand(size(w0_beams))-1) *0.1 .*w0_beams; % random perturbation


% -- CALLBACK OPTIMIZATION FUNCTION --
fun1 = @(x) get_penalty(x, optGoal_beam);
fun2 = @(x) get_penalty(x, optGoal);

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

% define opt limits, and make it fmincon progress
options = optimoptions('fmincon');
options.MaxFunctionEvaluations = N_fcallback1;
options.Display = 'iter';
options.PlotFcn = 'optimplotfval';
fprintf('\n running initial optimizer:')

%% Run the optimization
% -- GET FULL BEAM WEIGHTS --
tic
w_beam = fmincon(fun1,w0_beams,A,b,Aeq,beq,lb_beam,ub,nonlcon,options);
t=toc;
disp(['Optimization time for beams = ',num2str(t)]);
w_beamlets = ones(numBeamlet,1);
numBeam=numel(unique(beam_i_list));
for beam_i = 1:numBeam % assign weights to beamlets
    % beamlets from same beam get same initial weights
    w_beamlets(beam_i_list == beam_i) = w_beam(beam_i);
end
w_beamlets = w_beamlets + (2*rand(size(w_beamlets))-1) *0.1 .*w_beamlets; % small random perturbation

% -- GET FULL BEAMLET WEIGHTS --
options.MaxFunctionEvaluations = N_fcallback2;
tic
w_fin = fmincon(fun2,w_beamlets,A,b,Aeq,beq,lb,ub,nonlcon,options);
t=toc;
disp(['Optimization time for beamlets = ',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');

optGoal_idx=[1,3];
plot_DVH(D_full, optGoal, optGoal_idx)
% plot_DVH_robust(D_full, optGoal, optGoal_idx)
end

% orthoslice(D_full, [0,70])

%% support functions
% ---- PENALTY FUNCTION ----
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
% ------ supp: 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

% ---- GET BEAMLETS ----
function beamlets = get_beamlets(beamlet_cell_array, numVox);
    wbar1 = waitbar(0, 'Creating beamlet array');
    numBeam = size(beamlet_cell_array,2);
    batchSize=100;
    beamlets = sparse(0, 0);
    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;

        % break the beamlets into multiple batches
        if rem(beam_i, batchSize)==1;
            beamlet_batch = sparse(numVox, batchSize);
            beam_i_temp=1;
        end

        beamlet_batch(idx, beam_i_temp) = 1000*beamlet_cell_array{1, beam_i}.non_zero_values;
        waitbar(beam_i/numBeam, wbar1, ['Adding beamlet array: #', num2str(beam_i)])

        % add the batch to full set when filled
        if rem(beam_i, batchSize)==0;
            beamlets =[beamlets, beamlet_batch];
        end
        % crop and add the batch to full set when completed
        if beam_i==numBeam;
            beamlet_batch=beamlet_batch(:, 1:beam_i_temp);
            beamlets =[beamlets, beamlet_batch];
        end
        beam_i_temp=beam_i_temp+1;

    end
    close(wbar1)

end

function show_joint_beamlets(beamlets, IMGsize, Beam_list)
    % this function overlays and plots multiple beamlets. The goal is to
    % check whether some beamlets are part of the same beam manually.
    
    beam=sum(beamlets(:, Beam_list), 2);
    
    beamImg=reshape(full(beam), IMGsize);
        
    orthoslice(beamImg)
    
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.imgDim = size(Geometry.data);
    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.imgDim = size(Geometry.data);
    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 = 0.8 / 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.imgDim = size(Geometry.data);
    goal_3.D_final = 10;
    goal_3.function = 'max';
    goal_3.beamlets_pruned = beamlets(ROI_idx, :);
    goal_3.target   = ones(numel(ROI_idx), 1) * goal_3.D_final;
    goal_3.opt_weight = 2 / 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
function optGoal = make_ROI_goals_DOG(Geometry, beamlets)
    optGoal={};
    
    % -- START DEFINITION OF GOAL --
    goal_1.name = 'Target_min';
    goal_1.ROI_name = Geometry.ROIS{1, 1}.name;
    ROI_idx = Geometry.ROIS{1, 1}.ind;
    goal_1.ROI_idx = ROI_idx;
    goal_1.imgDim = size(Geometry.data);
    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 = 'Doggo_max';
    goal_2.ROI_name = Geometry.ROIS{1, 4}.name;
    ROI_idx = Geometry.ROIS{1, 1}.ind;
    goal_2.ROI_idx = ROI_idx;
    goal_2.imgDim = size(Geometry.data);
    goal_2.D_final = 20;
    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 = 0.2 / 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 --
    
end
% ---- MAKE ROI ROBUST ----
function optGoal = make_robust_optGoal(optGoal, RO_params, beamlets);
    % take regular optimal goal and translate it into several robust cases
    
    % RO_params - should have the information below
    % nrs - random scenarios
    % sss - system setup scenarios
    % rrs - random range scenarios
    
    % X - X>0 moves image right
    % Y - Y>0 moves image down
    % Z - in/out.
    nrs_scene_list={[0,0,0]};
    sss_scene_list={[0,0,0]};
    rrs_scene_list={[0,0,0]};
    
%     nrs_scene_list={[0,0,0]};
    sss_scene_list={[0,0,0],  [-15,0,0], [-10,0,0], [-5,0,0], [5,0,0], [10,0,0], [15,0,0]};
%     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)
        % get target
        idx=optGoal{goal_i}.ROI_idx;
        targetImg1=zeros(optGoal{goal_i}.imgDim);
        targetImg1(idx)=1;
        % get beamlets
        
        for nrs_i = 1:optGoal{goal_i}.NbrRandScenarios          % num. of random scenarios
            % modify target and beamlets
            targetImg2=targetImg1;
            % beamlets stay the same
            
            for sss_i = 1:optGoal{goal_i}.NbrSystSetUpScenarios   % syst. setup scenarios = sss
                % modify target and beamlets
                targetImg3=get_RO_sss(targetImg2, sss_scene_list{sss_i});
                % beamlets stay the same
                
                for rrs_i = 1:optGoal{goal_i}.NbrRangeScenarios   % range scenario = rrs
                    % modify target and beamlets
                    targetImg4=targetImg3;
                    % beamlets stay the same
                    
                    %% make new target and beamlets
                    ROI_idx=[];
                    ROI_idx=find(targetImg4>0);
                    target = ones(numel(ROI_idx), 1) * optGoal{goal_i}.D_final;
                    beamlets_pruned = beamlets(ROI_idx, :);
                    
                    % save to optGoal output
                    optGoal{goal_i}.nrs{nrs_i}.sss{sss_i}.rrs{rrs_i}.ROI_idx            = ROI_idx;
                    optGoal{goal_i}.nrs{nrs_i}.sss{sss_i}.rrs{rrs_i}.beamlets_pruned    = beamlets_pruned;
                    optGoal{goal_i}.nrs{nrs_i}.sss{sss_i}.rrs{rrs_i}.target             = target;
                end
            end
        end
    end
end
% ------ supp: RO case SSS ------
function targetImg3=get_RO_sss(targetImg2, sss_scene_shift);
    % translate the target image 
    targetImg3 = imtranslate(targetImg2,sss_scene_shift);
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 plot_DVH_robust(dose, optGoal, optGoal_idx)
    % this function plots the DVHs of the given dose
    nfrac=1;
    
    figure
    hold on
    for goal_i=optGoal_idx
        % for all the perturbations
        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
                    % plot the histogram
                    ROI_idx = optGoal{goal_i}.nrs{nrs_i}.sss{sss_i}.rrs{rrs_i}.ROI_idx;
                    
                    [dvh dosebins] = get_dvh_data(dose,ROI_idx,nfrac);
                    plot(dosebins, dvh,'Color', optGoal{goal_i}.dvh_col,'LineStyle', '-','DisplayName', optGoal{goal_i}.ROI_name);
                end
            end
        end
    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