% ---------- This program solves Stokes problem FA 2nd question -----------

clear; close all; clc
addpath('Func_ReferenceElement')

%---------------------------- DATA ---------------------------------------- 
dom = [0,1,0,1]; 
mu = 10^3; % pN*s/ m  

%-------------------- Elements and Inteporlation degree -------------------
% (0: quadrilaterals, 1: triangles, 11: triangles with bubble function)
elemV = 0; 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

%------------------------- SET THE ELEMENT DATA ---------------------------
referenceElement = SetReferenceElementStokes(elemV,degreeV,elemP,degreeP); 

nx = 10; % NUMBER OF ELEMENTS IN THE X DIRECTION
ny = 20; % NUMBER OF ELEMENTS IN THE Y DIRECTION

%[X,T,XP,TP,alpha] = CreateMeshesz(nx,ny,referenceElement);
[X,T,alpha] = createMesh_velocity(ny,nx); % MESH DATA FOR VELOCITIY 
[XP,TP] = createMesh(ny,nx);        % MESH DATA FOR PRESSURE

figure (1); PlotMesh(T,X,elemV,'b-'); % PLOT THE MESH FOR VELOCITY
%hold on 
figure (2); PlotMesh(TP,XP,elemP,'r-');       % PLOT THE MESH FOR PRESSURE

%------------ Matrices arising from the discretization --------------------

[K,G,f] = StokesSystem(X,T,XP,TP,referenceElement);
%[K,G,f] = StokesSystem2(X,T,XP,TP,referenceElement);

K = mu*K; % Viscocity Matrix
[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,nx,ny,alpha);

%Total system of equations
if confined
   nunkP = ndofP-1;
   disp(' ')
   disp('Confined flow. Pressure on lower left corner is set to zero');
   G(1,:) = [];
else
   nunkP = ndofP;
end
Atot = [K          A_DirBC'             G'
        A_DirBC    zeros(nDir,nDir)     zeros(nDir,nunkP)
        G          zeros(nunkP,nDir)    zeros(nunkP,nunkP)];
btot = [f ; b_DirBC ; zeros(nunkP,1)];

sol = Atot\btot; 

% Postprocess
velo = reshape(sol(1:ndofV), 2, [])'; 
if confined
    pres = [0; sol(ndofV+nDir+1:ndofV+nDir+nunkP)]; 
else
	pres = sol(ndofV+nDir+1:ndofV+nDir+nunkP); 
end

nPt = size(X,1); 
figure (3); 
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); 

PlotResults(X,T,velo(:,1),referenceElement.elemV,referenceElement.degreeV)

PlotResults(X,T,velo(:,2),referenceElement.elemV,referenceElement.degreeV)

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