classdef Gt6DVolumeToBlobProjector < Gt6DVolumeProjector
    properties
        offsets = {};
        ss_offsets = {};
        blobs_depths = [];

        % Weights
        forwardWeights = {};
        backwardWeights = {};

        blob_sinogram_c_functions = true;
    end

    methods (Access = public)
        function self = Gt6DVolumeToBlobProjector(initialVolumes, detector_size, blobs_depths, varargin)
            self = self@Gt6DVolumeProjector(initialVolumes, detector_size, varargin{:});

            self.blobs_depths = blobs_depths;

            self.statistics.add_task('fproj_create_sum', 'FP Create Blob time');
            self.statistics.add_task('bproj_create_sum', 'BP Create Blob time');

            self.statistics.add_task('fproj_permute_sum', 'FP Permute Blob time');
            self.statistics.add_task('bproj_permute_sum', 'BP Permute Blob time');

            self.statistics.add_task('fproj_permute_ss_single', 'SS FP Permute Blob time (single)');
            self.statistics.add_task('bproj_permute_ss_single', 'SS BP Permute Blob time (single)');

            self.statistics.add_task('fproj_pre_weight_sum', 'FP Pre-weight time');
            self.statistics.add_task('bproj_pre_weight_sum', 'BP Pre-weight time');

            self.statistics.add_task('fproj_post_weight_sum', 'FP Post-weight time');
            self.statistics.add_task('bproj_post_weight_sum', 'BP Post-weight time');
        end

        function backwardWeights = getColumnsSum(self)
            numBlobs = numel(self.blobs_depths);
            blobs = cell(1, numBlobs);
            for m = 1:numBlobs
                blobs{m} = ones(self.proj_size(1), ...
                    self.blobs_depths(m), self.proj_size(2));
            end

            backwardWeights = self.backprojectVolume(blobs);
        end
    end

    methods (Access = protected)
        function blobs = projectVolume(self, volumes)
        % We get volumes in input and we want the blobs as output, but the
        % output from the base function is stacks of sinograms relative to
        % the geometry of the volumes, so we need to actually decompose the
        % output and recompose the slices into blobs

            sinograms = self.projectVolume@Gt6DVolumeProjector(volumes);

            numBlobs = numel(self.blobs_depths);
            self.statistics.tic('fproj_create_sum');
            if (self.blob_sinogram_c_functions)
                blobs_sizes = cell(1, numBlobs);
                for m = 1:numBlobs
                    blobs_sizes{m} = [self.proj_size(1), ...
                        self.blobs_depths(m), self.proj_size(2)];
                end
            else
                blobs = cell(1, numBlobs);
                for m = 1:numBlobs
                    blobs{m} = zeros(self.proj_size(1), ...
                        self.blobs_depths(m), self.proj_size(2), 'single');
                end
            end
            self.statistics.toc('fproj_create_sum');

            self.statistics.tic('fproj_permute_sum');
            if (self.blob_sinogram_c_functions)
                blobs = internal_cell_sinograms_to_blobs(blobs_sizes, sinograms, self.offsets);
            else
                try
                    numSinograms = numel(sinograms);
                    for n = 1:numSinograms
                        sino_offsets = self.offsets{n}.sino_offsets(:);
                        blob_offsets = self.offsets{n}.blob_offsets(:);
                        proj_offsets = self.offsets{n}.proj_offsets(:);
                        proj_coeffs = self.offsets{n}.proj_coeffs(:);
                        for m = 1:numel(proj_offsets)
                            p = proj_offsets(m);
                            b = blob_offsets(m);
                            c = proj_coeffs(m);
                            s = sino_offsets(m);
                            blobs{b}(:, p, :) = ...
                                blobs{b}(:, p, :) + c * sinograms{n}(:, s, :);
                        end
                    end
                catch mexc
                    fprintf('\n\nError!!\n')
                    fprintf('  indices: b %d, p %d, n %d, s %d, m %d\n', ...
                        b, p, n, s, m)
                    fprintf('  num blobs %d, num sinos %d\n', ...
                        numel(blobs), numSinograms)
                    fprintf('  size blob (%d %d %d), size sino (%d %d %d)\n\n', ...
                        size(blobs{b}), size(sinograms{n}))
                    fprintf('  Limits over p: [%d, %d]\n', ...
                        min(proj_offsets), max(proj_offsets))
                    fprintf('  Limits over s: [%d, %d]\n\n', ...
                        min(sino_offsets), max(sino_offsets))
                    rethrow(mexc)
                end
            end
            self.statistics.toc('fproj_permute_sum');

            if (~isempty(self.ss_offsets))
                blobs = self.projectVolumeSS(blobs, volumes);
            end
        end

        function blobs = projectVolumeSS(self, blobs, volumes)
        % We get volumes in input and we want the blobs as output, but the
        % output from the base function is stacks of sinograms relative to
        % the geometry of the volumes, so we need to actually decompose the
        % output and recompose the slices into blobs

            num_ss_geoms = numel(self.ss_geometries);
            for n = 1:num_ss_geoms
                ss_sino = self.projectSingleVolumeSS(volumes, n);

                self.statistics.tic('fproj_permute_ss_single');
                if (self.blob_sinogram_c_functions)
                    coeffs = self.ss_offsets{n};
                    blobs = internal_cell_single_sinogram_to_blobs(blobs, ss_sino, coeffs);
                else
                    try
                        sino_offsets = self.ss_offsets{n}.sino_offsets(:);
                        blob_offsets = self.ss_offsets{n}.blob_offsets(:);
                        proj_offsets = self.ss_offsets{n}.proj_offsets(:);
                        proj_coeffs = self.ss_offsets{n}.proj_coeffs(:);
                        for m = 1:numel(proj_offsets)
                            p = proj_offsets(m);
                            b = blob_offsets(m);
                            c = proj_coeffs(m);
                            s = sino_offsets(m);
                            blobs{b}(:, p, :) = ...
                                blobs{b}(:, p, :) + c * ss_sino(:, s, :);
                        end
                    catch mexc
                        fprintf('\n\nError!!\n')
                        fprintf('  indices: b %d, p %d, n %d, s %d, m %d\n', ...
                            b, p, n, s, m)
                        fprintf('  num blobs %d, num sinos %d\n', ...
                            numel(blobs), numel(self.ss_offsets))
                        fprintf('  size blob (%d %d %d), size sino (%d %d %d)\n\n', ...
                            size(blobs{b}), size(ss_sino))
                        fprintf('  Limits over p: [%d, %d]\n', ...
                            min(proj_offsets), max(proj_offsets))
                        fprintf('  Limits over s: [%d, %d]\n\n', ...
                            min(sino_offsets), max(sino_offsets))
                        rethrow(mexc)
                    end
                end
                self.statistics.toc('fproj_permute_ss_single');
            end

            num_geoms = numel(self.geometries);
            renorm_factor = num_geoms / (num_geoms + num_ss_geoms);

            blobs = internal_cell_prod_scalar_assign(blobs, renorm_factor);
        end

        function volumes = backprojectVolume(self, blobs)
        % Similar problem to the projection function, with the only
        % difference that in this case the input is blobs and we want to
        % transform them into stack of sinograms

            numSinograms = numel(self.astra_geometries);
            self.statistics.tic('bproj_create_sum');
            if (self.blob_sinogram_c_functions)
                sinograms_sizes = cell(1, numSinograms);
                for n = 1:numSinograms
                    sinograms_sizes{n} = [self.proj_size(1), ...
                        size(self.geometries{n}, 1), self.proj_size(2)];
                end
            else
                sinograms = cell(1, numSinograms);
                for n = 1:numSinograms
                    sinograms{n} = zeros(self.proj_size(1), ...
                        size(self.geometries{n}, 1), self.proj_size(2), 'single');
                end
            end
            self.statistics.toc('bproj_create_sum');

            self.statistics.tic('bproj_permute_sum');
            if (self.blob_sinogram_c_functions)
                sinograms = internal_cell_blobs_to_sinograms(sinograms_sizes, blobs, self.offsets);
            else
                for n = 1:numSinograms
                    sino_offsets = self.offsets{n}.sino_offsets(:);
                    blob_offsets = self.offsets{n}.blob_offsets(:);
                    proj_offsets = self.offsets{n}.proj_offsets(:);
                    proj_coeffs = self.offsets{n}.proj_coeffs(:);
                    for m = 1:numel(proj_offsets)
                        p = proj_offsets(m);
                        b = blob_offsets(m);
                        c = proj_coeffs(m);
                        s = sino_offsets(m);
                        sinograms{n}(:, s, :) = ...
                            sinograms{n}(:, s, :) + c * blobs{b}(:, p, :);
                    end
                end
            end
            self.statistics.toc('bproj_permute_sum');

            volumes = self.backprojectVolume@Gt6DVolumeProjector(sinograms);

            if (~isempty(self.ss_offsets))
                volumes = self.backprojectVolumeSS(volumes, blobs);
            end
        end

        function volumes = backprojectVolumeSS(self, volumes, blobs)
        % Similar problem to the projection function, with the only
        % difference that in this case the input is blobs and we want to
        % transform them into stack of sinograms

            numSinograms = numel(self.ss_geometries);
            if (self.blob_sinogram_c_functions)
                sinograms_sizes = cell(1, numSinograms);
                for n = 1:numSinograms
                    sinograms_sizes{n} = [self.proj_size(1), ...
                        size(self.ss_geometries{n}, 1), self.proj_size(2)];
                end
            end

            num_ss_geoms = numel(self.ss_geometries);
            for n = 1:num_ss_geoms
                self.statistics.tic('bproj_permute_ss_single');
                if (self.blob_sinogram_c_functions)
                    ss_sino = internal_cell_blobs_to_single_sinogram( ...
                        sinograms_sizes{n}, blobs, self.ss_offsets{n});
                else
                    ss_sino = zeros(self.proj_size(1), ...
                        size(self.ss_geometries{n}, 1), self.proj_size(2), 'single');

                    sino_offsets = self.ss_offsets{n}.sino_offsets(:);
                    blob_offsets = self.ss_offsets{n}.blob_offsets(:);
                    proj_offsets = self.ss_offsets{n}.proj_offsets(:);
                    proj_coeffs = self.ss_offsets{n}.proj_coeffs(:);
                    for m = 1:numel(proj_offsets)
                        p = proj_offsets(m);
                        b = blob_offsets(m);
                        c = proj_coeffs(m);
                        s = sino_offsets(m);
                        ss_sino(:, s, :) = ...
                            ss_sino(:, s, :) + c * blobs{b}(:, p, :);
                    end
                end
                self.statistics.toc('bproj_permute_ss_single');

                volume = self.backprojectSingleVolumeSS(ss_sino, n);

                coeffs = self.ss_coeffs(n);
                num_inerp = numel(coeffs.indx);
                for c = 1:num_inerp
                    b = coeffs.indx(c);
                    volumes{b} = volumes{b} + coeffs.coeff(c) * volume;
                end
            end

            num_geoms = numel(self.geometries);
            renorm_factor = num_geoms / (num_geoms + num_ss_geoms);

            volumes = internal_cell_prod_scalar_assign(volumes, renorm_factor);
        end

        function y = project(self, x)
            self.statistics.tic('fproj_pre_weight_sum');
            y = internal_cell_prod_copy(x, self.backwardWeights);
            self.statistics.toc('fproj_pre_weight_sum');

            y = self.projectVolume(y);

            self.statistics.tic('fproj_post_weight_sum');
            y = internal_cell_prod_assign(y, self.forwardWeights);
            self.statistics.toc('fproj_post_weight_sum');
        end

        function y = backproject(self, x)
            self.statistics.tic('bproj_pre_weight_sum');
            y = internal_cell_prod_copy(x, self.forwardWeights);
            self.statistics.toc('bproj_pre_weight_sum');

            y = self.backprojectVolume(y);

            self.statistics.tic('bproj_post_weight_sum');
            y = internal_cell_prod_assign(y, self.backwardWeights);
            self.statistics.toc('bproj_post_weight_sum');
        end
    end
end