% [CloseInd,Neighbours]  = nonlocalTV_weight_GPU(Im0, Rpatch, Rwin, Nclose, sigma)
% FUNCTION 
%   find close indices for the nonlocalTV_GPU method , similar to the nonlocal means method
% Inputs:
%     Im0 - volume to be regularized
%     Rpatch - radius of the patch used to ocompare similarities 
%     Rwin - radius of the window used to find similar regions 
%     Nclose - number of most similar regions searched 
%     sigma - similarity threshold, (tunning constant, value < 1)
% Outputs: 
%     CloseInd - close indices generated by nonlocalTV_weight code  
%     Neighbours - list of neighboring pixels used for calculation and excluded from CloseInd


% RECOMPILE COMMAND
%  mexcuda -output +regularization/private/nonlocalTV_weight_tex +regularization/private/TV_cuda_texture.cu +regularization/private/nonlocalTV_mex_weight.cpp


%*-----------------------------------------------------------------------*
%|                                                                       |
%|  Except where otherwise noted, this work is licensed under a          |
%|  Creative Commons Attribution-NonCommercial-ShareAlike 4.0            |
%|  International (CC BY-NC-SA 4.0) license.                             |
%|                                                                       |
%|  Copyright (c) 2017 by Paul Scherrer Institute (http://www.psi.ch)    |
%|                                                                       |
%|       Author: CXS group, PSI                                          |
%*-----------------------------------------------------------------------*
% You may use this code with the following provisions:
%
% If the code is fully or partially redistributed, or rewritten in another
%   computing language this notice should be included in the redistribution.
%
% If this code, or subfunctions or parts of it, is used for research in a 
%   publication or if it is fully or partially rewritten for another 
%   computing language the authors and institution should be acknowledged 
%   in written form in the publication: “Data processing was carried out 
%   using the “cSAXS matlab package” developed by the CXS group,
%   Paul Scherrer Institut, Switzerland.” 
%   Variations on the latter text can be incorporated upon discussion with 
%   the CXS group if needed to more specifically reflect the use of the package 
%   for the published work.
%
% A publication that focuses on describing features, or parameters, that
%    are already existing in the code should be first discussed with the
%    authors.
%   
% This code and subroutines are part of a continuous development, they 
%    are provided “as they are” without guarantees or liability on part
%    of PSI or the authors. It is the user responsibility to ensure its 
%    proper use and the correctness of the results.


function [CloseInd,Neighbours]  = nonlocalTV_weight_GPU(Im0, Rpatch, Rwin, Nclose, sigma)
%% nonlocal TV 

%  mexcuda -output nonlocalTV_weight_tex TV_cuda_texture.cu nonlocalTV_mex_weight.cpp


Nwin= 2*Rwin+1;
Nclose_min=2;
Nclose_max = 256; 

if Nclose > Nclose_max
    error('Nclose is more than Nclose_max')
end
if ismatrix(Im0)
    assert(Nwin^2 < Nclose_max, 'Too large Nwin')
%     Neighbours=uint8(sub2ind([Nwin,Nwin],...
%     Rwin+1+[0,-1,1,0,0],Rwin+1+[0,0,0,-1,1]));
    Neighbours=uint8(sub2ind([Nwin,Nwin],...
    Rwin+1,Rwin+1));
else
    assert(Nwin^3 < Nclose_max, 'Too large Nwin')
     % include all neighbors 
%     Neighbours=uint8(sub2ind([Nwin,Nwin,Nwin],...
%     Rwin+1+[0,-1,1,0,0,0,0], ...
%     Rwin+1+[0,0,0,-1,1,0,0], ...
%     Rwin+1+[0,0,0,0,0,-1,1]));
    % only the central point !! 
    Neighbours=uint8(sub2ind([Nwin,Nwin,Nwin],...
    Rwin+1, Rwin+1,Rwin+1)); 
end

% keyboard

[Nx, Ny, Nz] = size(Im0);
MAX_SIZE = 512;   %% TESTED FOR TITAN X 12GB !!! 
% MAX_SIZE = 320;   %% roughly for Quadro K4200
split = ceil([Nx, Ny, Nz]/MAX_SIZE);




if prod(split) == 1
    CloseInd = nonlocalTV_weight_GPU_partial(Im0, Rpatch, Rwin, Nclose,Nclose_min,sigma, Neighbours);
else  %% in case of too large datasets 
    warning('Volume is too large, it will be sliced')
    CloseInd = zeros(Nx, Ny,Nz,Nclose, 'uint8');
    ind = {':',':',':'};
    
    for i = 1:split(1)
        ind{1} = (1+(i-1)*ceil(Nx/split(1))):min(Nx,i*ceil(Nx/split(1)));
        for j = 1:split(2)
            ind{2} = (1+(j-1)*ceil(Ny/split(2))):min(Ny,j*ceil(Ny/split(2)));
            for k = 1:split(3)
                 ind{3} = (1+(k-1)*ceil(Nz/split(3))):min(Nz,j*ceil(Nz/split(3)));
                 CloseInd(ind{:},:) = gather(nonlocalTV_weight_GPU_partial(Im0(ind{:}), Rpatch, Rwin, Nclose,Nclose_min, sigma, Neighbours));
                 utils.progressbar(k + (j-1)*split(2)+(i-1)*split(1)*split(2), prod(split))
            end
        end
    end
end



end

function CloseInd = nonlocalTV_weight_GPU_partial(Im0, Rpatch, Rwin, Nclose, Nclose_min, sigma, Neighbours)


    Im0_gpu = gpuArray(single(Im0));



    Npatch= 2*Rpatch+1;
    %% get simple estimate of the average gradients 

    if ismatrix(Im0_gpu)
        dx = diff(Im0_gpu,1,1);
        dx = dx(:,2:end).^2;
        dy = diff(Im0_gpu,1,2);
        dy = dy(2:end,:).^2; 
        ker = ones(Npatch, 'single');
        threshold = mean2(sqrt(conv2( (dx + dy)/2,ker, 'same')));
    elseif ndims(Im0_gpu) == 3
        dx = diff(Im0_gpu,1,1);
        dx = dx(:,2:end,2:end).^2;
        dy = diff(Im0_gpu,1,2);
        dy = dy(2:end,:,2:end).^2; 
        dz = diff(Im0_gpu,1,3);
        dz = dz(2:end,2:end,:).^2; 
        ker = ones([Npatch,Npatch,Npatch], 'single');
        threshold = mean2(sqrt(convn( (dx + dy + dz)/3,ker, 'same'))); 
    end


    % sigma is changing the threshold level 
    %% patched more different than threshold will be rejected !!! 
    threshold = gather(threshold ) * sigma;
    clear dx dy dz 

    % tic
    CloseInd = nonlocalTV_weight_tex(Im0_gpu,Neighbours,Nclose,Nclose_min,Rwin,Rpatch,threshold);
    % toc

end