% SART_PREPARE prepare cache and config for SART tomography reconstruction code
%
% [cache,cfg] = SART_prepare(cfg, vectors, grouping, varargin)
% 
% Inputs:
%     **cfg         - config struct from ASTRA_initialize
%     **vectors     - vectors of projection rotation generated by ASTRA_initialize
%     **split       - split parameter for the ASTRA projector 
%  *optional*
%     ** split =[1,1,1]           - split the solved volume, split(3 is used to split in separated blocks, split(1:2) is used inside Atx_partial to du subplitting for ASTRA 
%     ** deformation_fields = {}  - cell 3x1 of deformation arrays 
%     ** GPU =  []                - list of GPUs to be used in reconstruction
%     ** split_sub =  [1,1,1]     - splitting of the sub block on smaller tasks in the Atx_partial method , 1 == no splitting 
%     ** verbose =  1             - verbose = 0 : quiet, verbose : standard info , verbose = 2: debug 
%     ** keep_on_GPU = false      - if true, do not gather arrays from GPU  -> faster 
%     ** weights = []             - custom weights (3D array) that denote reliable region in the projection
%
% *returns*
%     ++cache       - precalcualted values needed for SART
%     ++cfg         - modified ASTRA config structure needed for SART

%*-----------------------------------------------------------------------*
%|                                                                       |
%|  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 [cache,cfg] = SART_prepare(cfg, vectors, grouping, varargin)
    
    par = inputParser;
    par.addOptional('split', 1)
    par.addParameter('deformation_fields',  {} )  % cell 3x1 of deformation arrays 
    par.addParameter('GPU', [])   % list of GPUs to be used in reconstruction
    par.addParameter('split_sub', 1)   % splitting of the sub block on smaller tasks in the Atx_partial method , 1 == no splitting 
    par.addParameter('verbose', 0)   % verbose = 0 : quiet, verbose : standard info , verbose = 2: debug 
    par.addParameter('keep_on_gpu', false);  % keep resuls on GPU or move to RAM
    par.addParameter('weights', []);  % custom weights (3D array) that denote reliable region in the projection
    par.parse(varargin{:})
    res = par.Results;
    
    Nangles = length(vectors); 
    cfg.Grouping = min(Nangles, grouping); 
    cfg.iProjGroups = ceil(Nangles/cfg.Grouping);
    
    if ~res.keep_on_gpu && gpuDeviceCount
        empty = ones(cfg.iVolX, cfg.iVolY, cfg.iVolZ, 'single');
        R = tomo.Ax_sup_partial(empty, cfg, vectors, res.split,  'GPU', res.GPU, 'split_sub', res.split_sub, 'verbose', res.verbose, 'deformation_fields', res.deformation_fields );
    else
        empty = gpuArray.ones(cfg.iVolX, cfg.iVolY, cfg.iVolZ, 'single');
        R = astra.Ax_partial(empty, cfg, vectors, res.split_sub, 'verbose', res.verbose, 'deformation_fields', res.deformation_fields );
    end

    % some "representative" sample of R 
    R_val = R(ceil(end/2), ceil(end/2), ceil(end/2));
    R = (R>0)./sqrt(R.^2+(1e-2 * R_val)^2 ); %% !! avoid division by 0 
  
    if ~isempty(res.weights)
       R = R.*res.weights;   % apply custom weights  
    end
    
    cache.R = R; 
    cfg.iProjAngles = Nangles;
    cache.Nbunches = cfg.iProjGroups;

end