% Navier Stokes - du/dt + (u.D)u = -(1/rho)Dp + visLPu
% D = divergence LP = Laplacian  u = {u, v}
% in continuous casep satisifes Poisson equation -
%     -1/rho*(d2p/dx2 + d2p/dy2) = (du/dx)^2 + 2du/dydv/dx + (dv/dy)^2
% discretize in time
% u/n+1 = n/n + dt[-u/n.Du/n - (1/rho)Dp/n + visLPu/n]
% set Du/n+1 = 0 to get Poisson eq for p
% D.p/n+1 = D.u/n + dt[-D.u/n.Du/n - (1/rho)D^2p/n + visLP(Du/n)]
% set D.u/n+1 = 0 to get Poisson eq for p
% D^2p/n+1 = rhoD.u/n/dt + [-rhoD.u/n.Du/n  + rho*visLP(Du/n)]
% discretize Poisson for p
%(p/i+1,j - 2p/i,j + p/i-1/j)/dx^2 + (p/i,j+1 - 2p/i,j + p/i/j-1)/dy^2
%    = (u/i+1,j - u/i-1,j)/2dx)^2 
%     + (u/i+1,j - u/i-1,j))/2dy*(v/i+1,j - v/i-1,j))/2dx 
%     + (v/i+1,j - v/i-1,j)/2dy)^2 
%     + (rho/dt)((u/i+1,j - u/i-1,j)/2dx) + (v/i+1,j - v/i-1,j)/2dy)                                  
% ME 702 Video Lesson 10-11
% IC - u,v,p = 0 everywhere
% BC u=1 @ y=2
% u,v = 0  x=0,2 & y=0
% p = 0 @ y = 2, dp/dy = 0 @ y = 0 & dp/dx = 0 @ x = 0,2



% from step 8?
% un+1/i.j = un/i,j + dt*vis*(un/i+1,j - 2un/i,j + un/i-1,j)/dx^2 + dt*vis*(un/i,j+1 - 2un/i,j + un/i,j-1)/dy^2  

%
clear
nx = 30; ny = 15;
nt = 1000;
nu = .1;
dt = .001;
nit= 10;
rho=1;

Vin=3

L=4; H=2;

dx = L/(nx-1); dy = H/(ny-1);

x=0:dx:L; y=0:dy:H;
[Y,X] = meshgrid(y,x);

u=zeros(nx,ny);
v=zeros(nx,ny);
p=zeros(nx,ny);
b=zeros(nx,ny);

nsy = 10;   % y index of step boundary
nsx = 10;   % x index of step boundary, when i <= 4 and j <= 5 the point is not in the channel

%initialize u
for i=1:nx
    for j=nsy+1:ny-1
        u(i,j) = Vin;
    end
end

for n=1:nt   % outer time step loop
    
    for i=2:nx-1    % fill b array (see step 10 for Poission eq)
        if(i <= nsx) j=nsy+1;
           else j=2;
        end
        for j=j:ny-1
             b(i,j) = rho*(1/dt*((u(i+1,j)-u(i-1,j))/(2*dx)+ (v(i,j+1)-v(i,j-1))/(2*dy)) ...
             -  ((u(i+1,j)-u(i-1,j))/(2*dx))^2 ...
              - 2*((u(i,j+1)-u(i,j-1))/(2*dy)*(v(i+1,j)-v(i-1,j))/(2*dx)) ...
              -  ((v(i,j+1)-v(i,j-1))/(2*dy))^2);
              
       
        end
    end
    

    for iit=1:nit   % solve poission eq for p
       pn = p;
       for i=2:nx-1    % fill b array (see step 10 for Poission eq)
        if(i <= nsx) j=nsy+1;
           else j=2;
        end
        for j=j:ny-1
               p(i,j) = [(pn(i+1,j) + pn(i-1,j))*dy^2 + (pn(i,j+1) + pn(i,j-1))*dx^2] ...
                          /  (2*(dx^2+dy^2)) ...
               - dx^2*dy^2/(2*(dx^2+dy^2))*b(i,j);
           end
       end
       p(nx,:) = p(nx-1,:);                  % dp/dx = 0 at x = 2       
       p(1,nsy:ny) = p(2,nsy:ny);            % dp/dx = 0 at x = 0 
       %p(nsx,1:nsy) = p(nsx+1,1:nsy);      
       
       %p(nsx:ny,1) = p(nsx:ny,2);               % dp/dy = 0 at y = 0
       %p(1:nsx,nsy) = p(1:nsx,nsy+1);         % dp/dy = 0 at y = nsy
       %p(:,ny) = p(:,ny-1);                    % dp/dy = 0 at y = 2
       
   
             
    end

    un=u;
    vn=v;
    for i=2:nx-1    % fill b array (see step 10 for Poission eq)
        if(i <= nsx) j=nsy+1;
           else j=2;
        end
        for j=j:ny-1
          u(i,j) = un(i,j) ...
                  - un(i,j)*dt/dx*(un(i,j) - un(i-1,j)) ...
                  - vn(i,j)*dt/dx*(un(i,j) - un(i,j-1)) ...
                  - dt/(2*rho*dx)*(p(i+1,j) - p(i-1,j)) ...
                  + nu*(dt/dx^2*(un(i+1,j) - 2*un(i,j) + un(i-1,j)) ...
                  +    dt/dy^2*(un(i,j+1) - 2*un(i,j) + un(i,j-1)));
                      
          v(i,j) = vn(i,j) ...
                  - un(i,j)*dt/dx*(vn(i,j) - vn(i-1,j)) ...
                  - vn(i,j)*dt/dy*(vn(i,j) - vn(i,j-1)) ...
                  - dt/(2*rho*dy)*(p(i,j+1) - p(i,j-1)) ...
                  + nu*(dt/dx^2*(vn(i+1,j) - 2*vn(i,j) + vn(i-1,j)) ...
                  +    dt/dy^2*(vn(i,j+1) - 2*vn(i,j) + vn(i,j-1)));                       
        end
    end
   
    
     u(1,nsy+1:ny-1)=Vin;     % u = 1 @ x = 0  inflow velocity
    u(1:nsx,nsy) = 0;      % u = 0 on top of step 
    u(nsx:nx,1) = 0;       % u = 0 on floor
    u(:,ny) = 0;            % u = 0 on ceiling
    u(nsx,1:nsy) = 0;       % u = 0 on side of setp
    u(nx,:) = u(nx-1,:);    % du/dx = 0 at outflow
    
    v(1,nsy+1:ny-1)=0;     % v = 0 @ x = 0  inflow velocity
    v(1:nsx,nsy) = 0;      % v = 0 on top of step 
    v(nsx:nx,1) = 0;     % v = 0 on floor
    v(:,ny) = 0;            % v = 0 on ceiling
    v(nsx,1:nsy) = 0;       % v = 0 on side of setp
    v(nx,:) = 0;            % v = 0 at outflow
end

%contour(X,Y,p,alpha=5)
%hold on
%colorbar()
%contour(X,Y,p,alpha=5)
%quiver(X(1:2:nx,1:2:ny),Y(1:2:nx,1:2:ny),u(1:2:nx,1:2:ny),v(1:2:nx,1:2:ny))
S = .1;     % scale u,v
for i=1:nx
    for j=1:ny
        l=sqrt(u(i,j)^2 + v(i,j)^2);
        if l > S
            u(i,j) = S*u(i,j)/l;
            v(i,j) = S*v(i,j)/l;
        end
    end
end

quiver(X,Y,u,v,0)
axis equal
hold on
contour(X,Y,p)
%plot([0,(nsx-1)*dx],[(nsy-1)*dy,(nsy-1)*dy])
%plot([(nsx-1)*dx,(nsx-1)*dx],[0, (nsy-1)*dy])

hold off