% SART tomography reconstruction code
%
% [rec] = SART(rec, sino, cfg, vectors, cache, varargin)
% 
% Inputs:
%     **rec         - initial guess of the reconstruction 
%     **sino        - sinogram (Nlayers x width x Nangles)
%     **cfg         - config struct from ASTRA_initialize
%     **vectors     - vectors of projection rotation generated by ASTRA_initialize
%     **cache       - precalcualated cache by SART_prepare
%  *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 
%     ** valid_angles =  []       - list of valid angles, []==all are valid
%     ** filter =   'ram-lak'     - name of the FBP filter 
%     ** filter_value =   1       - fitlering value for the FBP filter 
%     ** deformation_fields = {}  - cell 3x1 of deformation arrays 
%     ** inv_deformation_fields={}-cell 3x1 of inverse 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 
%     ** relax = 0                - relaxation constant, shorter steps , 1 == no relaxation 
%     ** constraint= []           - constraint function that will be applied after every step 
%
%  *returns*
%     ++rec   reconstruction 

%*-----------------------------------------------------------------------*
%|                                                                       |
%|  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 [rec, err] = SART(rec, sinogram, cfg,vectors,cache,varargin)
   
    import tomo.*
    import utils.*
    import astra.*
    
    G  = cfg.Grouping; 
    cfg.iProjAngles = size(vectors,1);
    Nsets = ceil(cfg.iProjAngles/G);
    
    % try to optimize the group sizes 
    G = ceil(cfg.iProjAngles / Nsets); 
    Nsets = ceil(cfg.iProjAngles/G);
    
    par = inputParser;
    par.addOptional('split', 1)
    par.addParameter('valid_angles', [])
    par.addParameter('deformation_fields', cell(Nsets,1) )  % cell 3x1 of deformation arrays 
    par.addParameter('inv_deformation_fields', cell(Nsets,1) )  % 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', 1)   % verbose = 0 : quiet, verbose : standard info , verbose = 2: debug 
    par.addParameter('relax', 0, @(x)(x < 1 && x >= -1))  % relaxation constant, shorter steps , 1 == no relaxation 
    par.addParameter('constraint', [])   % constraint function that will be applied after every step 
    par.addParameter('keep_on_GPU', true)   % constraint function that will be applied after every step 

    
    par.parse(varargin{:})
    res = par.Results;
    
    
    if ~isempty(res.valid_angles)
        sinogram = sinogram(:,:,res.valid_angles);
        vectors = vectors(res.valid_angles,:);
        cache.R = cache.R(:,:,res.valid_angles);
        try cfg.lamino_angle = cfg.lamino_angle(res.valid_angles); end 
    end
    if islogical(res.valid_angles); res.valid_angles = find(res.valid_angles); end

   
    err = nan(Nsets,1); 
    rng('default')

    if isempty(res.deformation_fields) && isempty(res.inv_deformation_fields)
        indices = randperm(cfg.iProjAngles);
        for i=randperm(Nsets)
            ind{i} = indices((1+(i-1)*G):min(cfg.iProjAngles,i*G));
        end
    else
        % for deformation tomography keep the blocks corresponding to the
        % subtomograms 
        for i=1:Nsets
            ind{i} = (1+(i-1)*G):min(cfg.iProjAngles,i*G);
        end
    end
    if ~isempty(res.valid_angles)
        for i=1:Nsets
            ind{i} = intersect(ind{i}, res.valid_angles); 
        end
    end
    
    for i = 1:Nsets
        if res.verbose > 0
            utils.progressbar(i, Nsets); 
        end
    
        
        cfg.iProjAngles = length(ind{i});
        data = sinogram(:,:,ind{i}); 
        R = cache.R(:,:,ind{i}); 
        if ~isempty(res.deformation_fields)
            R = R .* tukeywin(cfg.iProjV, 0.2); % avoid amplification of edge regions where deformation may be wrong 
        end  
        if res.keep_on_GPU
            data = gpuArray(data); 
        end
        
        if ~isa(data, 'gpuArray') && (isempty(res.deformation_fields) || isempty(res.deformation_fields{i})) 
            proj = Ax_sup_partial(rec, cfg, vectors(ind{i},:), res.split,  'GPU', res.GPU, 'split_sub', res.split_sub, 'verbose',0, 'deformation_fields',res.deformation_fields{i});
        else
            proj = Ax_partial(rec, cfg, vectors(ind{i},:), res.split_sub, 'verbose',0, 'deformation_fields', res.deformation_fields{i} );
        end

        
        D = data -  proj;
        D = D .* R; 
        
        if nargout > 1
            %% get error before update  
            err(ind{i}) = gather(sqrt(mean(mean(D.^2))));
        end


        if ~isa(D, 'gpuArray') && (isempty(res.deformation_fields) || isempty(res.deformation_fields{i})) 
            rec_upd = Atx_sup_partial(D, cfg, vectors(ind{i},:), res.split,  'GPU', res.GPU, 'split_sub', res.split_sub, 'verbose', res.verbose, 'deformation_fields', res.inv_deformation_fields{i} );
        else
            rec_upd = Atx_partial(D, cfg, vectors(ind{i},:), 1, 'verbose', res.verbose, 'deformation_fields', res.inv_deformation_fields{i} );
        end
        rec_upd = ((1-res.relax)/cfg.iProjAngles)*rec_upd;
        if ~isa(rec,'gpuArray')
            rec_upd = gather(rec_upd);
        end
        rec = rec + rec_upd;
        % apply constraints 
        if ~isempty(res.constraint)
            rec = res.constraint(rec);
        end
        
    end
    if ~isempty(res.constraint)
        rec = res.constraint(rec);
    end
end