% paths.patient_dir
% paths.Goal_dir (previously called DP_dir)
% paths.patient
% paths.goalsName

% colorwash(Geometry.data, D_full, [500, 1500], [0,70])
% orthoslice(D_full, [0,70])

function [D_full, w_fin, Geometry, optGoal] = NLP_optimizer_v2(varargin)
% This function performs the beamlet optimization
% [D_full, w_fin, Geometry, optGoal] = NLP_beamlet_optimizer;
% 
% Inputs:
%       () OR
%       (Pat_path, path2goal) OR
%       (Pat_path, path2goal, beamlet_weights)
%   Pat_path, path2goal = strings to patient folder and optimal goals
%   beamlet_weights = initial beamlet weights
%   
% Outputs:
%       full dose image dose: [D_full, w_fin, Geometry, optGoal]
% 
% Made by Peter Ferjancic 1. May 2018
% Last updated: 1. April 2019

if nargin<2
    load('WiscPlan_preferences.mat')
    [Pat_path] = uigetdir([WiscPlan_preferences.patientDataPath ], 'Select Patient folder');
    [Goal_file,Goal_path,indx] = uigetfile([Pat_path '\matlab_files\*.mat'], 'Select OptGoal file');
    
    path2geometry = [Pat_path, '\matlab_files\Geometry.mat'];
    path2goal = [Goal_path, Goal_file];
else
    Pat_path = varargin{1};
    path2geometry = [Pat_path, '\matlab_files\Geometry.mat'];
    path2goal = varargin{2};
end

str = inputdlg({'N of iterations for initial calc', 'N of iterations for full calc', ...
    'Use pre-existing NLP_result to initiate? (y/n)'}, 'input', [1,35], {'10000', '500000', 'n'});
N_fcallback1 = str2double(str{1}); % 10000  is a good guesstimate
N_fcallback2 = str2double(str{2}); % 500000 is a good guesstimate
pre_beamWeights = str{3};

switch pre_beamWeights
    case 'y'
        [NLP_file,NLP_path,indx] = uigetfile([Pat_path '\matlab_files\*.mat'], 'Select NLP_result file');
        load([NLP_path, NLP_file])
        w_beamlets = NLP_result.weights;
        
        load([Pat_path, '\all_beams.mat'])
        if numel(all_beams) ~= numel(w_beamlets)
            error('Provided weight number does not match beamlet number!')
        end
    case 'n'
        disp('Initial beam weights will be calculated.')
end

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

% % -- LOAD GEOMETRY, GOALS, BEAMLETS --
load(path2geometry)
load(path2goal)
[beamlets, beamlets_joined, numBeamlet, numBeam, beam_i_list] = get_beam_lets(Geometry, Pat_path);



%% -- OPTIMIZATION TARGETS --
% -- make the optimization optGoal structure --

for i_goal = 1:size(OptGoals.goals,1)
    answer = inputdlg(['# of supervoxels for "' OptGoals.data{i_goal}.name '" with ' num2str(numel(OptGoals.data{i_goal}.ROI_idx)) ' vox: ("0" to skip)'])
    
    switch answer{1}
        case '0'
        % if not supervoxel, just select provided ROI_idx
        optGoal{i_goal} = OptGoals.data{i_goal};
        optGoal{i_goal}.beamlets_pruned = sparse(beamlets(optGoal{i_goal}.ROI_idx, :));
        optGoal_beam{i_goal} = OptGoals.data{i_goal};
        optGoal_beam{i_goal}.beamlets_pruned = sparse(beamlets_joined(optGoal{i_goal}.ROI_idx, :));
    
        otherwise
        % -- if supervoxel, merge given columns
        % - make supervoxel map
        mask = zeros(OptGoals.data{i_goal}.imgDim);
        mask(OptGoals.data{i_goal}.ROI_idx) = 1;
        
        superMask = superpix_group(mask, str2double(answer{1})); 
        
        superVoxList = unique(superMask);
        superVoxList = superVoxList(superVoxList>0);
        
        optGoal{i_goal} = OptGoals.data{i_goal};
        optGoal{i_goal}.ROI_idx_old = optGoal{i_goal}.ROI_idx; % copy old index data
        optGoal{i_goal}.ROI_idx = zeros(numel(superVoxList), 1);
        optGoal{i_goal}.opt_weight = optGoal{i_goal}.opt_weight * numel(optGoal{i_goal}.ROI_idx_old)/numel(optGoal{i_goal}.ROI_idx);
        
        optGoal_beam{i_goal} = OptGoals.data{i_goal};
        optGoal_beam{i_goal}.ROI_idx_old = optGoal_beam{i_goal}.ROI_idx; 
        optGoal_beam{i_goal}.ROI_idx = zeros(numel(superVoxList), 1);
        optGoal_beam{i_goal}.opt_weight = optGoal_beam{i_goal}.opt_weight * numel(optGoal_beam{i_goal}.ROI_idx_old)/numel(optGoal_beam{i_goal}.ROI_idx);
        
        h_w1 = waitbar(0, 'merging superboxels');
        for i_supVox = 1:numel(superVoxList)
            supVox_idx = superVoxList(i_supVox);
            
            waitbar(i_supVox/numel(superVoxList), h_w1)
            
            idxList = find(superMask == supVox_idx);
            
            optGoal{i_goal}.beamlets_pruned(i_supVox,:) = sparse(mean(beamlets(idxList, :),1));
            optGoal_beam{i_goal}.beamlets_pruned(i_supVox,:) = sparse(mean(beamlets_joined(idxList, :),1));
            % -- make new indeces
            optGoal{i_goal}.ROI_idx(i_supVox) = idxList(1);
            optGoal_beam{i_goal}.ROI_idx(i_supVox) = idxList(1);
        end
        close(h_w1)
    end
    
    
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);

% -- 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, numBeam);
ub = [];
nonlcon = [];

% define opt limits, and make it fmincon progress
options = optimoptions('fmincon');
options.MaxFunctionEvaluations = N_fcallback1;
options.Display = 'iter';
options.PlotFcn = 'optimplotfval';
% options.UseParallel = true;
options.UseParallel = false;
% options.OptimalityTolerance = 1e-9;

%% -- INITIALIZE BEAMLET WEIGHTS --
switch pre_beamWeights
    case 'y'
        % should have been assigned previously.
        disp('Provided beamlet weights used for initial comparison')
    case 'n'
        % if initial beamlet weights are not provided, get quick estimate
        fprintf('\n running initial optimizer:')
        % initialize beamlet weights, OR
        w0_beams = ones(numBeam,1);
        w0_beams = mean(optGoal_beam{1}.D_final(optGoal{1}.ROI_idx) ./ (optGoal_beam{1}.beamlets_pruned*w0_beams+0.1)) * w0_beams;
        w0_beams = double(w0_beams);

        % -- GET BEAM WEIGHTS --
        tic
        w_beam = fmincon(fun1,w0_beams,A,b,Aeq,beq,lb_beam,ub,nonlcon,options);
        fprintf('  done!:')
        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
end

%% FULL OPTIMIZATION

% -- GET FULL BEAMLET WEIGHTS --
options.MaxFunctionEvaluations = N_fcallback2;
tic
fprintf('\n running full optimizer:')
w_fin = fmincon(fun2,w_beamlets,A,b,Aeq,beq,lb,ub,nonlcon,options);
fprintf('  done!:')
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([Pat_path, '\matlab_files\NLP_result.mat'], 'NLP_result');


plot_DVH(Geometry, D_full)
colorwash(Geometry.data, D_full, [-500, 500], [0, 35]);


end

%% 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 = sss
            for rgs_i = 1:optGoal{1}.NbrRangeScenarios % range scenario = rs
                fobj(sc_i)=eval_f(x, optGoal, nrs_i, sss_i, rgs_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, rgs_i)
    penalty = 0;
    % for each condition
    for goal_i = 1:numel(optGoal)
        switch optGoal{goal_i}.function
            % min, max, min_sq, max_sq, LeastSquare, min_perc_Volume, max_perc_Volume
            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}.rgs{rgs_i}.target) -...
                    (optGoal{goal_i}.nrs{nrs_i}.sss{sss_i}.rgs{rgs_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}.rgs{rgs_i}.beamlets_pruned * x)-...
                    (optGoal{goal_i}.nrs{nrs_i}.sss{sss_i}.rgs{rgs_i}.target)));
            case 'min_sq'
                % penalize if achieved dose is lower than target dose
                temp1=min(0, (optGoal{goal_i}.nrs{nrs_i}.sss{sss_i}.rgs{rgs_i}.beamlets_pruned * x)-...
                    (optGoal{goal_i}.nrs{nrs_i}.sss{sss_i}.rgs{rgs_i}.target));
                d_penalty = 1.0e0 * sum(temp1.*temp1);
            case 'max_sq'
                % penalize if achieved dose is higher than target dose
                temp1=max(0, (optGoal{goal_i}.nrs{nrs_i}.sss{sss_i}.rgs{rgs_i}.beamlets_pruned * x)-...
                    (optGoal{goal_i}.nrs{nrs_i}.sss{sss_i}.rgs{rgs_i}.target));
                d_penalty = 1.0e0 * sum(temp1.*temp1);
            case 'LeastSquare'
                % penalize with sum of squares any deviation from target
                % dose
                temp1 = (optGoal{goal_i}.nrs{nrs_i}.sss{sss_i}.rgs{rgs_i}.beamlets_pruned * x) - ...
                    optGoal{goal_i}.nrs{nrs_i}.sss{sss_i}.rgs{rgs_i}.target;
                d_penalty = 1.0e0* sum(temp1.^2);
            case 'min_perc_Volume'
                % penalize by amount of volume under threshold
                perc_vox = numel(find((optGoal{goal_i}.nrs{nrs_i}.sss{sss_i}.rgs{rgs_i}.target) -...
                    (optGoal{goal_i}.nrs{nrs_i}.sss{sss_i}.rgs{rgs_i}.beamlets_pruned * x) > 0)) ...
                    / numel(optGoal{goal_i}.nrs{nrs_i}.sss{sss_i}.rgs{rgs_i}.target);
                d_penalty = 3.0e4 * min(perc_vox-0.05, 0)
                
            case 'max_perc_Volume'
                % penalize by amount of volume under threshold
                perc_vox = numel(find((optGoal{goal_i}.nrs{nrs_i}.sss{sss_i}.rgs{rgs_i}.target) -...
                    (optGoal{goal_i}.nrs{nrs_i}.sss{sss_i}.rgs{rgs_i}.beamlets_pruned * x) < 0)) ...
                    / numel(optGoal{goal_i}.nrs{nrs_i}.sss{sss_i}.rgs{rgs_i}.target);
                d_penalty = 3.0e4 * min(perc_vox-0.05, 0)
                    
        end
        penalty = penalty + d_penalty * optGoal{goal_i}.opt_weight;
    end
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
    % rgs - random range scenarios
    
    % X - X>0 moves image right
    % Y - Y>0 moves image down
    % Z - in/out.
    
    shift_mag = 1; % vox of shift
    nrs_scene_list={[0,0,0]};

    
    % ----====#### CHANGE ROBUSTNESS HERE ####====----
    
    sss_scene_list={[0,0,0]};
%     sss_scene_list={[0,0,0], [-shift_mag,0,0], [shift_mag,0,0], [0,-shift_mag,0], [0,shift_mag,0], [0,0,-1], [0,0,1]};
%     sss_scene_list={[0,0,0], [-shift_mag,0,0], [shift_mag,0,0], [0,-shift_mag,0], [0,shift_mag,0],...
%         [-shift_mag*2,0,0], [shift_mag*2,0,0], [0,-shift_mag*2,0], [0,shift_mag*2,0]};

    % ----====#### CHANGE ROBUSTNESS HERE ####====----

    rgs_scene_list={[0,0,0]};

%     [targetIn, meta] = nrrdread('C:\010-work\003_localGit\WiscPlan_v2\data\archive\CDP_data\CDP5_DP_target.nrrd');
%     [targetIn, meta] = nrrdread('C:\010-work\003_localGit\WiscPlan_v2\data\PD_HD_dicomPhantom\Tomo_DP_target.nrrd');
%     [targetIn, meta] = nrrdread('C:\010-work\003_localGit\WiscPlan_v2\data\archive\CDP_data\CDP5_DP_target.nrrd');
    
    for i = 1:numel(optGoal)
        optGoal{i}.NbrRandScenarios     =numel(nrs_scene_list);
        optGoal{i}.NbrSystSetUpScenarios=numel(sss_scene_list);
        optGoal{i}.NbrRangeScenarios    =numel(rgs_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 idxValid]=get_RO_sss(targetImg2, sss_scene_list{sss_i});
                % beamlets stay the same
                
                for rgs_i = 1:optGoal{goal_i}.NbrRangeScenarios   % range scenario = rgs
                    % modify target and beamlets
                    targetImg4=targetImg3;
                    % beamlets stay the same
                    
                    %% make new target and beamlets
                    ROI_idx=[];
                    ROI_idx=find(targetImg4>0);
                    
                    target = optGoal{goal_i}.D_final(idxValid);
                    
                    beamlets_pruned = beamlets(ROI_idx, :);
                    
                    % save to optGoal output
                    optGoal{goal_i}.nrs{nrs_i}.sss{sss_i}.rgs{rgs_i}.ROI_idx            = ROI_idx;
                    optGoal{goal_i}.nrs{nrs_i}.sss{sss_i}.rgs{rgs_i}.beamlets_pruned    = beamlets_pruned;
                    optGoal{goal_i}.nrs{nrs_i}.sss{sss_i}.rgs{rgs_i}.target             = target;
                end
            end
        end
    end
end
% ------ supp: RO case SSS ------
function [targetImg3 ia]=get_RO_sss(targetImg2, sss_scene_shift);
    % translate the target image 
    
    targetImg3 = imtranslate(targetImg2,sss_scene_shift);
    
    % now we need to figure out if any target voxels fell out during the
    % shift
    imgValid = imtranslate(targetImg3,-sss_scene_shift);
    imgInvalid = (targetImg2-imgValid);
    
    idx_1 = find(targetImg2);
    idx_2 = find(imgInvalid);
    
    [idxValid,ia] = setdiff(idx_1,idx_2);
    
    [C,ia, ib] = intersect(idx_1,idxValid);
    
end