function vol_prep
% This series of functions take in segmented PET/CT image paths and outputs
% Dose Painting (DP) maps, as well as some intermediate results. 
% this function calls three other functions:
% - vol_prep_bin
% - vol_prep_fuzzy
% - vol_prep_DPmap
% that are needed for expanding the initial segmentations done in the
% segmentation GUI

close all

%% ---=== INPUT PARAMS ===---

patname = 'FET_FGL005';
% patname = 'FET_FGL015';
% patname = 'FET_FGL022';
timepoint = 'B1';

paths.in = ['\\Mpufs5\data_wnx1\_Data\Glioma_aus\' patname '\' timepoint '\Processed'];
paths.CT_in = [patname '_' timepoint '_CT2FET'];
paths.target_bin_in  = [ patname '_' timepoint '_seg_thr2.0'];
paths.target_fzy_in  = [ patname '_' timepoint '_seg_combined'];
paths.body_bin_in    = [ patname '_' timepoint '_head'];

% paths.target_bin_in  = [patname '_' timepoint '_thr2'];
% paths.target_fzy_in  = [patname '_' timepoint '_thr2'];
% paths.body_bin_in    = [patname '_' timepoint '_head'];


margins.CTV = 1.0; % GTV->CTV margin, in cm
margins.PTV = 0.5; % CTV->PTV margin, in cm
margins.body = 7.0; % how far from tumor do we still include body, in cm

dose.min = 60;
dose.max = 63;


%% ---=== FUNCTION CALLS ===---
fprintf('Starting binary volume expansion...')
vol_prep_bin(paths, margins)
fprintf(' done!\n')

fprintf('Starting fuzzy volume expansion...')
vol_prep_fuzzy(paths, margins)
fprintf(' done!\n')

fprintf('Creating dose plans ...')
vol_prep_DPmap(paths, dose)
fprintf(' done!\n')
end


function vol_prep_bin(paths, margins)
% this function takes in the segmentations and creates PTV/GTV/CTV binary volumes,
% as well as appropriately cropped healthy tissue (body) segmentation.

[CT_in, CT_meta] = nrrdread([paths.in '\' paths.CT_in '.nrrd']);
[GTV_in, GTV_meta] = nrrdread([paths.in '\' paths.target_bin_in '.nrrd']);
[body_in, body_meta] = nrrdread([paths.in '\' paths.body_bin_in '.nrrd']);

colorwash(CT_in, GTV_in,  [-500, 500], [0, 1.2]);

%% identify target voxels - these are used when identifiying valid beamlets
% voxels that need to be considered (above the threshold)
p_threshold = 0;
GTV_binary = zeros(size(GTV_in));
GTV_binary(GTV_in>p_threshold) = 1;

% Account for infiltration (GTV -> CTV)
CTV = zeros(size(GTV_binary));
bwD = bwdistsc(GTV_binary, GTV_meta.spacedirections);
CTV(bwD<margins.CTV) = 1;

% expand to PTV (CTV -> PTV)
PTV = zeros(size(GTV_binary));
bwD = bwdistsc(GTV_binary, GTV_meta.spacedirections);
PTV(bwD<margins.PTV) = 1;

%% identify body volume of interest
body = body_in;
bwD = bwdistsc(GTV_binary, GTV_meta.spacedirections);
body(PTV>0) = 0;
body(bwD>margins.body) = 0; % how far from the tumor are we still interested in the body

%% save outputs
filename = [paths.in '\RODP_files\' paths.target_bin_in '_binCTV.nrrd'];
matrix = single(CTV);
pixelspacing = GTV_meta.spacedirections;
origin = GTV_meta.spaceorigin;
nrrdWriter(filename, matrix, pixelspacing, origin, 'raw');

filename = [paths.in '\RODP_files\' paths.target_bin_in '_binPTV.nrrd'];
matrix = single(PTV);
pixelspacing = GTV_meta.spacedirections;
origin = GTV_meta.spaceorigin;
nrrdWriter(filename, matrix, pixelspacing, origin, 'raw');

filename = [paths.in '\RODP_files\' paths.body_bin_in '_crop_bin.nrrd'];
matrix = single(body);
pixelspacing = body_meta.spacedirections;
origin = body_meta.spaceorigin;
nrrdWriter(filename, matrix, pixelspacing, origin, 'raw');

colorwash(CT_in, PTV,  [-500, 500], [0, 1.2]);
colorwash(CT_in, body, [-500, 500], [0, 1.2]);
end

function vol_prep_fuzzy(paths, margins)
% this function takes in the segmentations and creates fuzzy (probabilistic)
% PTV/GTV/CTV volumes

infilt_range = margins.CTV; % range of infiltration, in cm. Should be value of sigma (or 1/3 decay range). Assuming normal infil.

[CT_in, CT_meta] = nrrdread([paths.in '\' paths.CT_in '.nrrd']);
[GTV_in, GTV_meta] = nrrdread([paths.in '\' paths.target_fzy_in '.nrrd']);
% [body_in, body_meta] = nrrdread([paths.in '\' paths.body_bin_in '.nrrd']);
colorwash(CT_in, GTV_in,  [-500, 500], [0, 1.2]);

%% Account for infiltration (GTV -> CTV)

% -- get distance from central dot --
ir_a = 1.5; % infiltration range factor

CTV = zeros(size(GTV_in));
bwD = bwdistsc(GTV_in>0.05, GTV_meta.spacedirections);
CTV(bwD< (ir_a * infilt_range) ) = 1;

% for each voxel in the area of interest, figure out the highest
% probability based on range from tumor * p that tumor is there
GTV_idx_lst = find(GTV_in   >0.05);
CTV_idx_lst = find(CTV      >0.0);
prob_CTV = zeros(size(GTV_in));

f = waitbar(0,'calculating expansion');

% define area of interest
Ksize = 1+2*ceil((ir_a * infilt_range)/min(GTV_meta.spacedirections)); % get kernel for convolution.
kernel_base = zeros(Ksize,Ksize,Ksize); 
kernel_base(ceil(Ksize/2),ceil(Ksize/2),ceil(Ksize/2))=1;
bwD = bwdistsc(kernel_base, GTV_meta.spacedirections);
kernel = zeros(size(kernel_base));
kernel(bwD< (ir_a * infilt_range)) = 1;

CTV_idx_lst_num = numel(CTV_idx_lst);
prob_CTV_solutions = zeros(1,CTV_idx_lst_num);

tic 
for i = 1:CTV_idx_lst_num
    i_idx = CTV_idx_lst(i);
    [y1, x1, z1] = ind2sub(size(GTV_in), i_idx);
    
    waitbar(i/numel(CTV_idx_lst),f);
    
    % get the crop of nearby voxels in tumor
%     tabula = zeros(size(GTV_in));
%     tabula(i_idx) = 1;
%     tabula2 = convn(tabula, kernel, 'same');
%     idx_nearby = find(tabula2>0);
    
    idx_nearby = f_neighbors(i_idx, size(GTV_in), ir_a, infilt_range, GTV_meta.spacedirections);
    
    % overlap
    idx_GTV_nearby = intersect (idx_nearby, GTV_idx_lst);
    
    if numel(idx_GTV_nearby) <1
        prob_CTV_solutions(i) = 0;
        error('this is fucky. no voxels within range')
    end
    
    p_list=zeros(1, numel(idx_GTV_nearby));
    p_list_max = 0;
    
    for j = 1:numel(idx_GTV_nearby)
        j_idx = idx_GTV_nearby(j);
        if GTV_in(j_idx) < p_list_max
            continue
        end
        
        [y2, x2, z2] = ind2sub(size(GTV_in), j_idx);
        d = sqrt((((y2-y1)*GTV_meta.spacedirections(1))^2 + ...
            ((x2-x1)*GTV_meta.spacedirections(2))^2 + ((z2-z1)*GTV_meta.spacedirections(3))^2));
        % get probability of invisible infiltration at given voxel, based
        % on single voxel from tumor
        p_list(j) = GTV_in(j_idx) * exp(-(d*d)/(infilt_range*infilt_range));
        
        p_list_max = max(p_list_max, p_list(j));
        
        if p_list(j) == 1
            prob_CTV_solutions(i) = 1;
            break
        end
    end
    prob_CTV_solutions(i) = max(p_list);
end
toc
prob_CTV(CTV_idx_lst) = prob_CTV_solutions;
close(f)

colorwash(CT_in, prob_CTV,  [-500, 500], [0, 1.2]);


%% save outputs
filename = [paths.in '\RODP_files\' paths.target_bin_in '_fuzzyCTV.nrrd'];
matrix = single(prob_CTV);
pixelspacing = GTV_meta.spacedirections;
origin = GTV_meta.spaceorigin;
nrrdWriter(filename, matrix, pixelspacing, origin, 'raw');

end
function neigh_mat_idx = f_neighbors(idx, mat_size, ir_a, infilt_range, voxsize);
% this function returns indeces of the box nearby around 
    [y, x, z] = ind2sub(mat_size, idx);
    
    y1 = max(1, floor(y - ir_a*infilt_range/(voxsize(1))));
    y2 = min(mat_size(1), ceil(y + ir_a*infilt_range/(voxsize(1))));
    
    x1 = max(1, floor(x - ir_a*infilt_range/(voxsize(2))));
    x2 = min(mat_size(1), ceil(x + ir_a*infilt_range/(voxsize(2))));
    
    z1 = max(1, floor(z - ir_a*infilt_range/(voxsize(3))));
    z2 = min(mat_size(1), ceil(z + ir_a*infilt_range/(voxsize(3))));
    
    tabula = zeros(mat_size);
    tabula(y1:y2, x1:x2, z1:z2)=1;
    
    neigh_mat_idx = find( tabula>0);

end

function vol_prep_DPmap(paths, dose)
% this function takes in the segmentations and creates fuzzy (probabilistic)
% PTV/GTV/CTV volumes

[CT_in, CT_meta] = nrrdread([paths.in '\' paths.CT_in '.nrrd']);
[CTV_fuzzy, CTV_meta] = nrrdread([paths.in '\RODP_files\' paths.target_bin_in '_fuzzyCTV.nrrd']);

dose_min = f_dose_min(CTV_fuzzy, dose.min);
dose_max = f_dose_max(CTV_fuzzy, dose.max);

orthoslice(dose_min, [0, dose.max+10]) % plot min dose boundary
orthoslice(dose_max, [0, dose.max+10]) % plot max dose boundary

%% save outputs
filename = [paths.in '\RODP_files\' paths.target_bin_in '_DP_minDose.nrrd'];
% filename = [pathIn 'RODP_files\FET_FGL005_B1_DP_minDose.nrrd'];
matrix = single(dose_min);
pixelspacing = CTV_meta.spacedirections;
origin = CTV_meta.spaceorigin;
nrrdWriter(filename, matrix, pixelspacing, origin, 'raw');

filename = [paths.in '\RODP_files\' paths.target_bin_in '_DP_maxDose.nrrd'];
% filename = [pathIn 'RODP_files\FET_FGL005_B1_DP_maxDose.nrrd'];
matrix = single(dose_max);
pixelspacing = CTV_meta.spacedirections;
origin = CTV_meta.spaceorigin;
nrrdWriter(filename, matrix, pixelspacing, origin, 'raw');

end
function dose_min = f_dose_min(CTV_fuzzy, D_min)
% get low boundary for ideal dose plan
dose_min = D_min* min(1, 2* max(0, (CTV_fuzzy - 0.25)));
dose_min(CTV_fuzzy>0.98) = 1.2*D_min;
% orthoslice(dose_min)
end
function dose_max = f_dose_max(CTV_fuzzy, D_max)
% get hihg boundary for ideal dose plan
dose_max = D_max* min(1, 2* max(0, (CTV_fuzzy - 0.0)));
dose_max(CTV_fuzzy>0.98) = 1.2*D_max;
end