function [A,rhs,MGData] = discretizePoissonEquation(n, mgLevel, numberOfSmootherSteps, uWithBC)
% Create nested multigrid structure of discretized data

% Convolute the cube to create the correct right-hand side using the same
% 7-way stencil as used in the discretization: 
% add the Left, Right, Forward, Backward, Down, Up neighbor boundary points
kernel = zeros(3, 3, 3);
Center = ceil(3*3*3/2);
Left = -1; Right = 1;
Forward = -3; Backward = 3;
Down = -(3*3); Up = (3*3);
idx = Center + [Left, Right, Forward, Backward, Down, Up];

kernel(idx)= ones(6,1);
C = convn(uWithBC, kernel, 'valid');

% Reshape into right-hand side vector
rhs=reshape(permute(C, [2, 1, 3]), n.X*n.Y*n.Z, 1);

% Create multigridData if requested
if mgLevel>0
    % Create structure for mgDepth+1 MG levels
    MGData = struct('MaxLevel', mgLevel+1, ...
        'Level', [], ...
        'NumberOfSmootherSteps', numberOfSmootherSteps, ...
        'Dimensions', [], ...
        'Matrices', [], ...
        'Fine2CoarseOperator', []);
    MGData = repmat(MGData, mgLevel+1, 1);
    
    MGData(1).Dimensions = n;
    % To speed up computation we explicitly store triangular factors of A
    needTriangularMatrices = true;
    % Create matrices for coarse levels and fine-to-coarse (f2c) operators
    for level = 1:mgLevel
        MGData(level).Level = level;
        MGData(level).Matrices = iGenerateMatrices(MGData(level).Dimensions, ...
            needTriangularMatrices, uWithBC, level);
        [MGData(level).Fine2CoarseOperator, MGData(level+1).Dimensions] = ...
            iCreateFine2CoarseOperator(MGData(level).Dimensions, level);
        % Create only A on coarsest level
        if level==mgLevel
            % On coarsest level, we don't need the triangular factors since
            % we solve with the system matrix directly
            needTriangularMatrices = false;
            MGData(level+1).Level = mgLevel;
            MGData(level+1).Matrices.A = iGenerateMatrices(...
                MGData(level+1).Dimensions, needTriangularMatrices, uWithBC, level+1);
        end
    end
    % Extract main system matrix A explicitly for convenience 
    A = MGData(1).Matrices.A;
else
    % If mgLevel==0, we only create the main system matrix A
    needTriangularMatrices = false;
    MGData = [];
    A = iGenerateMatrices(n, needTriangularMatrices, uWithBC, 1);
end
end

function matrices = iGenerateMatrices(dimensions, needTriangularMatrices, uWithBC, level)

A = iCreate3DGrid(dimensions, uWithBC);

if needTriangularMatrices
    % Create and save triangular structures to speed up computation
    L = tril(A, -1);
    U = triu(A, 1);
    LD = tril(A);
    DU = triu(A);
    D = diag(diag(A));

    matrices = struct('A', A, 'L', L, 'U', U, 'LD', LD, 'DU', DU, 'D', D);
else
    matrices = A;
end
nnzA = nnz(A);
if labindex==1
    fprintf('Level %d: The problem is of dimension %d with %d nonzeros.\n', level-1, size(A,1), nnzA);
end
end % End of generateCodistributedMatrices

function [linIdx, dimCoarse] = iCreateFine2CoarseOperator(dim, mgLevel)
% Helper to create the linear index that maps between the fine and corase
% grid of the x-y-z dimensional discretized cube.

dimCoarse = struct('X', dim.X/2, 'Y', dim.Y/2, 'Z', dim.Z/2);

% Ensure all entries in DIM are a multiple of 2
invalidDimCoarse = find(~structfun(@(x)(ceil(x)==floor(x)), dimCoarse));
if any(invalidDimCoarse)
    dimNames=fieldnames(dim);
    msg = sprintf('MG-Level %d cannot be created for input dimensions.\n', mgLevel);
    for i = invalidDimCoarse
        msg = [msg, sprintf('%s must be a multiple of 2^multigridLevels.\n', dimNames{i})]; %#ok<AGROW>
    end
    error(msg);
end

% Take every 2nd point in x-dimension
xIdx = 1:2:dim.X;

% Add those same points in every 2nd line of y-dimension
xyIdx = xIdx;
for i=1:dimCoarse.Y-1
    xyIdx = [xyIdx xIdx+i*2*dim.X]; %#ok<AGROW>
end

% Add those points in every 2nd plane of the z-dimension
xyzIdx = xyIdx;
for i=1:dimCoarse.Z-1
    xyzIdx = [xyzIdx xyIdx+i*2*(dim.X*dim.Y)]; %#ok<AGROW>
end
linIdx = xyzIdx;
end

function A = iCreate3DGrid(dims, uWithBC)
% Create a sparse array representing a 6-way connected regular 3D grid of
% size MxNxP.

m = dims.X;
n = dims.Y;
p = dims.Z;

% Create element offsets
Left = -1; Right = 1;
Forward = -m; Backward = m;
Down = -(m*n); Up = (m*n);
N = m*n*p;

offsets = [Left Right Forward Backward Down Up];
idxs = (1:N)';

% Form a numel-by-numConnections list of indices
cnx = idxs + offsets;

% We now need to remove connections off the sides
toDelete = false(size(cnx));
toDelete(mod(idxs, m) == 1,               offsets==Left) = true;     % Left
toDelete(mod(idxs-1, m)+1 == m,           offsets==Right) = true;    % Right
toDelete(mod(idxs-1, (m*n)) < m,          offsets==Forward) = true;  % Front
toDelete(mod(idxs-1, (m*n)) >= (m*(n-1)), offsets==Backward) = true; % Back
toDelete(idxs <= (m*n),                   offsets==Down) = true;     % Bottom
toDelete(idxs > (m*n*(p-1)),              offsets==Up) = true;       % Top

idxs = repmat(idxs, [size(cnx,2), 1]);
idxs(toDelete(:)) = [];
cnx(toDelete(:)) = [];

% idxs gives the row, cnx gives the column, to fill with -1. Diagonal is 6.
% idxs is casted like uWithBC, i.e., if we have a remote variable as input,
% idxs is send to remote and the subsequent call to SPARSE will create the
% matrix remotely.
idxs = cast(idxs, 'like', uWithBC);
A = sparse(idxs, cnx, -1, N, N) + 6*speye(N);
end
