% This program solves a convection problem 
% in a square domain [0,2]x[0,3] 

clear; close all; clc


disp(' ')
disp('This program solves a pure convection equation in [0,2]x[0,2] ')

dom = [0,2,0,3]; 
h= cinput('enter the size of elements',0.2);

% Element type
disp(' ')
elem = 0; p = 1; 
referenceElement = SetReferenceElement(elem,p); 

nx = abs(dom(1)-dom(2))/h; 
ny = abs(dom(3)-dom(4))/h;
 
[X,T] = CreateMesh(dom,nx,ny,referenceElement);
figure; PlotMesh(T,X,'b-');

disp(' ')
disp('The following velocity fields can be considered');
disp('   [1] v(x,y) = (-y,x)');
disp('   [2] v(x,y) = (1,0)');
disp('   [3] v(x,y) = (1,1)/sqrt(2)'); 
disp('   [4] v(x,y) = (-x,-y)');
velo = cinput('Choose the convection field to be used on the computation',4); 
Conv = ComputeVelocity(X,velo);

figure(1); hold on; 
quiver(X(:,1),X(:,2),Conv(:,1),Conv(:,2))
plot(dom([1,2,2,1,1]), dom([3,3,4,4,3]), 'k')
axis equal

% Time discretization
tEnd = cinput('End time', 3);
nStep = cinput('Numnber of time-steps', 120);
dt = tEnd / nStep; 
% Courant = a*dt/h; 
% disp(['Courant number: ', num2str(Courant)]); 

 
% Initial condition
u0 = InitialCondition(X); 

xx  = reshape(X(:,1), nx+1, ny+1)'; 
yy  = reshape(X(:,2), nx+1, ny+1)'; 
u0_dib = reshape(u0, nx+1, ny+1)';  
figure(2); clf; 
surface(xx,yy,u0_dib,'FaceColor','interp');
set(gca,'FontSize',16)
grid on
view(3)
figure(3); clf;
[C,h]=contour(xx,yy,u0_dib,[-0.1,-0.01,0.1:0.1:1.0]);
clabel(C,h,'FontSize',12);
axis equal; axis(dom);
set(gca,'FontSize',16,'XTick',-0.5:0.25:0.5,'YTick',-0.5:0.25:0.5)




% Boundary conditions (Lagrange multipliers)
[ADir,bDir,nDir] = BoundaryConditions(X,velo);


% Method used for solving the problem
disp(' ')
disp ('The problem can be solved using one of the following methods: ');
disp ('    [1] Crank-Nicolson ')
disp ('    [2] TG3-2 ')
method = cinput('Method = ',1);

[M,K,C] = FEM_matrices(X,T,Conv,referenceElement); 
[Mout,Cout] = Boundary_matrices(X,T,Conv,referenceElement,velo);
[A,B,B1,methodName] = system(method,M,K,C,Mout,Cout,dt); 




% Solution at each time step
nPt = length(u0);
u = zeros(nPt,nStep+1); 
u(:,1) = u0;
u_sol=cell(nStep,1);
u_sol{1}= reshape(u0, nx+1, ny+1)';


L=[A,ADir';ADir,zeros(nDir,nDir)];



if method == 1
for n = 1:nStep
    Lb = [B*u(:,n); bDir] ;
    Du = L\Lb; 
    u(:,n+1) = u(:,n) + Du(1:end-nDir);
    u_sol{n+1} = reshape(u(:,n+1), nx+1, ny+1)';
end

else
    for n = 1:nStep 
    Lb = [B*u(:,n); bDir] ;
    Du = L\Lb; 
    u_tel = u(:,n) + Du(1:end-nDir);
   
    Lb = [B1*u_tel; bDir] ;
    Du = L\Lb; 
    u(:,n+1) = u_tel + Du(1:end-nDir);
    u_sol{n+1} = reshape(u(:,n+1), nx+1, ny+1)';
        end
    
end
    figure(4);  
% figure(5); clf;
% [C,h]=contour(xx,yy,u_sol,[-0.1,-0.01,0.1:0.1:1.0]);
% clabel(C,h,'FontSize',12);
% axis equal; axis(dom);
% set(gca,'FontSize',16,'XTick',-0.5:0.25:0.5,'YTick',-0.5:0.25:0.5)


%%
for n = [1:round(nStep/30):nStep+1,nStep+1]
    clf;
    tt = (n-1)*dt;
   surface(xx,yy,u_sol{n},'FaceColor','interp');
set(gca,'FontSize',16)
zlim([-10,10])
grid on
view(3)
    pause()
end
