program NaverStokesSolver
  implicit none
  integer :: i, j
  integer, parameter :: nx = 60
  integer, parameter :: ny = 30
  double precision, parameter :: t_end = 0.5
  double precision :: t
  double precision, parameter :: dx = 4.0/(nx-1.0)
  double precision, parameter :: dy = 1.0/(ny-1.0)
  double precision, parameter :: nu = 0.025
  double precision, parameter :: rho = 1.0
  double precision, dimension(nx,ny) :: u, v, us, vs, un, vn, P
  double precision, dimension(nx,ny) :: us1, vs1, us2, vs2
  double precision, dimension(nx,ny) :: x, y
  double precision :: dt = 0.003
  double precision, dimension(nx,ny) :: uMtm, vMtm, uMtm1, vMtm1, uMtm2, vMtm2

  !Initial condition
  call CartesianGrid(x,y,dx,dy)
  call InitialCondition(u,v,P)

  t = 0.0
  do while (t <= t_end)
    call Momentum(u, v, uMtm, vMtm)

    ! TVD RK3 step 1
    do j=2,ny-1
        do i=2,nx-1
            us1(i,j) = u(i,j) + dt*uMtm(i,j)
            vs1(i,j) = v(i,j) + dt*vMtm(i,j)
        end do
    end do

    call IntermediateBC(us1,vs1,un,vn,P)
    call Momentum(us1,vs1,uMtm1,vMtm1)

    ! TVD RK3 step 2
    do j=2,ny-1
        do i=2,nx-1
            us2(i,j) = (3.0/4.0)*u(i,j) + (1.0/4.0)*us1(i,j) + (dt/4.0)*uMtm1(i,j)
            vs2(i,j) = (3.0/4.0)*v(i,j) + (1.0/4.0)*vs1(i,j) + (dt/4.0)*vMtm1(i,j)
        end do
    end do

    call IntermediateBC(us2,vs2,un,vn,P)
    call Momentum(us2,vs2,uMtm2,vMtm2)

    ! TVD RK3 step 3
    do j=2,ny-1
        do i=2,nx-1
            us(i,j) = (1.0/3.0)*u(i,j) + (2.0/3.0)*us2(i,j) + ((2.0*dt)/3.0)*uMtm2(i,j)
            vs(i,j) = (1.0/3.0)*v(i,j) + (2.0/3.0)*vs2(i,j) + ((2.0*dt)/3.0)*vMtm2(i,j)
        end do
    end do

    ! Solve for P
    call IntermediateBC(us,vs,un,vn,P)
    call PressureSolver(P,us,vs,un,vn)

    ! Corrector step
    do j=2,ny-1
      do i=2,nx-1
        un(i,j) = us(i,j) - (dt/(2.0*dx*rho))*(P(i+1,j) - P(i-1,j))
        vn(i,j) = vs(i,j) - (dt/(2.0*dy*rho))*(P(i,j+1) - P(i,j-1))
      end do
    end do

    call BoundaryConditions(un,vn)

    u = un
    v = vn

    t = t + dt

    print *, "t = ", t
  end do

  open(12,file='uvelocity.dat',status='replace')
    do i=1,nx
       write(12,'(1000F14.7)') (un(i,j),j=1,ny)
    end do
  close(12)

  open(10,file='pressure.dat',status='replace')
  do i=1,nx
     write(10,'(1000F14.7)') (P(i,j),j=1,ny)
  end do
  close(10)

  open(10,file='xcoord.dat',status='replace')
  do i=1,nx
     write(10,'(1000F14.7)') (x(i,j),j=1,ny)
  end do
  close(10)

  open(10,file='ycoord.dat',status='replace')
  do i=1,nx
     write(10,'(1000F14.7)') (y(i,j),j=1,ny)
  end do
  close(10)

contains

subroutine BoundaryConditions (u, v)
      implicit none
      double precision, dimension(nx,ny), intent(out) :: u, v
      integer :: i, j

      do i=1,nx
          u(i,1) = 0.0
          u(i,ny) = 0.0
          v(i,1) = 0.0
          v(i,ny) = 0.0
      end do

      do j=1,ny
          u(nx,j) = u(nx-1,j)
          v(1,j) = 0.0
          v(nx,j) = v(nx-1,j)
      end do

      do j=15,17
          u(1,j) = y(1,j)*(2.0 - y(1,j))
      end do
      return
end subroutine BoundaryConditions

!2nd-order BC as in Kim and Moin
subroutine IntermediateBC (us, vs, un, vn, P)
        implicit none
        integer :: i, j
        double precision, dimension(nx,ny), intent(out) :: un, vn, us, vs, P

        do j=1,ny
            us(1,j) = un(1,j) + (dt/dx)*(P(2,j) - P(1,j))
            us(nx,j) = un(nx,j) + (dt/dx)*(P(nx,j) - P(nx-1,j))

            vs(1,j) = vn(1,j) + (dt/dx)*(P(2,j) - P(1,j))
            vs(nx,j) = vn(nx,j) + (dt/dx)*(P(nx,j) - P(nx-1,j))
        end do

        do i=1,nx
            us(i,1) = un(i,1) + (dt/dy)*(P(i,2) - P(i,1))
            us(i,ny) = un(i,ny) + (dt/dy)*(P(i,ny) - P(i,ny-1))

            vs(i,1) = vn(i,1) + (dt/dy)*(P(i,2) - P(i,1))
            vs(i,ny) = vn(i,ny) + (dt/dy)*(P(i,ny) - P(i,ny-1))
        end do

        return
end subroutine IntermediateBC

subroutine CartesianGrid (x, y, dx, dy)
      implicit none
      integer :: i, j
      double precision, dimension(nx,ny), intent(out) :: x, y
      double precision, intent(in) :: dx, dy

      do j=1,ny
        do i=1,nx
          x(i,j) = (i-1.0)*dx
          y(i,j) = (j-1.0)*dy
        end do
      end do
end subroutine CartesianGrid

subroutine InitialCondition (u, v, P)
      implicit none
      double precision, dimension(nx,ny), intent(out) :: u, v, P

      do j=1,ny
        do i=1,nx
          u(i,j) = 0.0
          v(i,j) = 0.0
          P(i,j) = 0.0
        end do
      end do
end subroutine InitialCondition

! 2nd-order central differencing for diffusion and advection
subroutine Momentum (u, v, uMtm, vMtm)
    implicit none
    integer :: i, j
    double precision, dimension(nx,ny), intent(out) :: uMtm, vMtm
    double precision, dimension(nx,ny), intent(inout) :: u, v
    double precision, dimension(nx,ny) :: Cx, Cy;
    double precision, dimension(nx,ny) :: Vx, Vy
    double precision, dimension(nx,ny) :: uf, vf

    call FaceVelocities(u,v,uf,vf)
    call BoundaryConditions(uf,vf)

    do j=2,ny-1
        do i=2,nx-1
            Cx(i,j) = (1.0/dx)*(uf(i,j)*uf(i,j) - uf(i-1,j)*uf(i-1,j)) + (1.0/dy)*(uf(i,j)*vf(i,j) - uf(i,j-1)*vf(i,j-1))
            Cy(i,j) = (1.0/dx)*(uf(i,j)*vf(i,j) - uf(i-1,j)*vf(i-1,j)) + (1.0/dy)*(vf(i,j)*vf(i,j) - vf(i,j-1)*vf(i,j-1))

            Vx(i,j) = nu*((1.0/dx**2)*(u(i-1,j) - 2.0*u(i,j) + u(i+1,j)) + (1.0/dy**2)*(u(i,j-1) - 2.0*u(i,j) + u(i,j+1)))
            Vy(i,j) = nu*((1.0/dx**2)*(v(i-1,j) - 2.0*v(i,j) + v(i+1,j)) + (1.0/dy**2)*(v(i,j-1) - 2.0*v(i,j) + v(i,j+1)))

            uMtm(i,j) = -Cx(i,j) + Vx(i,j)
            vMtm(i,j) = -Cy(i,j) + Vy(i,j)
        end do
    end do
    return
end subroutine

subroutine PressureSolver (P, us, vs, u, v)
    implicit none
    integer :: i, j
    double precision, dimension(nx,ny), intent(out) :: P
    double precision, dimension(nx,ny), intent(in) :: us, vs
    double precision, dimension(nx,ny), intent(in) :: u, v
    double precision, dimension(nx,ny) :: vfs, ufs
    double precision, dimension(nx,ny) :: R, Res
    double precision, parameter :: inv_dxx = 1.0/dx**2
    double precision, parameter :: inv_dyy  = 1.0/dy**2
    double precision, parameter :: inv_dx = 1.0/dx
    double precision, parameter :: inv_dy = 1.0/dy
    double precision, parameter :: tol = 1e-6
    double precision, parameter :: iter_max = 16000
    double precision, parameter :: omega = 1.96
    double precision :: L2_Error
    double precision :: iter

    call FaceVelocities(us,vs,ufs,vfs)
    call BoundaryConditions(ufs,vfs)

    do j=3,ny-2
        do i=3,nx-2
            R(i,j) = (rho/dt)*(inv_dx*(ufs(i,j) - ufs(i-1,j)) + inv_dy*(vfs(i,j) - vfs(i,j-1)))
        end do
    end do

    Iter = 0

    do while (iter <= iter_max)
        do j=3,ny-2
            do i=3,nx-2
                P(i,j) = (1.0-omega)*P(i,j) + omega/(2.0*inv_dxx + 2.0*inv_dyy)*(inv_dxx*(P(i+1,j) + P(i-1,j)) + inv_dyy*(P(i,j+1) + P(i,j-1))- R(i,j))
            end do
        end do

        do j=3,ny-2
            R(2,j) = (rho/dt)*(inv_dx*(ufs(2,j)- u(1,j)) + inv_dy*(vfs(2,j) - vfs(2,j-1)))
            P(2,j) = (1.0-omega)*P(2,j) + omega/(inv_dxx + 2.0*inv_dyy)*(inv_dxx*P(3,j) + inv_dyy*(P(2,j+1) + P(2,j-1))-R(2,j))

            R(nx-1,j) = (rho/dt)*(inv_dx*(u(nx,j) - ufs(nx-2,j)) + inv_dy*(vfs(nx-1,j) - vfs(nx-1,j-1)))
            P(nx-1,j) = (1.0-omega)*P(nx-1,j) + omega/(inv_dxx + 2.0*inv_dyy)*(inv_dxx*P(nx-2,j) + inv_dyy*(P(nx-1,j+1) + P(nx-1,j-1)) - R(nx-1,j))
        end do

        do i=3,nx-2
            R(i,2) = (rho/dt)*(inv_dx*(ufs(i,2) - ufs(i-1,2)) + inv_dy*(vfs(i,2)- v(i,1)))
            P(i,2) = (1.0-omega)*P(i,2) + omega/(2.0*inv_dxx + inv_dyy)*(inv_dxx*(P(i+1,2) + P(i-1,2)) + inv_dyy*P(i,3) - R(i,2))

            R(i,ny-1) = (rho/dt)*(inv_dx*(ufs(i,ny-1)- ufs(i-1,ny-1)) + inv_dy*(v(i,ny) - vfs(i,ny-2)))
            P(i,ny-1) = (1.0-omega)*P(i,ny-1) + omega/(2.0*inv_dxx + inv_dyy)*(inv_dxx*(P(i+1,ny-1) + P(i-1,ny-1)) + inv_dyy*P(i,ny-2) - R(i,ny-1))
        end do

        R(2,ny-1) = (rho/dt)*(inv_dx*(ufs(2,ny-1) - u(1,ny-1)) + inv_dy*(v(2,ny) - vfs(2,ny-2)))
        P(2,ny-1) = (1.0-omega)*P(2,ny-1) + omega/(inv_dxx + inv_dyy)*(inv_dxx*P(3,ny-1) + inv_dyy*P(2,ny-2) - R(2,ny-1))

        R(nx-1,2) = (rho/dt)*(inv_dx*(u(nx,2) - ufs(nx-2,2)) + inv_dy*(vfs(nx-1,2) - v(nx-1,1)))
        P(nx-1,2) = (1.0-omega)*P(nx-1,2) + omega/(inv_dxx + inv_dyy)*(inv_dxx*P(nx-2,2) + inv_dyy*P(nx-1,3) - R(nx-1,2))

        R(nx-1,ny-1) = (rho/dt)*(inv_dx*(u(nx,ny-1) - ufs(nx-2,ny-1)) + inv_dy*(v(nx-1,ny) - vfs(nx-1,ny-2)))
        P(nx-1,ny-1) = (1.0-omega)*P(nx-1,ny-1) + omega/(inv_dxx + inv_dyy)*(inv_dxx*P(nx-2,ny-1) + inv_dyy*P(nx-1,ny-2) - R(nx-1,ny-1))

        !R(2,2) = (rho/dt)*(inv_dx*(ufs(2,2) - u(1,2)) + inv_dy*(vfs(2,2) - v(2,1)))
        !P(2,2) = (1.0-omega)*P(2,2) + omega/(inv_dxx + inv_dyy)*(inv_dxx*P(3,2) + inv_dyy*P(2,3) - R(2,2))
        P(2,2) = 1.0
        ! If I remove the P(2,2) BC and set P(2,2) = 1.0 then my simulation does not crash, however my soltuion does not satisfy the compatibility condition
        do j=3,ny-2
            do i=3,nx-2
                    Res(i,j) = (inv_dxx*(P(i-1,j) - 2*P(i,j) + P(i+1,j)) + inv_dyy*(P(i,j-1) - 2*P(i,j) + P(i,j+1))) - R(i,j);
            end do
        end do

        L2_Error = sum(dx*dy*abs(Res))

        if (L2_Error >= tol) then
          Iter = Iter + 1;
        else
            exit
        end if
    end do

    P(1,:) = P(2,:) + (dx/dt)*(us(1,:) - u(1,:))
    P(nx,:) = P(nx-1,:) - (dx/dt)*(us(nx,:) - u(nx,:))
    P(:,1) = P(:,2) + (dy/dt)*(vs(:,1) - v(:,1))
    P(:,ny) = P(:,ny-1) - (dy/dt)*(vs(:,ny) - v(:,ny))
end subroutine PressureSolver

subroutine FaceVelocities (u, v, uf, vf)
        implicit none
        double precision, dimension(nx,ny), intent(in) :: u, v
        double precision, dimension(nx,ny), intent(out) :: uf, vf
        integer :: i, j

        do j=1,ny
            do i=1,nx-1
                uf(i,j) = 0.5*(u(i+1,j) + u(i,j))
            end do
        end do

        do j=1,ny-1
            do i=1,nx
                vf(i,j) = 0.5*(v(i,j+1) + v(i,j))
            end do
        end do
end subroutine FaceVelocities

end program