% CGLS conjugate gradient tomo solver, solve tomography as least squares
% tasks
% Note: in contrast to SART, it does not accept additional constraints 
% 
% [rec] = CGLS(rec, sino, cfg, vectors, Niter, 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
%     **Niter       - number of iterations
%     **varargin    - see the code + parameters of tomo.Atx_sup_partial
% *returns*
%     ++rec         - tomography 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] = CGLS(rec, sino, cfg, vectors, Niter, varargin)


    par = inputParser;
    par.addOptional('split', 1)
    par.addParameter('valid_angles', [])
    par.addParameter('deformation_fields',  {} )  % cell 3x1 of deformation arrays 
    par.addOptional('GPU', [])   % list of GPUs to be used in reconstruction
    par.addOptional('split_sub', 1)   % splitting of the sub block on smaller tasks in the Atx_partial method , 1 == no splitting 
    par.addOptional('verbose', 1)   % verbose = 0 : quiet, verbose : standard info , verbose = 2: debug 
    par.parse(varargin{:})
    res = par.Results;
    
    if Niter == 0
        return
    end
    
    if res.verbose
        disp('====== CGLS ==========')
    end
    
    if ~isempty(res.valid_angles)
        sino = sino(:,:,res.valid_angles);
        vectors = vectors(res.valid_angles,:);
        try cfg.lamino_angle = cfg.lamino_angle(res.valid_angles); end 
    end

    [Nlayers,Nw,Nproj] = size(sino);
    cfg.iProjAngles = Nproj;
    assert(cfg.iProjU == Nw, 'Wrong sinogram width')
    assert(cfg.iProjV == Nlayers, 'Wrong sinogram height')

    
    import tomo.*

    varargin = {'deformation_fields',res.deformation_fields,'GPU',res.GPU, 'split_sub', res.split_sub,'verbose', res.verbose};

    % r = sino - A*x
    r = sino - Ax_sup_partial(rec, cfg, vectors, res.split, varargin{:});
    % p = A'*r         
    p =  Atx_sup_partial(r, cfg, vectors,res.split, varargin{:});

    norm_sino = sqrt(mean(sino(:).^2));
    
    gamma_0 = sum(p(:).^2);
    t0 = tic;
    for i = 1:Niter 
%         progressbar(i, Niter)
        fprintf('CGLS Iter %i/%i\n', i, Niter)
        q = Ax_sup_partial(p, cfg, vectors,  res.split, varargin{:});
        alpha = gamma_0 / sum(q(:).^2);
        rec = rec + alpha * p;
        r = r - alpha * q;
        err(i) = sqrt(mean(r(:).^2));
        s = Atx_sup_partial(r, cfg, vectors, res.split, varargin{:});
        gamma_1 = sum(s(:).^2);
        beta = gamma_1 / gamma_0;
        gamma_0 = gamma_1;
        p = s + beta * p;
        
        if toc(t0) >10 && res.verbose % plot every 5s 
            figure(239821)
            subplot(1,2,1)
            plotting.imagesc3D(-rec, 'init_frame', size(rec,3)/2 )
            axis off image ; 
            colormap bone 
            title('CLGS reconstruction preview')
            subplot(1,2,2)
            loglog(err/norm_sino)
            title('Relative data error')
            axis tight 
            drawnow
            t0 = tic;
        end
        if i > 1 && err(i) > err(i-1)
            disp('Error increased, finishing')
            break
        end
            
    end


end