% This program solves the 2D cavity flow Stokes problem 

clear; close all; clc

addpath('Func_ReferenceElement')

dom = [0,1,0,1]; 

% Element type and interpolation degree
% (0: quadrilaterals, 1: triangles, 11: triangles with bubble function)

% elemV = 0; degreeV = 2; degreeP = 1; st = 'Q2Q1';    % Q2Q1
elemV = 1; degreeV = 1; degreeP = 1; st = 'P1P1';    % P1P1
% elemV = 0; degreeV = 2; degreeP = 0; st = 'Q2Q0';    % Q2Q0
% elemV = 11; degreeV = 1;  degreeP = 1; st = 'MINI';  % MINI
if elemV == 11
    elemP = 1; 
else
    elemP = elemV; 
end
referenceElement = SetReferenceElementStokes(elemV,degreeV,elemP,degreeP); 

% nx = cinput('Number of elements in each direction',20);
nx = 20;
ny = nx; 
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] = StokesSystem(X,T,XP,TP,referenceElement);
[ndofP,ndofV] = size(G); 

[dofDir,valDir,dofUnk,confined] = BC_red(X,dom,ndofV); 
nunkV = length(dofUnk); 
if confined
   nunkP = ndofP-1;
   disp(' ')
   disp('Confined flow. Pressure on lower left corner is set to zero');
   G(1,:) = [];
else
   nunkP = ndofP;
end

% GLS = cinput('Use GLS stabilization? (1:Yes,0:No)',0);
GLS = 1;
if GLS == 1
    glsq = 'yes';
    tau = cinput('Define tau value',1.d-4);
    [TL,TR] = GLSMethod(XP,TP,tau,referenceElement);
    
    if confined
        TL = TL(2:end,2:end);
        TR = TR(2:end);
    end
else
    glsq = 'no';
    TL = zeros(nunkP);
    TR = zeros(nunkP,1);
end

f = f - K(:,dofDir)*valDir; 
Kred = K(dofUnk,dofUnk); 
Gred = G(:,dofUnk); 
fred = f(dofUnk); 

A = [Kred   Gred'; 
     Gred   -TL  ]; 
b = [fred;  -TR];

sol = A\b; 

velo = zeros(ndofV,1);
velo(dofDir) = valDir; 
velo(dofUnk) = sol(1:nunkV); 
velo = reshape(velo,2,[])'; 
pres = sol(nunkV+1:end); 
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),'AutoScale','on','AutoScaleFactor',2);
hold on 
plot(dom([1,2,2,1,1]),dom([3,3,4,4,3]),'k')
axis equal; axis tight
print(['Figures/Stokes/' st '_quiver_gls' glsq],'-dpng')

PlotStreamlines(X,velo,dom,s); 

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