function [Af,RHS,Arhs] = mk_DiffusionOperator(Nx, Ny, dx, dy, bcsD, bcsN, Node, stagr)
% function [Af,RHS,Arhs] = mk_DiffusionOperator(Nx, Ny, dx, dy,...
%                                           bcsD, bcsN, Node, stagr)
% This function creates the diffusion operator and its and RHS 
%
% INPUTS:
%       Nx:     Number of interior tracer nodes in the x-direction
%       Ny:     Number of interior tracer nodes in the y-direction
%       dx:     Step size in x direction
%       dy:     Step size in y direction
%       bcsD:   Vector containing the values of dirichlet boundary
%                   conditions. 
%       bcsN:   Vector containing the values of Neumann boundary
%                   conditions multiplied by the correct sign and length
%                   eg. bcsN = gN*dx, where du/dx = gN on the boundary
%       Node:   Is the node numbering matrix. size(Node) = [Ny+2,Nx+2]
%               NOTE: Node needs to include boundary conditions. In the 
%                       case for periodic boundaries, the periodicity needs
%                       to be included. See NodePad.m
%       stagr:  String: Which grid to use, 'P', 'u', or 'v'. 
% OUTPUTS:
%       Af:     The A matrix
%       RHS:    The Right-hand-side vector from BC contributions
%       Arhs:   The matrix used to create the Right-hand-side vector
%
% See also: NodePad.m
%
% Authors:   Matt Ueckermann, Pierre Lermusiaux, Pat Haley for MIT course 2.29

%% Preliminaries
%If X and Y increments are required, Uncomment appropriate lines below. For
%staggered u and v grids, the number of interior unknowns require
%modification.
if stagr=='P'
%     x = linspace(dx/2,a-dx/2,Nx);
%     y = fliplr(linspace(dy/2,b-dy/2,Ny));
elseif stagr=='u'
    Nx = Nx-1;
%     x = linspace(dx,   a-dx,   Nx);
%     y = fliplr(linspace(dy/2, b-dy/2, Ny));
else
    Ny = Ny-1;
%     x = linspace(dx/2, a-dx/2, Nx);
%     y = fliplr(linspace(dy,   b-dy,   Ny));
end

% Initialize the A matrix. Notice we are initializing it as a sparse
% matrix in order to reduce memory requirements. For a large number of
% nodes, this sparse matrix allocation would be essential, otherwise we may
% run out of memory.
Nbcs = length(bcsD)+length(bcsN);
TotalNodes = max(Node(:))-Nbcs;

% Initialize the RHS vector (in Ax=b, this is the 'b' vector)
RHS = zeros(TotalNodes,1);
% for i=1: TotalNodes
%     %Add source term contributions
%     RHS(i,1) = 0;
% end

%% Vectorization of loops below
Rowij = Node(2 : Ny + 1, 2 : Nx + 1) - Nbcs;
idsx = 2 : Nx + 1;
idsy = 2 : Ny + 1;
%Create the indices for the entries
I = [Rowij, Rowij, Rowij, Rowij, Rowij];
    
J = [Node(idsy, idsx),...
    Node(idsy, idsx - 1), Node(idsy, idsx + 1),...
    Node(idsy - 1, idsx), Node(idsy + 1, idsx)];
J(I <= 0) = [];
I(I <= 0) = [];
%Create dummy vector to ease building A
OO = ones(length(I(:)) / 5,1);

%Create the A matrix
A = sparse(I(:), J(:),...
    [(2 / dx ^ 2 + 2 / dy^2) * OO; 
    - 1 / dx ^ 2 * OO; - 1 / dx ^ 2 * OO; 
    - 1 / dy ^ 2 * OO; - 1 / dy ^ 2 * OO],...
    TotalNodes, TotalNodes + Nbcs, 5 * TotalNodes + Nbcs * TotalNodes);


%%  *****************   VECTORIZED   ********************
% Also note, this is a different way of implementing RHS vector (different
% that mk_DiffusionOperator.m in problem sets. We use the first Nbcs columns of A as
% a way to "gather" the contributions due to the boundary conditions.
% % A   = spalloc(TotalNodes, TotalNodes + Nbcs, 5*TotalNodes + Nbcs*TotalNodes);
% % 
% % %% Constructing the Inner part of the 'A' Matrix
% % 
% % for j = 2 : Nx + 1 %the range 2 to Nx+1 excludes boundary nodes
% %     for i = 2 : Ny + 1   %the range 2 to Ny+1 excludes boundary nodes
% %         %Now we stamp the computational cell into the 'A' Matrix
% % 
% %         %The rows correspond to the equation for solving the ij'th node.
% %         %First we obtain this node's number from the 'Node' matrix.
% %         Rowij = Node(i,j) - Nbcs;
% %         %Only do Rows of A for INTERIOR points. The first Nbcs columns will
% %         %include contributions from the BCs
% %         if Rowij > 0
% %             %The columns in the current row correspond to the variables used
% %             %in the equation for solving the ij'th node.
% % 
% %             %Start with the middle node
% %             A(Rowij, Node(i, j) ) =A(Rowij, Node(i, j) )...
% %                 +2/dx^2+2/dy^2; 
% % 
% %             %Now do the east component in the computational cell, and proceed
% %             %w-e, n-s;
% % 
% %             A(Rowij, Node(i, j-1) ) = A(Rowij, Node(i, j-1) )...
% %                 -1/dx^2; 
% %             A(Rowij, Node(i, j+1) ) = A(Rowij, Node(i, j+1) )...
% %                 -1/dx^2; 
% %             A(Rowij, Node(i-1, j) ) = A(Rowij, Node(i-1, j) )...
% %                 -1/dy^2; 
% %             A(Rowij, Node(i+1, j) ) = A(Rowij, Node(i+1, j) )...
% %                 -1/dy^2;
% %         end
% %     end
% % end

%% Now, constructing the Dirichlet boundary conditions is rather elegant:
%A quick formatting check
if ~isempty(bcsD)
    if size(bcsD,1)>size(bcsD,2)
        bcsD = bcsD';
    end
    %Then we add the boundary condition contributions to the RHS vector
    RHS = RHS - A( : , 1:Nbcs-length(bcsN) ) * bcsD';
end
% Neumann boundary conditions require a little more work. For Neumann on
% the west face, we have u_w = u_p + dx*bcN. We subtitute this in the
% governing equation, and we get a RHS contribution, AND an additional
% entry on the diagonal for u_p:
%A quick formatting check
if ~isempty(bcsN)
    if size(bcsN,1)>size(bcsN,2),   bcsN = bcsN'; end
    %Then we add the boundary condition contributions to the RHS vector
    RHS = RHS - A( : , length(bcsD)+1:Nbcs) * bcsN';
    %NOTE, we could combine the above two steps (i.e
    % RHS = RHS - A( : ,1:Nbcs) * [bcsD, bcsN]';
    A(:,Nbcs+1:end) =A(:,Nbcs+1:end) + diag(sum(A(:,length(bcsD)+1:Nbcs),2));
end
%Finally, distribute entries of A to outputs
Arhs=A(:,1:Nbcs);
Af  =A(:,Nbcs+1:end);
