% Onni Pohjavirta 04.04.2019

% Revision history:
%   - 02.01.2020 Cleaning up the code

%   2D p-FEM solver for a PDE with given weak form
%   applying zero boundary conditions
%   (expandable to other boundary conditions)
%
%       This is a FEM solver which uses p'th order polynomial basis
%       defined on quadrilateral elements. The mesh is imported in
%       standard form: matrices {p, e, t} where
%           -p holds coordinates of mesh nodes         (2-by-x)
%           -e holds indeces of edge nodes             (1-by-x)
%           -t holds indeces of nodes in each element  (4-by-x)
%
%


%% DEFINE THE FOLLOWING:

% Mesh via matrices p e t

    % For example:
    % L-domain - w/ refined corner
    p = [0.5 0; 0.75 0; 1 0; 1 0.25; 1 0.5; 1 0.75; 1 1; 0.75 1; 0.5 1; ...
        0.25 1; 0 1; 0 0.75; 0 0.5; 0.25 0.5; 0.5 0.25; ...
        0.75 0.25; 0.75 0.5; 0.75 0.75; 0.5 0.75; 0.25 0.75; ...
        0.5 0.45; 0.55 0.45; 0.55 0.5; 0.55 0.55; ...
        0.5 0.55; 0.45 0.55; 0.45 0.5; 0.5 0.5]';
    t = [1 2 16 15; 2 3 4 16; 4 5 17 16; 5 6 18 17; 6 7 8 18; ...
         8 9 19 18; 9 10 20 19; 10 11 12 20; 12 13 14 20; ...
         15 16 22 21; 16 17 23 22; 17 18 24 23; 18 19 25 24; ...
         19 20 26 25; 20 14 27 26; 21 22 23 28; 23 24 25 28; ...
         25 26 27 28]';
     e = [1:14 27 28 21 15];


% Path of auxiliary FEM functions
    addpath('C:\Users\Onni\Dropbox\Opinnot_Aallossa\FiniteElementMethod\ProjectClean\functions');

% Degree of polynomial basis
    pp = 1;

% Load function
    f = @(x,y) sin(pi*x).*sin(pi*y);

% Weak form B(u,v) = L(v):
    B = @(Xp,U,V,dU,dV) dU{1}.*dV{1} + dU{2}.*dV{2};
    L = @(Xp,V,dV) f(Xp{1},Xp{2}).*V;

% Exact solution - for error estimation (if known)
    fex = @(x,y) 1/(2*pi^2) * sin(pi*x).*sin(pi*y);
    fexdx = @(x,y) 1/(2*pi) * cos(pi*x).*sin(pi*y);
    fexdy = @(x,y) 1/(2*pi) * sin(pi*x).*cos(pi*y);




%% INITIALIZE

Nel = size(t,2);                % #elements
Nn  = size(p,2);                % #nodes
Nb  = 4 + 4*(pp-1) + (pp-1)^2;  % #local basis functions
N   = pp+2;                     % #1D quadrature points

% Fix node indeces in t to go anti-clockwise in each element
t = fixMeshIndeces(t,p);
p1 = p(1,:); p2 = p(2,:);   % Visualize
figure; patch(p1(t), p2(t), 0)

% 1D Quadrature points and weights
[qp1D,w1D] = gaussint(N, -1, 1);

% 2D Quadrature points and weights:
qp = {repmat(qp1D , 1, N) , repmat(qp1D', N, 1)};
w = w1D*w1D';

% Legendre pol. & der. coefficients
legc  = legendrepol(pp);
legdc = legc(:,2:end).*(1:pp);  % shift and scale

% Basis functions & their gradients (derivatives) at quadrature points
ph  = evaluateBasis(qp{1}, qp{2}, pp, legc);
dph = evaluateBasisGradients(qp{1}, qp{2}, pp, legc, legdc);

% Indeces of sides of elements in the mesh  &  number of sides
[s, Ns] = t2si(t,Nn);

% Index of global basis function corresponding to (k,j):
% j'th local basis function in k'th element
ind = lbi2gbi(t,s,pp,Nn,Ns,Nel);

% Edge function indeces (for setting boundary values later)
xei = edgefunind(pp,t,e,ind);

% Parity indeces (-1,+1) for (k,j) = k'th element, j'th (loc.) b. function
par = parityVec(pp,s,t,Ns,Nel,Nb);




%% ASSEMBLE

% Dimension of the global polynomial space
Np = Nn + Ns*(pp-1) + Nel*(pp-1)^2;

% Initialize A & b and find entries by looping over elements
gdph = cell(length(dph),Nel); gqp = cell(2,Nel); Jdets = cell(1,Nel);
b = zeros(Np, 1); A = sparse(Np, Np);
for k = 1:Nel

    % Current node indeces and nodes in global coordinates
    tc = t(:,k);
    P = p(:,tc);

    % Global quadrature points and f evaluated in them
    gqp(:,k) = F(qp{1}, qp{2}, P);

    % Jacobian matrix in qp's (2x2 cell) and determinants (square mat)
    J = JF(P, dph(1:4));
    Jdets{k} = J{1,1}.*J{2,2} - J{2,1}.*J{1,2};

    % Gradients of the global basis functions (manual inv. of 2x2 Jacobian)
    for j = 1:length(dph)
    gdph{j,k} = {(J{2,2}.*dph{j}{1}  -J{2,1}.*dph{j}{2})./Jdets{k}, ...
                (-J{1,2}.*dph{j}{1}  +J{1,1}.*dph{j}{2})./Jdets{k}};
    end

    % Assemble matrix
    for i = 1:Nb
        for j = 1:Nb
            pm = par(k,i)*par(k,j); % plus/minus 1, from parity fix
            Ae = sum(sum( pm*abs(Jdets{k}).*w.* B(gqp(:,k),ph{i},ph{j},gdph{i,k},gdph{j,k}) ));
            A(ind(k,i),ind(k,j)) = A(ind(k,i),ind(k,j)) + Ae;
        end
    end

    % Assemble load vector
    for j = 1:Nb
        b(ind(k,j)) = b(ind(k,j)) + sum(sum(par(k,j)*abs(Jdets{k}).*w.* L(gqp(:,k),ph{j},gdph{j}) ));
    end

end



%% BOUNDARY CONDITIONS AND SOLVING

% Zero-boundary conditions, solve, assemble global solution
xg = zeros(Np,1);
xsi = setdiff(1:Np, xei);   % Complement of boundary functions
xg(xsi) = A(xsi,xsi)\b(xsi);



%% PLOTTING & ERROR ESTIMATION

% For more explanation, see the loop in "ASSEMBLY" section
%

% Plot points on ref. elem.
npp = 20;
[xloc,yloc] = meshgrid(linspace(-1,1,npp), linspace(-1,1,npp));

% corresponding phi values
pph = evaluateBasis(xloc, yloc, pp, legc);

figure; hold on;
X = []; Y = []; Zfem = []; Zdxfem = []; Zdyfem = []; Jd = [];
for k = 1:Nel
    tc = t(:,k);
    P = p(:,tc);

    % Global plot and quadrature points
    xyg = F(xloc,yloc,P);

    % FEM solution values
    solp = zeros(length(xloc));
    solqp = zeros(size(w));
    soldxqp = solqp; soldyqp = solqp;

    for i = 1:Nb
        solp = solp + xg(ind(k,i))*par(k,i)*pph{i};
        solqp   = solqp   + xg(ind(k,i))*par(k,i)*ph{i};
        soldxqp = soldxqp + xg(ind(k,i))*par(k,i)*gdph{i,k}{1};
        soldyqp = soldyqp + xg(ind(k,i))*par(k,i)*gdph{i,k}{2};
    end

    X = [X gqp{1,k}];
    Y = [Y gqp{2,k}];
    Zfem   = [Zfem solqp];
    Zdxfem = [Zdxfem soldxqp];
    Zdyfem = [Zdyfem soldyqp];
    Jd = [Jd Jdets{k}];

    surf(xyg{1}, xyg{2}, solp); % Plot

end

% Evaluate exact solution (if known)
Zex   = fex(X,Y);
Zdxex = fexdx(X,Y);
Zdyex = fexdy(X,Y);

L2err = sqrt(sum(sum( abs(Jd).*repmat(w,1,Nel).*( Zex-Zfem ).^2 )));
L2graderr = sqrt(sum(sum( abs(Jd).*repmat(w,1,Nel).*( (Zdxex-Zdxfem).^2 + (Zdyex-Zdyfem).^2 ) )));
%}