#include "options.inc"


c=======================================================================
c  add effect of a prescribed background stratification in buoyancy
c  to buoyancy budget. Also the effect of background zonal and meridional
c  velocity (but not both).
c  linked in the code in driver.F 
c
c    the case 
c    b=b'(x,y,z,t) + B(z)
c    p=p'(x,y,t,z)+P(z) with P_z = -B,  P(z=0)=0
c    u=u'(x,y,z,t)+U(z)
c
c    b'_t = -(u'b')_x - (vb')_y - (wb')_z + diffusion(b')+convection(b+B)
c             -(wB)_z-(u'B)_x-(vB)_y-(Ub')_x
c    
c    u'_t =-p_x -(u'u')_x-(vu')_y-(wu')_z+friction(u')  -(Uu')_x  -(u'U)_x-(vU)_y-(wU)_z
c    v'_t =-p_y -(u'v)_x-(vv)_y-(wv)_z+friction(v)  -(Uv)_x
c    w'_t =-p_z -(u'w)_x-(vw)_y-(ww)_z+friction(w)  -(Uw)_x - b 
c
c   with Coriolis force: specify U(z) and B=f y U_z with U(0)=0
c   thus best to set U=U_0 z/h  and B=f y U_0/h 
c
c    P_y = -fU(z)   ->  P(y,z)= - f y U(z) + f(z),  note that U(0)=0, such that P(0)=0
c    P_z = -B, ->  B= f y U_z + f'Q
c
c   numerically consistent:  B_{j+1}=B_{j} + dx*U_z (coriolis_t(j  )+coriolis_t(j+1))/2.0
c
c    U(z), P(y,z), B(y,z)
c
c    b'_t = -(u'b')_x - (vb')_y - (wb')_z + diffusion(b')+convection(b+B)
c             -(wB)_z-(u'B)_x-(vB)_y-(Ub')_x
c    u'_t =fv  -p_x -(u'u')_x-(vu')_y-(wu')_z  -(Uu')_x  -(u'U)_x-(vU)_y-(wU)_z
c    v'_t =-fu' -p_y -(u'v)_x-(vv)_y-(wv)_z  -(Uv)_x 
c
c    non-hydrostatic and U(z), P(y,z) and B(y,z):
c    b'_t = -(u'b')_x - (vb')_y - (wb')_z + diffusion(b')+convection(b+B)
c             -(wB)_z-(u'B)_x-(vB)_y-(Ub')_x
c    u'_t =fv -f_h w -p_x -(u'u')_x-(vu')_y-(wu')_z  -(Uu')_x -(u'U)_x-(vU)_y-(wU)_z
c    v'_t =-fu'      -p_y -(u'v)_x-(vv)_y-(wv)_z     -(Uv)_x 
c    w'_t =f_h u'    -p_z -(u'w)_x-(vw)_y-(ww)_z     -(Uw)_x  -b + f_h U 
c                                  (the last term is not yet implemented!!!!)
c
c=======================================================================

#ifdef enable_back_stratification

c-----------------------------------------------------------------------
c  enable effect of a background stratification B
c
#define apply_back_stratification
c-----------------------------------------------------------------------
c   enable effect of background zonal or meridional advection 
c   no not use both
c
c#define apply_back_zonal_velocity
#define apply_back_meridional_velocity
c-----------------------------------------------------------------------
c    convection accounting for full buoyancy b+B
c
#define apply_back_convection
c-----------------------------------------------------------------------

      module back_stratification_module
      implicit none
      real, allocatable :: back(:,:,:,:)  ! background stratification B, depends on
                                          ! time and x/y only for technical reasons
      real, allocatable :: u0(:,:,:)      ! background zonal/meridional advection U
      end module back_stratification_module

      subroutine init_back_stratification
c-----------------------------------------------------------------------
c     initialize module
c-----------------------------------------------------------------------
      use cpflame_module
      use back_stratification_module
      implicit none
      integer :: k

      if (my_pe==0) print*,' using module background stratification '
#ifdef apply_back_stratification
      if (my_pe==0) print*,' applying background stratification '
#endif
#ifdef apply_back_zonal_velocity
      if (my_pe==0) print*,' applying background zonal velocity '
#endif
#ifdef apply_back_meridional_velocity
      if (my_pe==0) print*,' applying background meridional velocity '
#endif
#ifdef apply_back_convection
      if (my_pe==0) print*,' applying convection '
#endif
#if defined apply_back_meridional_velocity && defined apply_back_zonal_velocity
      if (my_pe==0) then
       print*,' ERROR: background meridional ',
     &           'plus zonal velocity does not work'
      endif
      call halt_stop('in back_stratification')
#endif
      allocate( back(imt,jmt,km,0:2), u0(imt,jmt,km) )
      back=0;u0=0
c    call template routine to set background stratification
      call set_back_stratification(back(:,:,:,0),u0)
      back(:,:,:,1)=back(:,:,:,0); back(:,:,:,2)=back(:,:,:,0)
      call diag_back_stratification
      end subroutine init_back_stratification



      subroutine back_stratification
c-----------------------------------------------------------------------
c    apply effect of background stratification on buoyancy budget
c-----------------------------------------------------------------------
      use cpflame_module
      use back_stratification_module
      implicit none
      integer :: i,j,k,js,je
      real :: adv_fe(imt,jmt,km), adv_ft(imt,jmt,km)
      real :: adv_fn(imt,jmt,km)
      real :: uback(imt,jmt,km,3)

      js=max(2,js_pe); je = min(je_pe,jmt-1)
#ifdef apply_back_stratification
c---------------------------------------------------------------------------------
c     background stratification  
c---------------------------------------------------------------------------------
      adv_fe(:,js_pe:je_pe,:)=0; adv_fn(:,js_pe:je_pe,:)=0
      adv_ft(:,js_pe:je_pe,:)=0
      call adv_flux(adv_fe,adv_fn,adv_ft,back)
      call border_exchg3D(adv_fn,1)
      call setcyclic3D(adv_fe); call setcyclic3D(adv_fn)
      do k=2,km-1
       do j=js,je
        do i=2,imt-1
         b(i,j,k,taup1)=b(i,j,k,taup1)+maskT(i,j,k)*c2dt*( 
     &  -(adv_fe(i,j,k)-adv_fe(i-1,j,k))/dx
     &  -(adv_fn(i,j,k)-adv_fn(i,j-1,k))/dx
     &  -(adv_ft(i,j,k)-adv_ft(i,j,k-1))/dz )
        enddo
       enddo
      enddo
#endif

#ifdef apply_back_zonal_velocity
c---------------------------------------------------------------------------------
c     effect of zonal background velocity
c---------------------------------------------------------------------------------
      do j=js_pe,je_pe
       uback(:,j,:,:)=u(:,j,:,1:3,tau);
       u(:,j,:,1,tau)=u0(:,j,:); u(:,j,:,2:3,tau)=0.
      enddo
      adv_fe(:,js_pe:je_pe,:)=0; adv_fn(:,js_pe:je_pe,:)=0
      adv_ft(:,js_pe:je_pe,:)=0
      call adv_flux(adv_fe,adv_fn,adv_ft,b)
      u(:,js_pe:je_pe,:,1:3,tau)=uback(:,js_pe:je_pe,:,:)
      call border_exchg3D(adv_fn,1)
      call setcyclic3D(adv_fe); call setcyclic3D(adv_fn)
      do k=2,km-1
       do j=js,je
        do i=2,imt-1
         b(i,j,k,taup1)=b(i,j,k,taup1)+maskT(i,j,k)*c2dt*( 
     &  -(adv_fe(i,j,k)-adv_fe(i-1,j,k))/dx
     &  -(adv_fn(i,j,k)-adv_fn(i,j-1,k))/dx
     &  -(adv_ft(i,j,k)-adv_ft(i,j,k-1))/dz )
        enddo
       enddo
      enddo
#endif

#ifdef apply_back_meridional_velocity
c---------------------------------------------------------------------------------
c     effect of zonal background velocity
c---------------------------------------------------------------------------------
      do j=js_pe,je_pe
       uback(:,j,:,:)=u(:,j,:,1:3,tau);
       u(:,j,:,2,tau)=u0(:,j,:); u(:,j,:,1,tau)=0.;u(:,j,:,3,tau)=0.
      enddo
      adv_fe(:,js_pe:je_pe,:)=0; adv_fn(:,js_pe:je_pe,:)=0
      adv_ft(:,js_pe:je_pe,:)=0
      call adv_flux(adv_fe,adv_fn,adv_ft,b)
      u(:,js_pe:je_pe,:,1:3,tau)=uback(:,js_pe:je_pe,:,:)
      call border_exchg3D(adv_fn,1)
      call setcyclic3D(adv_fe); call setcyclic3D(adv_fn)
      do k=2,km-1
       do j=js,je
        do i=2,imt-1
         b(i,j,k,taup1)=b(i,j,k,taup1)+maskT(i,j,k)*c2dt*( 
     &  -(adv_fe(i,j,k)-adv_fe(i-1,j,k))/dx
     &  -(adv_fn(i,j,k)-adv_fn(i,j-1,k))/dx
     &  -(adv_ft(i,j,k)-adv_ft(i,j,k-1))/dz )
        enddo
       enddo
      enddo
#endif

#ifdef apply_back_convection
c---------------------------------------------------------------------------------
c      Convection
c---------------------------------------------------------------------------------
       if (enable_hydrostatic) then
        K_b(:,js:je,:) = 0.0
        do j=js,je
         do k=1,km-1
          do i=2,imt
           if ( b(i,j,k+1,taum1)+back(i,j,k+1,taum1) > 
     &         b(i,j,k   ,taum1) +back(i,j,k,taum1) )
     &         K_b(i,j,k)=1000.*maskW(i,j,k)
          enddo
         enddo
        enddo
        call implicit_tracer_mix(b,K_b)
       endif
#endif
c---------------------------------------------------------------------------------
c      Boundary exchange 
c---------------------------------------------------------------------------------
      call border_exchg3D(b(:,:,:,taup1),2);
      call setcyclic3D(b(:,:,:,taup1) )
      end subroutine back_stratification



      subroutine back_strat_momentum
c-----------------------------------------------------------------------
c    apply effect of background velocity on momentum
c-----------------------------------------------------------------------
      use cpflame_module
      use back_stratification_module
      implicit none
      integer :: i,j,k,js,je
      real :: adv_fe(imt,jmt,km)
      real :: adv_fn(imt,jmt,km)
      real :: adv_ft(imt,jmt,km)

      js=max(2,js_pe); je = min(je_pe,jmt-1)

c-----------------------------------------------------------------------
c    zonal momentum equation
c-----------------------------------------------------------------------
#ifdef apply_back_zonal_velocity
      adv_fe(:,js_pe:je_pe,:)=0.0; adv_fn(:,js_pe:je_pe,:)=0.0; 
      adv_ft(:,js_pe:je_pe,:)=0.0; 
      do k=2,km-1
       do j=js,je
        do i=1,imt-1
         adv_fe(i,j,k)=2*
     &  0.5*( u0(i  ,j,k) + u0(i+1,j,k) )*
     &   (u(i,j,k,1,tau)+u(i+1,j,k,1,tau))
     &        *0.5*maskU(i+1,j,k)*maskU(i,j,k)
        enddo
       enddo
      enddo
      call setcyclic3D(adv_fe)
      do k=2,km-1
       do j=js,je
        do i=2,imt-1
         adv_fn(i,j,k)=
     &  0.5*( u0(i,j,k) + u0(i,j+1,k) )*
     &   (u(i,j,k,2,tau)+u(i+1,j,k,2,tau))
     &        *0.5*maskU(i,j+1,k)*maskU(i,j,k)
        enddo
       enddo
      enddo
      call border_exchg3D(adv_fn,1); call setcyclic3D(adv_fn)
      do k=1,km-1
       do j=js,je
        do i=2,imt-1
         adv_ft(i,j,k)=
     &  0.5*(u0(i,j,k)+u0(i,j,k+1))*
     &   (u(i,j,k,3,tau)+u(i+1,j,k,3,tau))
     &        *0.5*maskU(i,j,k+1)*maskU(i,j,k)
        enddo
       enddo
      enddo
      do k=2,km-1
       do j=js,je
        do i=2,imt-1
         fu(i,j,k)= fu(i,j,k)+maskU(i,j,k)*(
     &    -(adv_fe(i,j,k) - adv_fe(i-1,j,k) )/dx
     &    -(adv_fn(i,j,k) - adv_fn(i,j-1,k) )/dx
     &    -(adv_ft(i,j,k) - adv_ft(i,j,k-1) )/dz  )
        enddo
       enddo
      enddo
#endif
#ifdef apply_back_meridional_velocity
      adv_fn(:,js_pe:je_pe,:)=0.0; 
      do k=2,km-1
       do j=js,je
        do i=2,imt-1
         adv_fn(i,j,k)=
     &  0.5*( u(i,j,k,1,tau) + u(i,j+1,k,1,tau) )*
     &   (u0(i,j,k)+u0(i+1,j,k))
     &        *0.5*maskU(i,j+1,k)*maskU(i,j,k)
        enddo
       enddo
      enddo
      call border_exchg3D(adv_fn,1); call setcyclic3D(adv_fn)
      do k=2,km-1
       do j=js,je
        do i=2,imt-1
         fu(i,j,k)= fu(i,j,k)+maskU(i,j,k)*(
     &    -(adv_fn(i,j,k) - adv_fn(i,j-1,k) )/dx )
        enddo
       enddo
      enddo
#endif

c-----------------------------------------------------------------------
c    meridional momentum equation
c-----------------------------------------------------------------------
#ifdef apply_back_zonal_velocity
      adv_fe(:,js_pe:je_pe,:)=0.0; 
      do k=2,km-1
       do j=js,je
        do i=1,imt-1
         adv_fe(i,j,k)=
     &  0.5*( u(i  ,j,k,2,tau) + u(i+1,j,k,2,tau) )*
     &    (u0(i,j,k)+u0(i,j+1,k))
     &        *0.5*maskV(i+1,j,k)*maskV(i,j,k)
        enddo
       enddo
      enddo
      call setcyclic3D(adv_fe)
      do k=2,km-1
       do j=js,je
        do i=2,imt-1
         fv(i,j,k)= fv(i,j,k)+maskV(i,j,k)*(
     &    -(adv_fe(i,j,k) - adv_fe(i-1,j,k) )/dx)
        enddo
       enddo
      enddo
#endif

#ifdef apply_back_meridional_velocity
      adv_fe(:,js_pe:je_pe,:)=0.0; adv_fn(:,js_pe:je_pe,:)=0.0; 
      adv_ft(:,js_pe:je_pe,:)=0.0; 
      do k=2,km-1
       do j=js,je
        do i=1,imt-1
         adv_fe(i,j,k)=
     &  0.5*( u0(i  ,j,k) + u0(i+1,j,k) )*
     &    (u(i,j,k,1,tau)+u(i,j+1,k,1,tau))
     &        *0.5*maskV(i+1,j,k)*maskV(i,j,k)
        enddo
       enddo
      enddo
      call setcyclic3D(adv_fe)
      do k=2,km-1
       do j=js,je
        do i=2,imt-1
         adv_fn(i,j,k)=2*
     &  0.5*( u0(i,j,k) + u0(i,j+1,k) )*
     &   (u(i,j,k,2,tau)+u(i,j+1,k,2,tau))
     &        *0.5*maskV(i,j+1,k)*maskV(i,j,k)
        enddo
       enddo
      enddo
      call border_exchg3D(adv_fn,1); call setcyclic3D(adv_fn)
      do k=1,km-1
       do j=js,je
        do i=2,imt-1
         adv_ft(i,j,k)=
     &  0.5*( u0(i,j,k) + u0(i,j,k+1) )*
     &   (u(i,j,k,3,tau)+u(i,j+1,k,3,tau))
     &        *0.5*maskV(i,j,k+1)*maskV(i,j,k)
        enddo
       enddo
      enddo
      do k=2,km-1
       do j=js,je
        do i=2,imt-1
         fv(i,j,k)= fv(i,j,k)+maskV(i,j,k)*(
     &    -(adv_fe(i,j,k) - adv_fe(i-1,j,k) )/dx
     &    -(adv_fn(i,j,k) - adv_fn(i,j-1,k) )/dx
     &    -(adv_ft(i,j,k) - adv_ft(i,j,k-1) )/dz  )
        enddo
       enddo
      enddo
#endif

      if (.not. enable_hydrostatic) then
c-----------------------------------------------------------------------
c    vertical momentum equation
c-----------------------------------------------------------------------
#ifdef apply_back_zonal_velocity
       adv_fe(:,js_pe:je_pe,:)=0.0
       do k=2,km-1
        do j=js,je
         do i=1,imt-1
          adv_fe(i,j,k)=
     &     0.5*( u(i  ,j,k,3,tau) + u(i+1,j,k,3,tau) )*
     &         (u0(i,j,k)+u0(i,j,k+1))
     &        *0.5*maskW(i+1,j,k)*maskW(i,j,k)
         enddo
        enddo
       enddo
       call setcyclic3D(adv_fe)
       do k=2,km-1
        do j=js,je
         do i=2,imt-1
          fw(i,j,k)= fw(i,j,k)+maskW(i,j,k)*(
     &     -(adv_fe(i,j,k) - adv_fe(i-1,j,k) )/dx)
         enddo
        enddo
       enddo
#endif
#ifdef apply_back_meridional_velocity
      adv_fn(:,js_pe:je_pe,:)=0.0; 
      do k=2,km-1
       do j=js,je
        do i=2,imt-1
         adv_fn(i,j,k)=
     &     0.5*( u(i,j,k,3,tau) + u(i,j+1,k,3,tau) )*
     &     (u0(i,j,k)+u0(i,j,k+1))
     &        *0.5*maskW(i,j+1,k)*maskW(i,j,k)
        enddo
       enddo
      enddo
      call border_exchg3D(adv_fn,1); call setcyclic3D(adv_fn)
      do k=2,km-1
       do j=js,je
         do i=2,imt-1
          fw(i,j,k)= fw(i,j,k)+maskW(i,j,k)*(
     &     -(adv_fn(i,j,k) - adv_fn(i,j-1,k) )/dx)
        enddo
       enddo
      enddo
#endif
      endif
#if defined apply_back_zonal_velocity || defined apply_back_meridional_velocity
c---------------------------------------------------------------------------------
c       boundary exchange for result
c---------------------------------------------------------------------------------
      call border_exchg3D(fu,1); call setcyclic3D(fu)
      call border_exchg3D(fv,1); call setcyclic3D(fv)
      if (.not. enable_hydrostatic) then
       call border_exchg3D(fw,1); call setcyclic3D(fw)
      endif
#endif
      end subroutine back_strat_momentum



      subroutine diag_back_stratification
c-----------------------------------------------------------------------
c    write diagnostics
c-----------------------------------------------------------------------
      use cpflame_module
      use back_stratification_module
      implicit none
#include "netcdf.inc"
      integer :: ncid,iret,i,j,k,n,npe
      integer :: lon_tdim,lon_udim,z_tdim,z_udim,itimedim
      integer :: lat_tdim,lat_udim,varid
      integer :: dims(4), corner(4), edges(4)
      character :: name*24, unit*16
      real :: a(imt,js_pe:je_pe,km)
      if (my_pe==0) then
       call def_grid_cdf('back_stratification.cdf')
       iret=nf_open('back_stratification.cdf',NF_WRITE,ncid)
       iret=nf_set_fill(ncid, NF_NOFILL, iret)
       call ncredf(ncid, iret)
       iret=nf_inq_dimid(ncid,'xt',lon_tdim)
       iret=nf_inq_dimid(ncid,'xu',lon_udim)
       iret=nf_inq_dimid(ncid,'yt',lat_tdim)
       iret=nf_inq_dimid(ncid,'yu',lat_udim)
       iret=nf_inq_dimid(ncid,'zt',z_tdim)
       iret=nf_inq_dimid(ncid,'zu',z_udim)
#ifdef apply_back_stratification
       dims = (/Lon_tdim,lat_tdim, z_tdim, 1/)
       varid  = ncvdef (ncid,'back', NCFLOAT,3,dims,iret)
       name = 'Background buoyancy'; unit = 'm/s^2'
       call dvcdf(ncid,varid,name,24,unit,16,spval)
#endif
#ifdef apply_back_zonal_velocity
       dims = (/Lon_udim,lat_tdim, z_tdim, 1/)
       varid  = ncvdef (ncid,'u0', NCFLOAT,3,dims,iret)
       name = 'Background zonal velocity'; unit = 'm/s'
       call dvcdf(ncid,varid,name,24,unit,16,spval)
#endif
#ifdef apply_back_meridional_velocity
       dims = (/Lon_tdim,lat_udim, z_tdim, 1/)
       varid  = ncvdef (ncid,'v0', NCFLOAT,3,dims,iret)
       name = 'Background meridional velocity'; unit = 'm/s'
       call dvcdf(ncid,varid,name,24,unit,16,spval)
#endif
       call ncclos (ncid, iret)
      endif

      do npe=0,n_pes-1
       call barrier
       if (my_pe==npe) then
        iret=nf_open('back_stratification.cdf',NF_WRITE,ncid)
        Corner = (/1,js_pe,1,1/); edges=(/imt,je_pe-js_pe+1,km,1/)
#ifdef apply_back_stratification
        iret=nf_inq_varid(ncid,'back',varid)
        a=back(:,js_pe:je_pe,:,tau)
        where( maskT(:,js_pe:je_pe,:) == 0.) a = spval
        iret= nf_put_vara_double(ncid,varid,corner,edges,a)
#endif
#ifdef apply_back_zonal_velocity
        iret=nf_inq_varid(ncid,'u0',varid)
        a=u0(:,js_pe:je_pe,:)
        where( maskU(:,js_pe:je_pe,:) == 0.) a = spval
        iret= nf_put_vara_double(ncid,varid,corner,edges,a)
#endif
#ifdef apply_back_meridional_velocity
        iret=nf_inq_varid(ncid,'v0',varid)
        a=u0(:,js_pe:je_pe,:)
        where( maskV(:,js_pe:je_pe,:) == 0.) a = spval
        iret= nf_put_vara_double(ncid,varid,corner,edges,a)
#endif
        call ncclos (ncid, iret)
       endif
      enddo
      end subroutine diag_back_stratification





#else
      subroutine dummy_back_stratification
      end subroutine dummy_back_stratification
#endif
