%% Vorticity Stream function
   clear all %clear all variables
   clc % clear output screen

   %Constants
   dpic	  = 0.5;
   tfinal = 3.0;
   mu	  = 0.1;  %i.e. Re=10
   nx	  = 10+1;
   ny	  = nx;  %different nx & ny gives non-square
   itmax  = 20;
   
   w	= zeros(nx,ny); % form a matrix of zeros (1..nx, 1..ny) 
   p = w; % w = vorticty, p = streamfunction
   wt = w; % wt = dw/dt
   
   %Initialization
   h = 1/(nx-1); %h = dx (=dy); so x-size is 1, y-size is ny/nx
   t = 0;
   r = 2/(1+ pi/(nx-1)); %SOR parameter
   dt = 0.2*h*h/mu; 
   
   for i= 1:nx
      for j=1:ny
	      p(i,j) = 0;
	      w(i,j) = 0;
      end
   end
   
   %% Main program starts here
   tpic = t + dpic - 0.5*t;
   while(t<tfinal)
      while(t<tpic)
           % Main Poisson solver: solve nabla^2 p = - w
           for it = 1:itmax % iterations start here
              for i = 2:nx-1 
                  for j = 2:ny-1
                    p(i,j) = (1-r)*p(i,j) + r*0.25*(p(i-1,j) + p(i+1,j) + ...
                            p(i,j-1) + p(i,j+1) + h*h*w(i,j) );  
                  end
              end
           end % iterations end here
       
           % Time-step vorticity
           
           % BC 1st order
           for i= 2:nx-1
              w(i,1) =  -((p(i,2)-p(i,1))/h - 0)/(0.5*h);
              w(i,ny)= -(sin(pi*(i-1)*h)*sin(pi*(i-1)*h)... 
                        -(p(i,ny)-p(i,ny-1))/h)/(0.5*h);
           end
           for j= 2:ny-1
              w(1,j) = ((-(p(2,j)-p(1,j))/h) - 0)/(0.5*h);
              w(nx,j) = (0 - (-(p(nx,j)-p(nx-1,j))/h))/(0.5*h);
           end
           
           %BC 2nd order correction
           for i = 2:nx-1
              w(i,1) =  (4*w(i,1)  - w(i,2))/3;
              w(i,ny) = (4*w(i,ny) - w(i,ny-1))/3;
           end
           for j = 2:ny-1
              w(1,j) =  (4*w(1,j)  - w(2,j))/3;
              w(nx,j) = (4*w(nx,j) - w(nx-1,j))/3;
           end
        
           for i = 2:nx-1
              for j = 2:ny-1  
                wt(i,j) =-(p(i,j+1)-p(i,j-1))*(w(i+1,j)-w(i-1,j))/(4*h*h)...
                    + (p(i+1,j)-p(i-1,j))*(w(i,j+1)-w(i,j-1))/(4*h*h)...
                    + mu*(w(i+1,j)+w(i-1,j)+w(i,j+1)+w(i,j-1)-4*w(i,j))/(h*h);
              end
           end
        
           for i = 2:nx-1
              for j = 2:ny-1
                w(i,j) = w(i,j) + dt*wt(i,j);       
              end
           end
           t = t + dt;
           % End of BC
        endwhile
   
%       % Force calculation   
%       q = 0;
%       for i = 2:nx-1 
%          q = q + (2*p(i,ny) - 5*p(i,ny-1) + 4*p(i,ny-2) -p(i,ny-3))/h;
%          %1st order: (p[i,ny] - 2*p[i,ny-1] + p[i,ny-2])/h;
%          %2nd order:(2*p[i,ny] - 5*p[i,ny-1] + 4*p[i,ny-2] -p[i,ny-3])/h;
%       end
%       q = q*mu;
%       formatSpec1='t=%10.4f; q=%10.5f; q*sqrt(t)=%10.4f\n ';
%       fprintf(formatSpec1, t, q, q*sqrt(t));
%       % End of Force calculation
       
       formatSpec2='t=%10.5f; w(0.5, 0.5)=%10.5f\n ';
       fprintf(formatSpec2, t, w(round(nx/2),round(ny/2)));
       tpic = t + dpic - 0.5*dt;
    endwhile
    
    %Plot 'p', streamfunction 
    figure(1)   
    tx = linspace (0, (nx-1)*h, nx)';
    ty = linspace (0, (ny-1)*h, ny)';
    
    [xx, yy] = meshgrid (tx, ty);
    mesh (tx, ty, p');
    xlabel ("x");
    ylabel ("y");
    zlabel ("p(x,y)");
    title ("Streamfunction Plot");
    % End of plot
    
    %Plot 'v'
    figure(2)    
    tx = linspace (0.5*h, (nx-1+0.5)*h, nx-1)';
    ty = linspace (0, (ny-1)*h, ny)';
    for i=1:nx-1;
      for j = 1:ny
        tz(j,i) = (-p(i,j)+p(i+1,j))/h;
      end
    end
    [xx, yy] = meshgrid (tx, ty);
    mesh (tx, ty, tz);
    xlabel ("x");
    ylabel ("y");
    zlabel ("v(x,y)");
    title ("Velocity Plot");
    % End of plot
    
    
    
    % midu: 'u' along x=0.5
    i = round(nx/2);
    for j = 2:ny
      formatSpec3='y=%10.5f; u=%10.5f\n ';
      fprintf(formatSpec3, h*(j-1-0.5), (p(i,j)-p(i,j-1))/h);
    end
    
    % Final output
    formatSpec4='h=%10.5f; h*h=%10.5f; w(0.5,0.5)=%10.5f\n ';
    fprintf(formatSpec4, h, h*h, w(round(nx/2),round(ny/2)));
      
   
