!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
!
! RayleighBernard2D 
! press.f90 module
! James Sandham 
! 11 Mar 2014
!
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!


module press
use def
use fft
use solvers
use mpi

contains
  subroutine pressure(uh,vh,ph,ik)
  complex, dimension(bs), intent(in) :: uh,vh,ik
  complex, dimension(bs), intent(inout) :: ph

  integer :: status(MPI_STATUS_SIZE)
  integer :: i,j,k,tag=100
  complex, dimension(bs) :: div=0.0
  complex, dimension(n) :: lower_bc=0.0,upper_bc=0.0
  real, dimension(n) :: a,b,c

  if (mod(rank,2)==0) then
    if (rank==0) then
      call MPI_SEND(vh((bs-n)+1:bs),n,MPI_COMPLEX,rank+1,tag,MPI_COMM_WORLD,ierr)
      call MPI_RECV(upper_bc,n,MPI_COMPLEX,rank+1,tag,MPI_COMM_WORLD,status,ierr)
    else
      call MPI_SEND(vh((bs-n)+1:bs),n,MPI_COMPLEX,rank+1,tag,MPI_COMM_WORLD,ierr)
      call MPI_SEND(vh(1:n),n,MPI_COMPLEX,rank-1,tag,MPI_COMM_WORLD,ierr)
      call MPI_RECV(lower_bc,n,MPI_COMPLEX,rank-1,tag,MPI_COMM_WORLD,status,ierr)
      call MPI_RECV(upper_bc,n,MPI_COMPLEX,rank+1,tag,MPI_COMM_WORLD,status,ierr)
    endif
  else
    if (rank==nproc-1) then
      call MPI_RECV(lower_bc,n,MPI_COMPLEX,rank-1,tag,MPI_COMM_WORLD,status,ierr)
      call MPI_SEND(vh(1:n),n,MPI_COMPLEX,rank-1,tag,MPI_COMM_WORLD,ierr)
    else
      call MPI_RECV(lower_bc,n,MPI_COMPLEX,rank-1,tag,MPI_COMM_WORLD,status,ierr)
      call MPI_RECV(upper_bc,n,MPI_COMPLEX,rank+1,tag,MPI_COMM_WORLD,status,ierr)
      call MPI_SEND(vh((bs-n)+1:bs),n,MPI_COMPLEX,rank+1,tag,MPI_COMM_WORLD,ierr)
      call MPI_SEND(vh(1:n),n,MPI_COMPLEX,rank-1,tag,MPI_COMM_WORLD,ierr)
    endif
  endif

  do i=1,bs
    j=i-n ; k=i+n
    if (j<=0) then
      if (rank==0) then
        div(i) = ik(i)*uh(i) + (-vh(i+2*n)+4*vh(k)-3*vh(i))/(2*dy)
        !div(i) = ik(i)*uh(i) + (vh(k)-vh(i))/dy
      else
        div(i) = ik(i)*uh(i)+(vh(k)-lower_bc(i))/(2*dy)
      endif
    elseif (k>bs) then
      if (rank==nproc-1) then
        div(i) = ik(i)*uh(i) + (3*vh(i)-4*vh(j)+vh(i-2*n))/(2*dy)
        !div(i) = ik(i)*uh(i) + (vh(i)-vh(j))/dy
      else
        div(i) = ik(i)*uh(i)+(upper_bc(k-bs)-vh(j))/(2*dy)
      endif
    else
      div(i) = ik(i)*uh(i)+(vh(k)-vh(j))/(2*dy)
    endif
  enddo

  div = (1.0/dt)*div

  a=1.0/dy2 ; c=1.0/dy2
  a(1)=0.0; c(n)=0.0

  call mpitranspose(div)

  do i=1,bs/n
    j = rank*bs/n+i
    if (j<=n/2) then
       k=j-1
    else
       k=-(n-j+1)
    end if
    b=(-2.0/dy2-k**2)
    b(n)=(-1.0/dy2-k**2);
    if (k==0) then
      b(1)=-2.0/dy2
    else
      b(1)=-1.0/dy2-k**2
    endif
    
    call tridiag(div(n*(i-1)+1:n*i),ph(n*(i-1)+1:n*i),a,b,c)
  enddo

  call mpitranspose(ph)  

  end subroutine pressure



  subroutine update_pressure(uh,vh,ph,ik)
  complex, dimension(bs), intent(inout) :: uh, vh
  complex, dimension(bs), intent(in) :: ph,ik

  integer :: status(MPI_STATUS_SIZE)
  integer :: i,j,k,tag=100
  complex, dimension(bs) :: dpdy=0.0
  complex, dimension(n) :: lower_bc=0.0,upper_bc=0.0
  
  if (mod(rank,2)==0) then
    if (rank==0) then
      call MPI_SEND(ph((bs-n)+1:bs),n,MPI_COMPLEX,rank+1,tag,MPI_COMM_WORLD,ierr)
      call MPI_RECV(upper_bc,n,MPI_COMPLEX,rank+1,tag,MPI_COMM_WORLD,status,ierr)
    else
      call MPI_SEND(ph((bs-n)+1:bs),n,MPI_COMPLEX,rank+1,tag,MPI_COMM_WORLD,ierr)
      call MPI_SEND(ph(1:n),n,MPI_COMPLEX,rank-1,tag,MPI_COMM_WORLD,ierr)
      call MPI_RECV(lower_bc,n,MPI_COMPLEX,rank-1,tag,MPI_COMM_WORLD,status,ierr)
      call MPI_RECV(upper_bc,n,MPI_COMPLEX,rank+1,tag,MPI_COMM_WORLD,status,ierr)
    endif
  else
    if (rank==nproc-1) then
      call MPI_RECV(lower_bc,n,MPI_COMPLEX,rank-1,tag,MPI_COMM_WORLD,status,ierr)
      call MPI_SEND(ph(1:n),n,MPI_COMPLEX,rank-1,tag,MPI_COMM_WORLD,ierr)
    else
      call MPI_RECV(lower_bc,n,MPI_COMPLEX,rank-1,tag,MPI_COMM_WORLD,status,ierr)
      call MPI_RECV(upper_bc,n,MPI_COMPLEX,rank+1,tag,MPI_COMM_WORLD,status,ierr)
      call MPI_SEND(ph((bs-n)+1:bs),n,MPI_COMPLEX,rank+1,tag,MPI_COMM_WORLD,ierr)
      call MPI_SEND(ph(1:n),n,MPI_COMPLEX,rank-1,tag,MPI_COMM_WORLD,ierr)
    endif
  endif

  do i=1,bs
    j=i-n ; k=i+n
    if (j<=0) then
      if (rank==0) then
        dpdy(i) = 0.0
      else
        dpdy(i) = (ph(k)-lower_bc(i))/(2*dy)
      endif
    elseif (k>bs) then
      if (rank==nproc-1) then
        dpdy(i) = 0.0
      else
        dpdy(i) = (upper_bc(k-bs)-ph(j))/(2*dy)
      endif
    else
      dpdy(i) = (ph(k)-ph(j))/(2*dy)
    endif
  enddo

  !update u-velocity with x-pressure gradient
  if (rank==0) then
    do i=n+1,bs
      uh(i) = uh(i) - dt*ik(i)*ph(i)
    enddo
  elseif (rank==nproc-1) then
    do i=1,bs-n
      uh(i) = uh(i) - dt*ik(i)*ph(i)
    enddo
  else
    do i=1,bs
      uh(i) = uh(i) - dt*ik(i)*ph(i)
    enddo
  endif

  !update v-velocity with y-pressure gradient
  vh = vh - dt*dpdy
  end subroutine update_pressure

end module press
