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


module scalar
use def
use fft
use solvers
use mpi

contains
  subroutine density(uh,vh,rhoh,sc,sp,ik)
  complex, dimension(bs), intent(inout) :: rhoh
  complex, dimension(bs), intent(inout) :: sc,sp
  complex, dimension(bs), intent(in) :: uh,vh,ik

  integer :: status(MPI_STATUS_SIZE)
  integer :: i,j,k,ii,tag=100
  complex, dimension(bs) :: urhoh=0.0,vrhoh=0.0
  complex, dimension(bs) :: urhor=0.0,vrhor=0.0,ur=0.0,vr=0.0,rhor=0.0
  complex, dimension(n) :: lower_bc=0.0,upper_bc=0.0
  complex, dimension(n) :: rhoht
  real, dimension(n) :: a,b,c
  
  call ifft1(uh,ur)
  call ifft1(vh,vr)
  call ifft1(rhoh,rhor)

  do i=1,bs
    urhor(i) = ur(i)*rhor(i)
    vrhor(i) = vr(i)*rhor(i)
  enddo
  
  call fft1(urhor,urhoh)
  call fft1(vrhor,vrhoh)
  
  if (mod(rank,2)==0) then
    if (rank==0) then
      call MPI_SEND(vrhoh((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(vrhoh((bs-n)+1:bs),n,MPI_COMPLEX,rank+1,tag,MPI_COMM_WORLD,ierr)
      call MPI_SEND(vrhoh(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(vrhoh(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(vrhoh((bs-n)+1:bs),n,MPI_COMPLEX,rank+1,tag,MPI_COMM_WORLD,ierr)
      call MPI_SEND(vrhoh(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
        sc(i) = 0.0 
      else
        sc(i) = -ik(i)*urhoh(i)-(vrhoh(k)-lower_bc(i))/(2*dy)
      endif
    elseif (k>bs) then
      if (rank==nproc-1) then
        sc(i) = 0.0 
      else
        sc(i) = -ik(i)*urhoh(i)-(upper_bc(k-bs)-vrhoh(j))/(2*dy)
      endif
    else
      sc(i) = -ik(i)*urhoh(i)-(vrhoh(k)-vrhoh(j))/(2*dy)
    endif
  enddo


  if (tt==1) then
    rhoh = rhoh + dt*sc
  else
    rhoh = rhoh + (dtl*sc - dts*sp)
  endif
  sp = sc;


  call mpitranspose(rhoh)

  a=-dt*kappa/dy2 ; c=-dt*kappa/dy2
  a(1)=0.0; a(n)=0.0; c(1)=0.0; c(n)=0.0

  do i=1,bs/n
    rhoht = rhoh(n*(i-1)+1:n*i)

    j = rank*bs/n+i
    if (j<=n/2) then
      k=j-1
    else
      k=-(n-j+1)
    end if

    do ii=2,n-1
      rhoht(ii) = rhoht(ii)-0.5*kappa*dt*(k**2*rhoht(ii)-(rhoht(ii+1)-2*rhoht(ii)+rhoht(ii-1))/dy2)
    enddo

    b=1.0+dt*kappa*(2.0/dy2+k**2); b(1)=1.0; b(n)=1.0  
    call tridiag(rhoht,rhoh(n*(i-1)+1:n*i),a,b,c)
  enddo

  call mpitranspose(rhoh)
  end subroutine density


  subroutine update_density(vh,rhoh)
  complex, dimension(bs), intent(inout) :: vh
  complex ,dimension(bs), intent(in) :: rhoh

  integer :: i

  if (rank==0) then
    do i=n+1,bs
      vh(i) = vh(i) - dt*g*rhoh(i)
    enddo
  elseif (rank==nproc-1) then
    do i=1,bs-n
      vh(i) = vh(i) - dt*g*rhoh(i)
    enddo
  else
    do i=1,bs
      vh(i) = vh(i) - dt*g*rhoh(i)
    enddo
  endif
  end subroutine update_density

end module scalar
