% 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
disp('Select Element Type');
disp('   [1] Q2Q0');
disp('   [2] Q2Q1');
disp('   [3] P1P1');
disp('   [4] MINI(P1+P1)');
elem_type = cinput('Choose the convection field to be used on the computation',3);
% (0: quadrilaterals, 1: triangles, 11: triangles with bubble function)
switch elem_type
    case 1
        elemV = 0; degreeV = 2; degreeP = 0;
    case 2
        elemV = 0; degreeV = 2; degreeP = 1;
    case 3
        elemV = 1; degreeV = 1; degreeP = 1;
    case 4
        elemV = 11; degreeV = 1;  degreeP = 1;
end
if elemV == 11
    elemP = 1; 
else
    elemP = elemV; 
end
referenceElement = SetReferenceElementStokes(elemV,degreeV,elemP,degreeP);

% GLS Stabilization for P1P1
if elem_type == 3
    disp('Apply Stabilization using GLS?');
    stbl = cinput('0:NO 1:YES ',0);
else
    stbl = 0;
end

% Stabilization parameter
if stbl == 1
    tau1 = 6.25e-4;
    tau2 = 4;
else
    tau1 = 0;
    tau2 = 0;  
end

nx = cinput('Number of elements in each direction',20);
ny = nx; 
adapted = 1; % 0: uniform, 1:adapted
[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,K2,G,L,f,f2] = StokesSystem(X,T,XP,TP,referenceElement,stbl);
K = K + tau2*K2;
[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,:) = [];
   L(1,:) = [];
   L(:,1) = [];
   f2(1) = [];
else
   nunkP = ndofP;
end

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

if stbl == 1
    A22 = tau1*L;
    b2 = tau1*f2;
else
    A22 = zeros(nunkP);
    b2 = zeros(nunkP,1);
end

A = [Kred   Gred'; 
     Gred   -A22]; 
b = [fred; -b2];

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));
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