#include "options.inc"

#define remove_coriolis_expl_free_surf

      subroutine solve_pressure
c=======================================================================
c      solve for surface pressure or implicit linear free surface
c=======================================================================
      use cpflame_module
      use congrad2D_module
      use congrad3D_module
      implicit none
      integer :: i,j,k,n,js,je,js2,je2, maxitt=5000
      real :: forc2D(imt,jmt), forc3D(imt,jmt,km)
      real :: fpx(imt,jmt), fpy(imt,jmt)

      real :: etam(imt,jmt),bum(imt,jmt),bvm(imt,jmt)

      logical, save :: first = .true.

      js =max(2,js_pe);   je  = min(je_pe,jmt-1)
      js2=max(1,js_pe-1); je2 = min(je_pe+1,jmt)
c---------------------------------------------------------------------------------
c      hydrostatic pressure
c---------------------------------------------------------------------------------
      do j=js2,je2
       p_hydro(:,j,km) = 0
       do k=km-1,1,-1
        p_hydro(:,j,k)= maskT(:,j,k)*
     &    (p_hydro(:,j,k+1)+(b(:,j,k+1,tau)+b(:,j,k,tau))/2.0*dz)
       enddo
      enddo
      call setcyclic3D(p_hydro)
c---------------------------------------------------------------------------------
c       surface pressure forcing, etc
c---------------------------------------------------------------------------------
      fpx(:,js_pe:je_pe)=0.;fpy(:,js_pe:je_pe)=0.
      do k=1,km-1
       do j=js,je
        do i=2,imt-1
           fpx(i,j)=fpx(i,j)
     &         +(fu(i,j,k)-(p_hydro(i+1,j,k)-p_hydro(i,j,k))/dx
     &                )*maskU(i,j,k)*dz
           fpy(i,j)=fpy(i,j)
     &         +(fv(i,j,k)-(p_hydro(i,j+1,k)-p_hydro(i,j,k))/dx
     &                )*maskV(i,j,k)*dz
        enddo
       enddo
      enddo

      if (enable_free_surface) then
      do k=1,km-1
       do j=js,je
        do i=2,imt-1
           fpx(i,j)=fpx(i,j)+(u(i,j,k,1,taum1)/c2dt
     &                )*maskU(i,j,k)*dz
           fpy(i,j)=fpy(i,j)+(u(i,j,k,2,taum1)/c2dt
     &                )*maskV(i,j,k)*dz
        enddo
       enddo
      enddo
      endif

      if (.not.enable_hydrostatic) then
       do k=1,km-1
        do j=js,je
         do i=2,imt-1
           fpx(i,j)=fpx(i,j)
     &         -(psi(i+1,j,k)-psi(i,j,k))/dx*maskU(i,j,k)*dz
           fpy(i,j)=fpy(i,j)
     &         -(psi(i,j+1,k)-psi(i,j,k))/dx*maskV(i,j,k)*dz
         enddo
        enddo
       enddo
      endif

#ifdef remove_coriolis_expl_free_surf
      if (enable_expl_free_surf) then
       do k=1,km-1
        do j=js,je
         do i=2,imt-1
           fpx(i,j)=fpx(i,j)-
     &    maskU(i,j,k)*dz*
     &     coriolis_t(j)*(u(i,j  ,k,2,tau)+u(i+1,j  ,k,2,tau)+
     &                    u(i,j-1,k,2,tau)+u(i+1,j-1,k,2,tau))/4.0
           fpy(i,j)=fpy(i,j)
     &              +maskV(i,j,k)*dz*
     &     (coriolis_t(j  )*(u(i-1,j  ,k,1,tau)+u(i,j  ,k,1,tau)) 
     &     +coriolis_t(j+1)*(u(i-1,j+1,k,1,tau)+u(i,j+1,k,1,tau)))/4.0
         enddo
        enddo
       enddo
      endif
#endif

      call border_exchg2D(fpx,1); call border_exchg2D(fpy,1)
      call setcyclic2D(fpx); call setcyclic2D(fpy)

      if (enable_expl_free_surf) then
        n=nint(dt_in/dtex)
        do j=js2,je2
         etam(:,j)=0;
         bum(:,j)=0;
         bvm(:,j)=0;
        enddo
        do k=1,n*2
c        (eta'-eta)/dt + U_x + V_y = 0
c        (U'-U)/dt + g H eta_x = fu
c        (V'-V)/dt + g H eta_y = fv
         do j=js,je
          do i=2,imt-1
           eta(i,j,taum1)=eta(i,j,tau)+maskT(i,j,km-1)*dtex*(
     &         -(bu(i,j,tau)-bu(i-1,j,tau))/dx
     &         -(bv(i,j,tau)-bv(i,j-1,tau))/dx )
           bu(i,j,taum1)=bu(i,j,tau)+maskU(i,j,km-1)*dtex*(
     &          - g*hu(i,j)*(eta(i+1,j,tau)-eta(i,j,tau))/dx
#ifdef remove_coriolis_expl_free_surf
     &    +coriolis_t(j)*(bv(i,j,tau)+bv(i+1,j,tau)+
     &                           bv(i,j-1,tau)+bv(i+1,j-1,tau))/4.0
#endif
     &          + fpx(i,j) )
           bv(i,j,taum1)=bv(i,j,tau)+maskV(i,j,km-1)*dtex*(
     &          - g*hv(i,j)*(eta(i,j+1,tau)-eta(i,j,tau))/dx
#ifdef remove_coriolis_expl_free_surf
     &     -(coriolis_t(j  )*(bu(i-1,j,tau)+bu(i,j,tau)) 
     &      +coriolis_t(j+1)*(bu(i-1,j+1,tau)+bu(i,j+1,tau)))/4.0
#endif
     &          + fpy(i,j) )
          enddo
         enddo
         call border_exchg2D(eta(:,:,taum1),1); 
         call setcyclic2D(eta(:,:,taum1))
         call border_exchg2D(bu(:,:,taum1),1); 
         call setcyclic2D(bu(:,:,taum1))
         call border_exchg2D(bv(:,:,taum1),1); 
         call setcyclic2D(bv(:,:,taum1))
c        (eta''-eta)/dt+U'_x+V'_y = 0
c        (U''-U)/dt + g H eta'_x = fu
c        (V''-V)/dt + g H eta'_y = fv
         do j=js,je
          do i=2,imt-1
           eta(i,j,taup1)=eta(i,j,tau)+maskT(i,j,km-1)*dtex*(
     &         -(bu(i,j,taum1)-bu(i-1,j,taum1))/dx
     &         -(bv(i,j,taum1)-bv(i,j-1,taum1))/dx )
           bu(i,j,taup1)=bu(i,j,tau)+maskU(i,j,km-1)*dtex*(
     &          - g*hu(i,j)*(eta(i+1,j,taum1)-eta(i,j,taum1))/dx
#ifdef remove_coriolis_expl_free_surf
     &    +coriolis_t(j)*(bv(i,j,taum1)+bv(i+1,j,taum1)+
     &                           bv(i,j-1,taum1)+bv(i+1,j-1,taum1))/4.0
#endif
     &          + fpx(i,j) )
           bv(i,j,taup1)=bv(i,j,tau)+maskV(i,j,km-1)*dtex*(
     &          - g*hv(i,j)*(eta(i,j+1,taum1)-eta(i,j,taum1))/dx
#ifdef remove_coriolis_expl_free_surf
     &     -(coriolis_t(j  )*(bu(i-1,j,taum1)+bu(i,j,taum1)) 
     &      +coriolis_t(j+1)*(bu(i-1,j+1,taum1)+bu(i,j+1,taum1)))/4.0
#endif
     &          + fpy(i,j) )
          enddo
         enddo
         call border_exchg2D(eta(:,:,taup1),1); 
         call setcyclic2D(eta(:,:,taup1))
         call border_exchg2D(bu(:,:,taup1),1); 
         call setcyclic2D(bu(:,:,taup1))
         call border_exchg2D(bv(:,:,taup1),1); 
         call setcyclic2D(bv(:,:,taup1))
c       shift time positions
         do j=js2,je2
          eta(:,j,tau)=eta(:,j,taup1)
          bu(:,j,tau)=bu(:,j,taup1)
          bv(:,j,tau)=bv(:,j,taup1)
         enddo
c        accumulate eta,u and v
         do j=js2,je2
          etam(:,j)=etam(:,j)+eta(:,j,taup1)
          bum(:,j)=bum(:,j)+bu(:,j,taup1)
          bvm(:,j)=bvm(:,j)+bv(:,j,taup1)
         enddo
        enddo! n
c       use averaged eta, u and v
        do j=js2,je2
          eta(:,j,tau)=etam(:,j)/(2.*n); eta(:,j,taup1)=eta(:,j,tau)
          bu(:,j,tau)=bum(:,j)/(2.*n); bu(:,j,taup1)=bu(:,j,tau)
          bv(:,j,tau)=bvm(:,j)/(2.*n); bv(:,j,taup1)=bv(:,j,tau)
        enddo
        do j=js2,je2
         p_surf(:,j)=eta(:,j,taup1)*g
        enddo
      else ! enable_expl_free_surf
       forc2D(:,js_pe:je_pe)=0.0
       do j=js,je
        do i=2,imt-1
         forc2D(i,j)=(fpx(i,j)-fpx(i-1,j))/dx+(fpy(i,j)-fpy(i,j-1))/dx 
        enddo
       enddo
       if (enable_free_surface) then
        do j=js,je
         do i=2,imt-1
          forc2D(i,j)=forc2D(i,j)-eta(i,j,taum1)/(c2dt**2)
         enddo
        enddo
       endif
       call setcyclic2D(forc2D)
c---------------------------------------------------------------------------------
c      solve for surface pressure or free surface
c---------------------------------------------------------------------------------
       if (first) call make_coef2D(cf2D)
       if (enable_free_surface) then
        p_surf(:,js_pe:je_pe)=eta(:,js_pe:je_pe,taup1)
        call border_exchg2D(p_surf,1); call setcyclic2D(p_surf)
       endif
       call congrad2D(cf2D,forc2D,maxitt,sor2D_itts,eps2D_sor)
c      call border_exchg2D(p_surf,1); call setcyclic2D(p_surf)
       if (enable_free_surface) then
        do j=js,je
         do i=2,imt-1
          eta(i,j,taup1)=p_surf(i,j)
          p_surf(i,j)=eta(i,j,taup1)*g
          eta(i,j,tau) = eta(i,j,tau) + gamma*
     &      (0.5*(eta(i,j,taup1) + eta(i,j,taum1)) - eta(i,j,tau)) 
         enddo
        enddo
       endif
       call border_exchg2D(p_surf,1); call setcyclic2D(p_surf)
      endif

c---------------------------------------------------------------------------------
c       full pressure is surface pressure plus hydrostatic pressure
c---------------------------------------------------------------------------------
      do k=1,km
       do j=js,je
        p_full(:,j,k,tau) =(p_hydro(:,j,k)+p_surf(:,j))*maskT(:,j,k)
       enddo
      enddo
      call border_exchg3D(p_full(:,:,:,tau),1)
      call setcyclic3D(p_full(:,:,:,tau))

      if (.not. enable_hydrostatic) then
c---------------------------------------------------------------------------------
c        forcing for non-hydrostatic pressure
c---------------------------------------------------------------------------------
       forc3D=0.
       do k=2,km-1
        do j=js,je
         do i=2,imt-1
          forc3D(i,j,k)= (fw(i,j,k)-fw(i,j,k-1))/dz
     &                 + (fu(i,j,k)-fu(i-1,j,k))/dx
     &                 + (fv(i,j,k)-fv(i,j-1,k))/dx
     &  -((p_full(i+1,j,k,tau)-p_full(i,j,k,tau))/dx*maskU(i  ,j,k)
     &   -(p_full(i,j,k,tau)-p_full(i-1,j,k,tau))/dx*maskU(i-1,j,k))/dx
     &  -((p_full(i,j+1,k,tau)-p_full(i,j,k,tau))/dx*maskV(i,j  ,k)
     &   -(p_full(i,j,k,tau)-p_full(i,j-1,k,tau))/dx*maskV(i,j-1,k))/dx
         enddo
        enddo
       enddo
       call setcyclic3D(forc3D)
c---------------------------------------------------------------------------------
c        solve for non hydrostatic pressure
c---------------------------------------------------------------------------------
       if (first) call make_coef3D(cf3D)
       call congrad3D(cf3D,forc3D,maxitt,sor3D_itts,eps3D_sor)
       call border_exchg3D(psi,1)
       call setcyclic3D(psi)
c---------------------------------------------------------------------------------
c       full pressure 
c---------------------------------------------------------------------------------
       do j=js,je
        p_full(:,j,:,tau) =(p_full(:,j,:,tau)+psi(:,j,:))*maskT(:,j,:)
       enddo
       call border_exchg3D(p_full(:,:,:,tau),1)
       call setcyclic3D(p_full(:,:,:,tau))
      endif

      first=.false.
      end subroutine solve_pressure



      subroutine make_coef3D(cf)
      use cpflame_module
      implicit none
c-----------------------------------------------------------------------
c             A * dpsi = forc
c                       res = A * p
c          res = res +  cf(...,ii,jj,kk) * p(i+ii,j+jj,k+kk) 
c
c          forc = p_xx + p_yy + p_zz
c         forc = (p(i+1) - 2p(i) + p(i-1))  /dx^2 ...
c              = [ (p(i+1) - p(i))/dx - (p(i)-p(i-1))/dx ] /dx 
c-----------------------------------------------------------------------
      real :: cf(imt,jmt,km,-1:1,-1:1,-1:1)
      real :: maskM(imt,jmt,km),mp,mm
      integer :: i,j,k

      cf=0.
      do k=2,km-1
       do j=2,jmt-1
        do i=2,imt-1
         mp=maskT(i,j,k)*maskT(i+1,j,k)/dx**2
         mm=maskT(i,j,k)*maskT(i-1,j,k)/dx**2
         cf(i,j,k, 0, 0, 0)= cf(i,j,k, 0, 0, 0)-mp
         cf(i,j,k, 1, 0, 0)= cf(i,j,k, 1, 0, 0)+mp
         cf(i,j,k, 0, 0, 0)= cf(i,j,k, 0, 0, 0)-mm
         cf(i,j,k,-1, 0, 0)= cf(i,j,k,-1, 0, 0)+mm

         mp=maskT(i,j,k)*maskT(i,j+1,k)/dx**2
         mm=maskT(i,j,k)*maskT(i,j-1,k)/dx**2
         cf(i,j,k, 0, 0, 0)= cf(i,j,k, 0, 0, 0)-mp
         cf(i,j,k, 0, 1, 0)= cf(i,j,k, 0, 1, 0)+mp
         cf(i,j,k, 0, 0, 0)= cf(i,j,k, 0, 0, 0)-mm
         cf(i,j,k, 0,-1, 0)= cf(i,j,k, 0,-1, 0)+mm

         mp=maskT(i,j,k)*maskT(i,j,k+1)/dz**2
         mm=maskT(i,j,k)*maskT(i,j,k-1)/dz**2
         cf(i,j,k, 0, 0, 0)= cf(i,j,k, 0, 0, 0)-mp
         cf(i,j,k, 0, 0, 1)= cf(i,j,k, 0, 0, 1)+mp
         cf(i,j,k, 0, 0, 0)= cf(i,j,k, 0, 0, 0)-mm
         cf(i,j,k, 0, 0,-1)= cf(i,j,k, 0, 0,-1)+mm

        end do
       end do
      end do
      end subroutine make_coef3D


      subroutine make_coef2D(cf)
      use cpflame_module
      implicit none
c-----------------------------------------------------------------------
c         A * dpsi = forc
c         res = A * p
c         res = res +  cf(...,ii,jj,kk) * p(i+ii,j+jj,k+kk) 
c
c         forc = (h p_x)_x +(h p_y)_y + (hp_z)_z
c         forc = [ hu(i)(p(i+1) - p(i))/dx - hu(i-1)(p(i)-p(i-1))/dx ] /dx 
c         forc = hu(i) p(i+1)/dx^2  - p(i) (hu(i)+hu(i-1))/dx^2  + hu(i-1) p(i-1)/dx^2
c-----------------------------------------------------------------------
      real :: cf(imt,jmt,-1:1,-1:1)
      real :: maskM(imt,jmt),mp,mm
      integer :: i,j

      cf=0.
      maskM=maskT(:,:,km-1)
      do j=2,jmt-1
       do i=2,imt-1
         mp=maskM(i,j)*maskM(i+1,j)/dx**2
         mm=maskM(i,j)*maskM(i-1,j)/dx**2
         cf(i,j, 0, 0)= cf(i,j, 0, 0)-mp*hu(i  ,j)
         cf(i,j, 1, 0)= cf(i,j, 1, 0)+mp*hu(i  ,j)
         cf(i,j, 0, 0)= cf(i,j, 0, 0)-mm*hu(i-1,j)
         cf(i,j,-1, 0)= cf(i,j,-1, 0)+mm*hu(i-1,j)

         mp=maskM(i,j)*maskM(i,j+1)/dx**2
         mm=maskM(i,j)*maskM(i,j-1)/dx**2
         cf(i,j, 0, 0)= cf(i,j, 0, 0)-mp*hv(i,j)
         cf(i,j, 0, 1)= cf(i,j, 0, 1)+mp*hv(i,j)
         cf(i,j, 0, 0)= cf(i,j, 0, 0)-mm*hv(i,j-1)
         cf(i,j, 0,-1)= cf(i,j, 0,-1)+mm*hv(i,j-1)
       end do
      end do
      if (enable_free_surface) then
       cf=cf*g
       do j=2,jmt-1
        do i=2,imt-1
         cf(i,j, 0, 0)= cf(i,j,0,0)-1/( c2dt )**2
        end do
       end do
      endif
      end subroutine make_coef2D


