gtDefShapeFunctionsCreateNW.m

function sf = gtDefShapeFunctionsCreateNW(sampler, eta_resoltion, varargin)
% FUNCTION sf = gtDefShapeFunctionsCreateNW(sampler, eta_resoltion, varargin)
%

    conf = struct( ...
        'data_type', 'double', ...
        'fast_inverse', false, ...
        'fast_lines', false, ...
        'compute_w_sfs', false );
    conf = parse_pv_pairs(conf, varargin);

    p = sampler.parameters;

    if (~exist('eta_resoltion', 'var') || isempty(eta_resoltion))
        eta_resoltion = 1;
    end

    det_ind = sampler.detector_index;

    om_step = gtAcqGetOmegaStep(p, det_ind);
    om_step_rod = tand(om_step / 2);

    space_res = tand(sampler.estimate_maximum_resolution() ./ 2);
    eta_step = 2 * atand(space_res(1)) / eta_resoltion;

    oversampling = 11;
    oversampling_w = oversampling;
    step_w = 1 / oversampling_w;
    shift_w = 0.5 - step_w / 2;
    oversampling_n = (oversampling - 1) / eta_resoltion + 1;
    step_n = 1 / oversampling_n;
    shift_n = 0.5 - step_n / 2;

    sub_space_size = tand(sampler.sampling_res / 2);

    % Let's first compute the W extent
    half_rspace_sizes = sub_space_size / 2;

    ref_gr = sampler.get_reference_grain();

    ref_ond = ref_gr.proj(det_ind).ondet;
    ref_inc = ref_gr.proj(det_ind).included;
    ref_sel = ref_gr.proj(det_ind).selected;

    ref_ab_sel = ref_ond(ref_inc(ref_sel));
    ref_ab_inc = ref_ond(ref_inc);

    num_blobs = numel(find(ref_sel));

    pl_samd = ref_gr.allblobs(det_ind).plorig(ref_ab_sel, :);

    g = gtMathsRod2OriMat(ref_gr.R_vector');
    pl_cry = gtVectorLab2Cryst(pl_samd, g)';

%     pl_labd = gtGeoSam2Lab(pl_samd, eye(3), self.parameters.labgeo, self.parameters.samgeo, true)';

    fprintf('Computing bounds of NW shape functions: ')
    c = tic();
    tot_orientations = numel(sampler.lattice.gr);

    oris_bounds_size = [size(sampler.lattice.gr, 1), size(sampler.lattice.gr, 2), size(sampler.lattice.gr, 3)] + 1;
    oris_lims_min = sampler.lattice.gr{1, 1, 1}.R_vector - half_rspace_sizes;
    oris_lims_max = sampler.lattice.gr{end, end, end}.R_vector + half_rspace_sizes;

    oris_steps_x = linspace(oris_lims_min(1), oris_lims_max(1), oris_bounds_size(1));
    oris_steps_y = linspace(oris_lims_min(2), oris_lims_max(2), oris_bounds_size(2));
    oris_steps_z = linspace(oris_lims_min(3), oris_lims_max(3), oris_bounds_size(3));

    oris_bounds = cell(oris_bounds_size);
    for dz = 1:oris_bounds_size(3)
        for dy = 1:oris_bounds_size(2)
            for dx = 1:oris_bounds_size(1)
                r_vec = [oris_steps_x(dx), oris_steps_y(dy), oris_steps_z(dz)];
                oris_bounds{dx, dy, dz} = struct( ...
                    'phaseid', ref_gr.phaseid, ...
                    'center', ref_gr.center, 'R_vector', r_vec );
            end
        end
    end

    chunk_size = 250;
    for ii = 1:chunk_size:numel(oris_bounds)
        end_chunk = min(ii+chunk_size, numel(oris_bounds));
        num_chars_gr_ii = fprintf('%05d/%05d', ii, numel(oris_bounds));

        oris_bounds(ii:end_chunk) = gtCalculateGrain_p( ...
            oris_bounds(ii:end_chunk), p, ...
            'ref_omind', ref_gr.allblobs(det_ind).omind, ...
            'included', ref_ab_inc );

        fprintf(repmat('\b', [1 num_chars_gr_ii]));
    end
    fprintf('Done in %g seconds.\n', toc(c))

    fprintf('Precomputing shared values..')
    % Forming the grid in orientation space, on the plane that is
    % perpendicular to the axis of rotation (z-axis in the w case),
    % where to sample the sub-regions of the orientation-space voxels

    if (conf.compute_w_sfs)
        num_poins = max(ceil(sub_space_size ./ space_res ./ 2), 5);
        pos_x = linspace(-half_rspace_sizes(1), half_rspace_sizes(1), num_poins(1)+1);
        pos_y = linspace(-half_rspace_sizes(2), half_rspace_sizes(2), num_poins(2)+1);

        renorm_coeff_w = (pos_x(2) - pos_x(1)) * (pos_y(2) - pos_y(1)) / prod(sub_space_size);

        pos_x = (pos_x(1:end-1) + pos_x(2:end)) / 2;
        pos_y = (pos_y(1:end-1) + pos_y(2:end)) / 2;

        pos_x_grid_w = reshape(pos_x, 1, [], 1);
        pos_x_grid_w = pos_x_grid_w(1, :, ones(1, num_poins(2)));
        pos_y_grid_w = reshape(pos_y, 1, 1, []);
        pos_y_grid_w = pos_y_grid_w(1, ones(1, num_poins(1)), :);
        pos_z_grid_w = zeros(1, num_poins(1), num_poins(2));

        tot_sampling_points_w = num_poins(1) * num_poins(2);
        ones_tot_sp_w = ones(tot_sampling_points_w, 1);
    end

    ref_ws = ref_gr.allblobs(det_ind).omega(ref_ab_sel) ./ om_step;

    % There is no point in computing the same rotation matrices each
    % time, so we pre compute them here and store them for later
    c = tic();
    % selecting Ws to find the interval
    ori_bounds = [oris_bounds{:}];
    allblobs = [ori_bounds(:).allblobs];
    w_tabs_int = [allblobs(:).omega];
    w_tabs_int = w_tabs_int(ref_sel, :) ./ om_step;
    w_tabs_int = gtGrainAnglesTabularFix360deg(w_tabs_int, ref_ws, p);
    w_tabs_int = reshape(w_tabs_int, [], oris_bounds_size(1), oris_bounds_size(2), oris_bounds_size(3));

    ref_ns = ref_gr.allblobs(det_ind).eta(ref_ab_sel);
    n_tabs_int = [allblobs(:).eta];
    n_tabs_int = n_tabs_int(ref_sel, :);
    n_tabs_int = gtGrainAnglesTabularFix360deg(n_tabs_int, ref_ns);
    n_tabs_int = n_tabs_int ./ eta_step;
    n_tabs_int = reshape(n_tabs_int, [], oris_bounds_size(1), oris_bounds_size(2), oris_bounds_size(3));

    rotcomp_w = gtMathsRotationMatrixComp(p.labgeo.rotdir', 'col');

    if (conf.compute_w_sfs)
        rotcomp_z = gtMathsRotationMatrixComp([0, 0, 1]', 'col');

        % Precomputing A matrix components for all the used omegas
        ref_round_int_ws = round(ref_ws + 0.5) - 0.5;
        ref_round_ws = ref_round_int_ws .* om_step;
        rot_samp_w = gtMathsRotationTensor(ref_round_ws, rotcomp_w);

        beamdirs = p.labgeo.beamdir * reshape(rot_samp_w, 3, []);
        beamdirs = reshape(beamdirs, 3, [])';

        Ac_part_w = beamdirs * rotcomp_z.cos;
        As_part_w = beamdirs * rotcomp_z.sin;
        Acc_part_w = beamdirs * rotcomp_z.const;

        Ac_part_w = Ac_part_w(:, :, ones_tot_sp_w);
        As_part_w = As_part_w(:, :, ones_tot_sp_w);
        Acc_part_w = Acc_part_w(:, :, ones_tot_sp_w);
    end

    fprintf('\b\b: Done in %g seconds.\n', toc(c))

    sf = cell(tot_orientations, 1);

    [o_x, o_y, o_z] = ind2sub(oris_bounds_size-1, 1:tot_orientations);

    fprintf('Computing NW shape functions: ')
    c = tic();
    for ii_g = 1:tot_orientations
        if (mod(ii_g, chunk_size/10) == 1)
            currtime = toc(c);
            num_chars = fprintf('%05d/%05d (%g s, ETA %g s)', ...
                ii_g, tot_orientations, currtime, ...
                toc(c) / (ii_g - 1) * tot_orientations - currtime);
        end

        gr = sampler.lattice.gr{ii_g};

        t = gr.allblobs(det_ind).theta(ref_sel);
        n = gr.allblobs(det_ind).eta(ref_sel);
        w = gr.allblobs(det_ind).omega(ref_sel);
        rot_w = gtMathsRotationTensor(w, rotcomp_w);

        n_1 = gr.allblobs(det_ind).eta(ref_sel) + eta_step;
        rot_w_1 = gtMathsRotationTensor(w + om_step, rotcomp_w);

        ominds = gr.allblobs(det_ind).omind(ref_sel);
        ssp = ((ominds == 1) | (ominds == 2));
        ss12  = ssp - ~ssp;

        h = ss12(:, [1 1 1]) .* pl_cry';

        % Extracting ospace boundaries for the given voxel
        w_tab_int = w_tabs_int(:, o_x(ii_g):o_x(ii_g)+1, o_y(ii_g):o_y(ii_g)+1, o_z(ii_g):o_z(ii_g)+1);
        w_tab_int = reshape(w_tab_int, [], 8);

        n_tab_int = n_tabs_int(:, o_x(ii_g):o_x(ii_g)+1, o_y(ii_g):o_y(ii_g)+1, o_z(ii_g):o_z(ii_g)+1);
        n_tab_int = reshape(n_tab_int, [], 8);

        lims_w_proj_g_int = round([min(w_tab_int, [], 2), max(w_tab_int, [], 2)]);
        lims_n_proj_g_int = round([min(n_tab_int, [], 2), max(n_tab_int, [], 2)]);

        num_ws = lims_w_proj_g_int(:, 2) - lims_w_proj_g_int(:, 1) + 1;
        num_ns = lims_n_proj_g_int(:, 2) - lims_n_proj_g_int(:, 1) + 1;

        [ref_r0, ref_r_dir] = compute_rod_line(t, n, rot_w, h, p.labgeo, p.samgeo);
        [ref_r0_n, ref_r_dir_n] = compute_rod_line(t, n_1, rot_w, h, p.labgeo, p.samgeo);
        [ref_r0_w, ref_r_dir_w] = compute_rod_line(t, n, rot_w_1, h, p.labgeo, p.samgeo);

        ref_r0_diff_n = ref_r0_n - ref_r0;
        ref_r0_diff_w = ref_r0_w - ref_r0;

        ref_r_dir_diff_n = ref_r_dir_n - ref_r_dir;
        ref_r_dir_diff_w = ref_r_dir_w - ref_r_dir;

        ref_inv_r_dir = 1 ./ ref_r_dir;

        % Translating everything to the origin
        ref_r0 = ref_r0 - gr.R_vector(ones(numel(t), 1), :);

        if (conf.compute_w_sfs)
            num_dws = num_ws + 1;
            num_dns = num_ns + 1;

            % Forming the grid in orientation space, on the plane that is
            % perpendicular to the axis of rotation (z-axis in the w case),
            % where to sample the sub-regions of the orientation-space voxel
            % that we want to integrate.
            r_vecs_w_samp = [pos_x_grid_w; pos_y_grid_w; pos_z_grid_w];
            r_vecs_w_samp = reshape(r_vecs_w_samp, 3, []);
            r_vecs_w_samp = gr.R_vector(ones_tot_sp_w, :)' + r_vecs_w_samp;

            gs_w = gtMathsRod2OriMat(r_vecs_w_samp);
            % Unrolling and transposing the matrices at the same time
            % because we need g^-1
            gs_w = reshape(gs_w, 3, [])';

            ys = gs_w * pl_cry;
            ys = reshape(ys, 3, [], num_blobs);
            ys = permute(ys, [3 1 2]);

            % Starting here to compute initial shifts aligned with the
            % rounded omega, by frst doing three scalar products
            Ac = sum(Ac_part_w .* ys, 2);
            As = sum(As_part_w .* ys, 2);
            Acc = sum(Acc_part_w .* ys, 2);

            Ac = reshape(Ac, num_blobs, tot_sampling_points_w);
            As = reshape(As, num_blobs, tot_sampling_points_w);
            Acc = reshape(Acc, num_blobs, tot_sampling_points_w);

            D = Ac .^ 2 + As .^ 2;
            sinths = ss12 .* gr.allblobs(det_ind).sintheta(ref_sel);
            sinths = sinths(:, ones(tot_sampling_points_w, 1));

            CC = Acc + sinths;
            DD = D - CC .^ 2;
            E  = sqrt(DD);
            ok = DD > 0;

            % Apparently it is inverted, compared to the omega prediction
            % function. Should be investigated
            ssp = ~((ominds == 1) | (ominds == 3));
            ss13  = ssp - ~ssp;
            ss13 = ss13(:, ones(tot_sampling_points_w, 1));

            % Shift along z in orientation space, to align with the images
            % in proections space
            dz = (-As + ss13 .* ok .* E) ./ (CC - Ac);

            ws_disps = [ ...
                (lims_w_proj_g_int(:, 1) - 0.5 - ref_round_int_ws), ...
                (lims_w_proj_g_int(:, 2) + 0.5 - ref_round_int_ws) ];
            ws_disps = tand(ws_disps .* om_step / 2);
        end

        sf_bls_nw = gtFwdSimBlobDefinition('sf_nw', num_blobs);

        % This is the most expensive operation because the different
        % sizes of the omega spread, which depends on the blob itself,
        % don't allow to have vectorized operations on all the blobs at
        % the same time
        for ii_b = 1:num_blobs
            ws = ((lims_w_proj_g_int(ii_b, 1)-shift_w):step_w:(lims_w_proj_g_int(ii_b, 2)+shift_w)) .* om_step;
            ones_ws = ones(oversampling_w * num_ws(ii_b), 1);

            ns = ((lims_n_proj_g_int(ii_b, 1)-shift_n):step_n:(lims_n_proj_g_int(ii_b, 2)+shift_n)) .* eta_step;
            ones_ns = ones(oversampling_n * num_ns(ii_b), 1);
            % It is always in x, because this is still in lab
            % coords!

            tot_vecs = oversampling_w * num_ws(ii_b) * oversampling_n * num_ns(ii_b);
            ones_tvecs = ones(tot_vecs, 1);

            if (conf.fast_lines)
                vars_ns = (ns - n(ii_b)) / eta_step;
                vars_ns = vars_ns([1 1 1], :);
                vars_ws = (ws - w(ii_b)) / om_step;
                vars_ws = vars_ws([1 1 1], :);

                % The system has already been translated to the origin
                rep_r0s = ref_r0(ii_b, :)';
                rep_r0s = rep_r0s(:, ones_ns, ones_ws);

                vars_r0_ns = vars_ns .* ref_r0_diff_n(ii_b * ones_ns, :)';
                vars_r0_ns = vars_r0_ns(:, :, ones_ws);

                vars_r0_ws = vars_ws .* ref_r0_diff_w(ii_b * ones_ws, :)';
                vars_r0_ws = reshape(vars_r0_ws, 3, 1, []);
                vars_r0_ws = vars_r0_ws(:, ones_ns, :);

                % r0 should be defined = (h x y) / (1 + h . y)
                approx_r0 = rep_r0s + vars_r0_ns + vars_r0_ws;
                r0 = reshape(approx_r0, 3, [])';

                rep_r_dirs = ref_r_dir(ii_b * ones_tvecs, :);

                vars_r_dir_ns = vars_ns .* ref_r_dir_diff_n(ii_b * ones_ns, :)';
                vars_r_dir_ns = vars_r_dir_ns(:, :, ones_ws);
                vars_r_dir_ns = reshape(vars_r_dir_ns, 3, [])';

                vars_r_dir_ws = vars_ws .* ref_r_dir_diff_w(ii_b * ones_ws, :)';
                vars_r_dir_ws = reshape(vars_r_dir_ws, 3, 1, []);
                vars_r_dir_ws = vars_r_dir_ws(:, ones_ns, :);
                vars_r_dir_ws = reshape(vars_r_dir_ws, 3, [])';

                d_vers_nw = vars_r_dir_ns + vars_r_dir_ws;

                % r should be defined = r0 + t * (h + y) / (1 + h . y)
                r_dir = rep_r_dirs + d_vers_nw;
            else
                y_lab = gtGeoPlLabFromThetaEta(t(ii_b * ones_ns), ns', p.labgeo);
                rot_ws = gtMathsRotationTensor(ws, rotcomp_w);
                rot_ws = reshape(rot_ws, 3, []);
                y_sam = y_lab * rot_ws;
                y_sam = reshape(y_sam, oversampling_n * num_ns(ii_b), 3, oversampling_w * num_ws(ii_b));
                y_sam = permute(y_sam, [1 3 2]);
                y_sam = reshape(y_sam, [], 3);

                renorms = 1 + y_sam * h(ii_b, :)';
                renorms = renorms(:, [1 1 1]);

                hh = h(ii_b * ones_tvecs, :);

                % r0 is defined = (h x y) / (1 + h . y)
                r0 = gtMathsCross(hh, y_sam) ./ renorms - gr.R_vector(ones_tvecs, :);

                % r is defined = r0 + t * (h + y) / (1 + h . y)
                r_dir = (hh + y_sam) ./ renorms;
%                 r_dir = gtMathsNormalizeVectorsList(r_dir);
            end

            % Translating r0s to the closest point to the average R
            % vector, on the line r
            dot_r0_avgR = sum((0 - r0) .* r_dir, 2);
            r0 = r0 + r_dir .* dot_r0_avgR(:, [1 1 1]);

            % Let's now find directions and limits
            minus_dirs = r_dir < 0;
            u_lims = half_rspace_sizes(ones_tvecs, :);
            u_lims(minus_dirs) = -u_lims(minus_dirs);
            l_lims = -u_lims;

            if (conf.fast_inverse)
                inv_r_dir = ref_inv_r_dir(ii_b * ones_tvecs, :);
                tmp_d_inv_r = d_vers_nw .* inv_r_dir;

                % Approximate inverse: taylor expansion to the 1st
                % order, which is NOT good enough close to zero
%                 inv_r_dir = (1 - tmp_d_inv_r) .* inv_r_dir;

                % Approximate inverse: taylor expansion to the 2nd
                % order, which seems to be good enough close to zero
                inv_r_dir = (1 - (1 - tmp_d_inv_r) .* tmp_d_inv_r) .* inv_r_dir;
            else
                inv_r_dir = 1 ./ r_dir;
            end

            len_to_ulims = (u_lims - r0) .* inv_r_dir;
            len_to_llims = (l_lims - r0) .* inv_r_dir;

            len_to_ulims = min(len_to_ulims, [], 2);
            len_to_llims = max(len_to_llims, [], 2);

            n_ints = len_to_ulims - len_to_llims;
            n_ints(n_ints < 0) = 0;
            n_ints = reshape(n_ints, [oversampling_n, num_ns(ii_b), oversampling_w, num_ws(ii_b)]);
            n_ints = sum(sum(n_ints, 1), 3);
            n_ints = reshape(n_ints, [num_ns(ii_b), num_ws(ii_b)]);
            sum_n_ints = gtMathsSumNDVol(n_ints);
            n_ints = n_ints ./ sum_n_ints;

            n_ints = cast(n_ints, conf.data_type);

            if (conf.compute_w_sfs)
                d = (ws_disps(ii_b, 1):om_step_rod:ws_disps(ii_b, 2))';
                d = d(:, ones(tot_sampling_points_w, 1)) + dz(ii_b * ones(num_dws(ii_b), 1), :);

                w_upper_limits = min(d(2:end, :), half_rspace_sizes(3));
                w_lower_limits = max(d(1:end-1, :), -half_rspace_sizes(3));

                % integration is simple and linear: for every w and
                % xy-position, we consider the corresponding segment on the
                % z-direction
                w_ints = w_upper_limits - w_lower_limits;
                w_ints(w_ints < 0) = 0;
                w_ints = renorm_coeff_w * sum(w_ints, 2)';

                disp([sum(n_ints, 1), w_ints])
                pause
            end

            sf_bls_nw(ii_b).bbnim = lims_n_proj_g_int(ii_b, :) .* eta_step;
            sf_bls_nw(ii_b).bbwim = lims_w_proj_g_int(ii_b, :);
            sf_bls_nw(ii_b).intm = n_ints;
            sf_bls_nw(ii_b).bbsize = [num_ns(ii_b) num_ws(ii_b)];
        end

        sf{ii_g} = sf_bls_nw;

        if (mod(ii_g, chunk_size/10) == 1)
            fprintf(repmat('\b', [1, num_chars]));
        end
    end
    fprintf(repmat(' ', [1, num_chars]));
    fprintf(repmat('\b', [1, num_chars]));
    fprintf('Done in %g seconds.\n', toc(c))
end

function [r0, r_dir] = compute_rod_line(thetas, etas, rot_w, h, labgeo, samgeo)
    y_lab = gtGeoPlLabFromThetaEta(thetas, etas, labgeo);
    y_sam = gtGeoLab2Sam(y_lab, rot_w, labgeo, samgeo, true);

    renorms = 1 + sum(y_sam .* h, 2);
    renorms = renorms(:, [1 1 1]);

    % r0 is defined = (h x y) / (1 + h . y)
    r0 = gtMathsCross(h, y_sam) ./ renorms;

    % r is defined = r0 + t * (h + y) / (1 + h . y)
    r_dir = (h + y_sam) ./ renorms;
    r_dir = gtMathsNormalizeVectorsList(r_dir);
end