function G = nodal_gradient_matrix(mesh)
%NODAL_GRADIENT_MATRIX Matrix that converts nodal values to gradients
%   G = NODAL_GRADIENT_MATRIX(MESH) Computes a matrix G for converting
%       nodal values in a finite element mesh into nodal gradients. If P
%       is a column vector with the coefficients of a discretized scalar 
%       field over MESH, then the discretized gradient field in the nodes
%       is found as DP = G * P. The rows of the vector DP will contain 
%       the gradient in each node with the x, y, z components being in the 
%       consecutive rows of the vector.
%       The size of the matrix G is (d*N)xN, where d is the number of 
%       dimensions of the mesh (1, 2, or 3).
%
%   The nodal gradient matrix is computed by evaluating the gradients of 
%   the shape functions in each corner of all elements and then averaging
%   the resulting gradients for each node. If a node is a corner of E
%   elements, the computed gradient is 1/E times the sum of the gradients
%   found for the E elements.

%   Peter Rucz, 2021.

% Vector of possible element types
elemTypes = [12, 23, 24, 34, 36, 38];

% Corner coordinates for each element type
corners = { ...
    [-1  1];                                    % 12 - Line
    [ 0  0;  1  0;  0  1];                      % 23 - Tria
    [-1 -1;  1 -1;  1  1; -1  1];               % 24 - Quad
    [ 0  0  0;  1  0  0;  0  1  0;  0  0  1];   % 34 - Tetra
    [ 0  0 -1;  1  0 -1;  0  1 -1; ...
      0  0  1;  1  0  1;  0  1  1];             % 36 - Wedge (penta)
    [-1 -1 -1;  1 -1 -1;  1  1 -1; -1  1 -1; 
     -1 -1  1;  1 -1  1;  1  1  1; -1  1  1];   % 38 - Hexa
};

% Drop element identifiers
elements = drop_IDs(mesh);

% Mesh dimensionality
dim = unique(floor(elements(:, 2) / 10));

% Check if mesh dimension is homogeneous
if numel(dim) ~= 1
    error('nodal_gradient_matrix:inconsistent_mesh', ...
        'Mesh has element types with different dimensions');
end

nNodes = size(mesh.Nodes, 1);               % Number of nodes in mesh
nElems = size(elements, 1);                 % Number of elements in mesh
nElemNodes = mod(elements(:,2), 10);        % Number of nodes per element
nElemValues = sum(nElemNodes.^2 * dim);     % Number of values per element

% Sparse storage preallocation
R = zeros(nElemValues, 1);                  % Row indices
C = zeros(nElemValues, 1);                  % Column indices
V = zeros(nElemValues, 1);                  % Values

% Counter for each node, counting the number of elements the node is a
% corner of.
elemCount = zeros(nNodes, 1);

c = 0;          % Value counter (for filling sparse storage)
e = 0;          % Element counter (for progress bar)

% Go through the element types
for iType = 1 : length(elemTypes)
    type = elemTypes(iType);                % Current element type
    nnod = mod(type, 10);                   % Number of nodes for type
    % Check if there are enough columns in the elements matrix
    % (This prevents a possible overindex error for the next check)
    if size(elements, 2) < 4 + nnod
        continue
    end
    % Select the elements of the current type
    elems = elements(elements(:,2) == type, 4 + (1 : nnod));
    % Skip element type, if there are none in the mesh
    if isempty(elems)
        continue
    end
    
    % Pre-compute element type specific quantities
    % Number of values in the sparse matrix
    nval = nnod^2 * dim;
    % Shape function derivatives w.r.t. local coordinates in corners
    [~, dNxi] = shapefun(corners{iType}, type);
    dNxi = reshape(shiftdim(dNxi,2), numel(dNxi)/nnod, nnod);
    % Preallocate for derivatives w.r.t global coordinates in corners
    dNx = zeros(size(dNxi));
    for iElem = 1 : size(elems, 1)
        nodes = elems(iElem, :);
        % Nodal coordinates of the element as columns
        xnode = mesh.Nodes(nodes, 1+(1:dim)).';
        J = dNxi * xnode.';
        for iNode = 1 : size(xnode, 2)
            ind = dim*(iNode-1) + (1:dim);
            Jac = J(ind, :);
            dNx(ind, :) = (Jac) \ dNxi(ind, :);
        end
        % Column indices
        cols = reshape(repmat(nodes, dim*nnod, 1), [], 1);
        % Row indices
        rows = reshape(repmat(bsxfun(@plus, (nodes-1)*dim, (1 : dim).'), 1, nnod), [], 1);
        
        idx = c + (1 : nval).';
        R(idx) = rows;
        C(idx) = cols;
        V(idx) = dNx(:);
        
        % Increase counter
        c = c + nval;
        
        % Increase elemCounter
        elemCount(nodes) = elemCount(nodes) + 1;
        
        e = e + 1;
        progbar(1, nElems, e);
    end % of cycle through elements of one type
end % of cycle through element types

% Nodal gradient matrix, without averaging
G = sparse(R, C, V, nNodes*dim, nNodes);
% Pre-weighting matrix that performs the averaging for each node
W = spdiags(kron(1./elemCount, ones(dim,1)), 0, dim*nNodes, dim*nNodes);

% Final nodal gradient matrix
G = W * G;

end % of function nodal_gradient_matrix


