#include "options.inc"


c=======================================================================
c       overturning streamfunction diagnostics
c       linked in the code in driver only
c=======================================================================

#ifdef enable_diag_over

      module overturning_module
      implicit none
      real, allocatable :: bax(:)
      integer :: bdim 
      real, allocatable :: phiz(:,:)
      real, allocatable :: phiy(:,:)
      real, allocatable :: phix(:,:)
#ifdef enable_bolus_velocity
      real, allocatable :: phiz_b(:,:),phiy_b(:,:),phix_b(:,:)
#endif
      end module overturning_module


      subroutine init_overturning
      use cpflame_module
      use overturning_module
      implicit none
#include "netcdf.inc"
      integer :: ncid,iret,i
      integer :: z_udim,itimedim,varid,bdimid,bid
      integer :: lon_udim,lat_udim, dims(4)
      character :: name*64, unit*32
      real    :: db,b1=-0.032,b2=0.002

      if (my_pe==0) then
       print*,''
       print*,' Initializing overturning module '
       print*,''
      endif

      bdim=40
      allocate(bax(bdim));bax=0.
      db=(b2-b1)/bdim
      bax(1)=b1
      do i=2,bdim
       bax(i)=bax(i-1)+db
      enddo

      allocate( phiz(bdim,km) )
      allocate( phiy(bdim,jmt) )
      allocate( phix(bdim,imt) )
#ifdef enable_bolus_velocity
      allocate( phiz_b(bdim,km) )
      allocate( phiy_b(bdim,jmt) )
      allocate( phix_b(bdim,imt) )
#endif

      if (my_pe==0) then 

       call def_grid_cdf('over.cdf')
       iret=nf_open('over.cdf',NF_WRITE,ncid)
       iret=nf_set_fill(ncid, NF_NOFILL, iret)
       call ncredf(ncid, iret)
       bdimid= ncddef(ncid, 'bdim', bdim, iret)
       bid   = ncvdef (ncid,'bdim',NCFLOAT,1,bdimid,iret)
       name = 'buoyancy'; unit = 'm/s^2'
       call ncaptc(ncid,bid,'long_name',NCCHAR,len_trim(name),name,iret) 
       call ncaptc(ncid,bid,'units',NCCHAR,len_trim(unit),unit,iret) 

       iret=nf_inq_dimid(ncid,'xu',lon_udim)
       iret=nf_inq_dimid(ncid,'yu',lat_udim)
       iret=nf_inq_dimid(ncid,'zu',z_udim)
       iret=nf_inq_dimid(ncid,'Time',itimedim)

       dims = (/bdimid, z_udim, iTimedim,1/)
       varid = ncvdef (ncid,'psiz', NCFLOAT,3,dims,iret)
       name = 'Streamfunction'; unit = 'Sv'
       call dvcdf(ncid,varid,name,len_trim(name),
     &            unit,len_trim(unit),spval) 

       dims = (/bdimid, lat_udim, iTimedim,1/)
       varid = ncvdef (ncid,'psiy', NCFLOAT,3,dims,iret)
       call dvcdf(ncid,varid,name,len_trim(name),
     &            unit,len_trim(unit),spval) 

       dims = (/bdimid, lon_udim, iTimedim,1/)
       varid = ncvdef (ncid,'psix', NCFLOAT,3,dims,iret)
       call dvcdf(ncid,varid,name,len_trim(name),
     &            unit,len_trim(unit),spval) 

#ifdef enable_bolus_velocity
       dims = (/bdimid, z_udim, iTimedim,1/)
       varid = ncvdef (ncid,'psiz_b', NCFLOAT,3,dims,iret)
       name = 'Bolus streamfunction'; unit = 'Sv'
       call dvcdf(ncid,varid,name,len_trim(name),
     &            unit,len_trim(unit),spval) 

       dims = (/bdimid, lat_udim, iTimedim,1/)
       varid = ncvdef (ncid,'psiy_b', NCFLOAT,3,dims,iret)
       call dvcdf(ncid,varid,name,len_trim(name),
     &            unit,len_trim(unit),spval) 

       dims = (/bdimid, lon_udim, iTimedim,1/)
       varid = ncvdef (ncid,'psix_b', NCFLOAT,3,dims,iret)
       call dvcdf(ncid,varid,name,len_trim(name),
     &            unit,len_trim(unit),spval) 
#endif

       call ncendf(ncid, iret)
       iret= nf_put_vara_double(ncid,bid,1,bdim,bax)
       call ncclos (ncid, iret)
      endif ! my_pe ==0
      end subroutine init_overturning




      subroutine overturning
      use cpflame_module
      use overturning_module
#ifdef enable_bolus_velocity
      use bolus_module
#endif
      implicit none
#include "netcdf.inc"
      integer :: i,j,k,js,je,n
      integer :: ncid,iret,corner(4), edges(4)
      integer :: itdimid,ilen,itimeid,varid
      type(time_type) :: time
      character :: name*64, unit*32
      real :: ww,fxa, bb

      if (snapshot_time_step.or.initial_time==current_time) then

       js=max(2,js_pe); je = min(je_pe,jmt-1)
       phiz=0;phiy=0;phix=0
#ifdef enable_bolus_velocity
       phiz_b=0;phiy_b=0;phix_b=0
#endif

       do n=1,bdim

        do k=2,km-1
         do j=js,je 
          do i=2,imt-1
           ww=u(i,j,k,3,tau)*maskW(i,j,k)
           bb=0.5*( b(i,j,k,tau)+b(i,j,k+1,tau) )
           if( bb<bax(n) ) ww=0
           phiz(n,k)=phiz(n,k)+ww*dx*dx*1e-6
          enddo 
         enddo 
         call global_sum_vec(phiz(n,k),1)
        enddo 

        do k=2,km-1
         do j=js,je 
          do i=2,imt-1
           ww=u(i,j,k,1,tau)*maskU(i,j,k)
           bb=0.5*( b(i,j,k,tau)+b(i+1,j,k,tau) )
           if( bb<bax(n) ) ww=0
           phix(n,i)=phix(n,i)+ww*dx*dx*1e-6
          enddo 
         enddo 
        enddo 

        do j=js,je 
         do k=2,km-1
          do i=2,imt-1
           ww=u(i,j,k,2,tau)*maskV(i,j,k)
           bb=0.5*( b(i,j,k,tau)+b(i,j+1,k,tau))
           if( bb<bax(n) ) ww=0
           phiy(n,j)=phiy(n,j)+ww*dx*dz*1e-6
          enddo 
         enddo 
        enddo 

#ifdef enable_bolus_velocity
        do k=2,km-1
         do j=js,je 
          do i=2,imt-1
           ww=wb(i,j,k)*maskW(i,j,k)
           bb=0.5*( b(i,j,k,tau)+b(i,j,k+1,tau) )
           if( bb<bax(n) ) ww=0
           phiz_b(n,k)=phiz_b(n,k)+ww*dx*dx*1e-6
          enddo 
         enddo 
         call global_sum_vec(phiz_b(n,k),1)
        enddo 

        do k=2,km-1
         do j=js,je 
          do i=2,imt-1
           ww=ub(i,j,k)*maskU(i,j,k)
           bb=0.5*( b(i,j,k,tau)+b(i+1,j,k,tau) )
           if( bb<bax(n) ) ww=0
           phix_b(n,i)=phix_b(n,i)+ww*dx*dx*1e-6
          enddo 
         enddo 
        enddo 

        do j=js,je 
         do k=2,km-1
          do i=2,imt-1
           ww=vb(i,j,k)*maskV(i,j,k)
           bb=0.5*( b(i,j,k,tau)+b(i,j+1,k,tau))
           if( bb<bax(n) ) ww=0
           phiy_b(n,j)=phiy_b(n,j)+ww*dx*dz*1e-6
          enddo 
         enddo 
        enddo 
#endif
       enddo 

       do i=2,imt-1
         call global_sum_vec(phix(:,i),bdim)
#ifdef enable_bolus_velocity
         call global_sum_vec(phix_b(:,i),bdim)
#endif
       enddo 

      do n=1,n_pes-1
        if (my_pe==n) then
         call send_integer(js_pe,1,0,0)
         call send_integer(je_pe,1,0,0)
         do j=js_pe,je_pe
          call send_real(phiy(:,j),bdim,0,0)
#ifdef enable_bolus_velocity
          call send_real(phiy_b(:,j),bdim,0,0)
#endif
         enddo
        elseif (my_pe==0) then
         call recv_integer(js,1,n,0)
         call recv_integer(je,1,n,0)
         do j=js,je
          call recv_real(phiy(:,j),bdim,n,0)
#ifdef enable_bolus_velocity
          call recv_real(phiy_b(:,j),bdim,n,0)
#endif
         enddo
        endif
        call barrier
      enddo

       if (my_pe==0) then

         iret=nf_open('over.cdf',NF_WRITE,ncid)
         iret=nf_set_fill(ncid, NF_NOFILL, iret)
         iret=nf_inq_dimid(ncid,'Time',itdimid)
         iret=nf_inq_dimlen(ncid, itdimid,ilen)
         iret=nf_inq_varid(ncid,'Time',itimeid)
         ilen=ilen+1
         time = current_time-initial_time
         fxa = time%days + time%seconds/86400.
         iret= nf_put_vara_double(ncid,itimeid,ilen,1,fxa)

         Corner = (/1,1,ilen,1/); edges  = (/bdim,km,1,1/)
         iret=nf_inq_varid(ncid,'psiz',varid)
         iret= nf_put_vara_double(ncid,varid,corner,edges,phiz)
#ifdef enable_bolus_velocity
         iret=nf_inq_varid(ncid,'psiz_b',varid)
         iret= nf_put_vara_double(ncid,varid,corner,edges,phiz_b)
#endif

         Corner = (/1,1,ilen,1/); edges  = (/bdim,jmt,1,1/)
         iret=nf_inq_varid(ncid,'psiy',varid)
         iret= nf_put_vara_double(ncid,varid,corner,edges,phiy)
#ifdef enable_bolus_velocity
         iret=nf_inq_varid(ncid,'psiy_b',varid)
         iret= nf_put_vara_double(ncid,varid,corner,edges,phiy_b)
#endif

         Corner = (/1,1,ilen,1/); edges  = (/bdim,imt,1,1/)
         iret=nf_inq_varid(ncid,'psix',varid)
         iret= nf_put_vara_double(ncid,varid,corner,edges,phix)
#ifdef enable_bolus_velocity
         iret=nf_inq_varid(ncid,'psix_b',varid)
         iret= nf_put_vara_double(ncid,varid,corner,edges,phix_b)
#endif


         call ncclos (ncid, iret)
        endif
        call barrier
      endif
      end subroutine overturning

#else
      subroutine overturning_dummy
      end
#endif












