% This program solves Stokes problem in a square domain


clear; close all; clc

addpath('Func_ReferenceElement')

dom = [0,2,0,3]; 

Re = 1000;
nu = 1/Re;

% Element type and interpolation degree
% (0: quadrilaterals, 1: triangles, 11: triangles with bubble function)
elemV = 1; degreeV = 2; degreeP = 1;
% elemV = 1; degreeV = 2; degreeP = 1;
% elemV = 11; degreeV = 1;  degreeP = 1; 
if elemV == 11
    elemP = 1; 
else
    elemP = elemV; 
end

referenceElement = SetReferenceElementStokes(elemV,degreeV,elemP,degreeP); 

nx = cinput('Number of elements in X direction',10);
ny = cinput('Number of elements in Y direction',15); 
adapted = 0; 
[X,T,XP,TP] = CreateMeshes(dom,nx,ny,referenceElement,adapted);

figure; PlotMesh(T,X,elemV,'b-');
figure; PlotMesh(TP,XP,elemP,'r-');

% Matrices arising from the discretization
[K,G,f] = Stokes_system(X,T,XP,TP,referenceElement);
K = K*nu; 
[ndofP,ndofV] = size(G); 

% Matrix and r.h.s vector to impose Dirichlet boundary conditions using
% Lagrange multipliers
[A_DirBC, b_DirBC, nDir, confined] = BC(X,dom,ndofV);

if confined
   nunkP = ndofP-1;
   disp(' ')
   disp('Confined flow. Pressure on lower left corner is set to zero');
   G(1,:) = [];
else
   nunkP = ndofP;
end
nunkV = ndofV; 
btot = [f ; b_DirBC ; zeros(nunkP,1)];

% Initial solution
disp(' ')
IniVelo_file = input('.mat file with the initial velocity = ','s'); 
tic
if isempty(IniVelo_file)
    velo = zeros(nunkV/2,2); 
    y2 = dom(4); 
    nodesY2 = find(abs(X(:,2)-y2) < 1e-6); 
    velo(nodesY2,1) = 1; 
else
    load(IniVelo_file); 
end
lambda = zeros(nDir,1); 
pres = zeros(nunkP,1);
sol0  = [reshape(velo',nunkV,1);lambda;pres(1:nunkP)];

iter = 0; tol = 0.5e-4; 
while iter < 1000
    fprintf('Iteration = %d\n',iter);
    C = ConvectionMatrix(X,T,referenceElement,velo);
    Atot = [K+C        A_DirBC'             G'
            A_DirBC    zeros(nDir,nDir)     zeros(nDir,nunkP)
            G          zeros(nunkP,nDir)    zeros(nunkP,nunkP)];
    % Computation of residual
    res = btot - Atot*sol0;
    % Computation of velocity and pressure increment
    solInc = Atot\res;
    
    % Update the solution
    veloInc = solInc(1:nunkV); 
    lambdaInc = solInc(nunkV+1:nunkV+nDir); 
    presInc = solInc(nunkV+nDir+1:nunkV+nDir+nunkP); 
    velo = velo + reshape(veloInc,2,nunkV/2)';
    lambda = lambda + lambdaInc;
    pres = pres + presInc; 
    
    delta1 = max(abs(veloInc)); 
    delta2 = max(abs(res)); 
    fprintf('Velocity increment=%8.6e, Residue max=%8.6e\n',delta1,delta2); 
    if delta1 < tol*max(max(abs(velo))) && delta2 < tol
        fprintf('\nConvergence achieved in iteration number %g\n',iter); 
        break
    end
    sol0 = [reshape(velo',nunkV,1); lambda; pres]; 
    
    iter = iter + 1; 
end

% Postprocess
if confined
    pres = [0; pres]; 
end

nPt = size(X,1); 
figure; 
quiver(X(1:nPt,1),X(1:nPt,2),velo(1:nPt,1),velo(1:nPt,2));
hold on 
plot(dom([1,2,2,1,1]),dom([3,3,4,4,3]),'k')
axis equal; axis tight

PlotStreamlines(X,velo,dom); 

if degreeP == 0
    PlotResults(X,T,pres,referenceElement.elemP,referenceElement.degreeP)
else
    PlotResults(XP,TP,pres,referenceElement.elemP,referenceElement.degreeP)
end
toc