#include "options.inc"
    

c=======================================================================
c       calculate numbers
c       linked in the code in driver.F
c=======================================================================

#ifdef enable_diag_numbers

c#define write_threedim

      module numbers_module
      implicit none
      real,allocatable :: cfl_h(:,:,:)
      real,allocatable :: cfl_v(:,:,:)
      real :: cfl_h_max ,cfl_v_max

      real,allocatable :: pec_h(:,:,:)
      real,allocatable :: pec_v(:,:,:)
      real :: pec_h_max ,pec_v_max

      real,allocatable :: rey_h(:,:,:)
      real,allocatable :: rey_v(:,:,:)
      real :: rey_h_max ,rey_v_max

      real,allocatable :: ro(:,:,:)
      real :: ro_max
      real,allocatable :: ri(:,:,:)
      real :: ri_min
      real,allocatable :: Ek_h(:),Ek_v(:)

      real, parameter :: eps = 1e-20
      end module numbers_module
 
      subroutine init_numbers
c=======================================================================
c       allocate work space
c=======================================================================
      use cpflame_module
      use numbers_module
      implicit none
      real :: L_r, c, cfl

      if (my_pe==0) then
       print*,''
       print*,' intializing diagnostic for numerical'
       print*,' and physical critria and numbers '
       print*,''
      endif

      allocate( cfl_h(imt,jmt,km) ); cfl_h=0.0
      allocate( cfl_v(imt,jmt,km) ); cfl_v=0.0
      allocate( pec_h(imt,jmt,km) ); pec_h=0.0
      allocate( pec_v(imt,jmt,km) ); pec_v=0.0
      allocate( rey_h(imt,jmt,km) ); rey_h=0.0
      allocate( rey_v(imt,jmt,km) ); rey_v=0.0
      allocate( ro(imt,jmt,km)    ); ro=0.0
      allocate( ri(imt,jmt,km)    ); ri=0.0
      allocate( Ek_h(jmt),Ek_v(jmt));Ek_h=0.0;Ek_v=0.0
      call init_numbers_diag
c
c=======================================================================
c      check CFL criterion for barotropic Rossby wave speed
c=======================================================================
c
      L_r = sqrt( g * maxval(ht )) /minval(coriolis_t)
      L_r = min( L_r, sqrt( sqrt( g * maxval(ht )) / beta ) )
      c = beta * L_r**2
      cfl = dt_in/dx *c
      if (my_pe==0) then
        print *,' barotropic Rossby radius     ', L_r/1e3 , ' km'
        print *,' barotropic Rossby wave speed ', c , ' m/s'
        print *,' Courant number for that speed ', cfl
        if (cfl >0.5) then
          print*,'---------------------------------------------'
          print*,'!!!!! WARNING CFL FOR ROSSBY WAVES > 0.5 !!!!'
          print*,'---------------------------------------------'
        endif
      endif

      if (my_pe==0) print*,' Checking horizontal diffusive criteria'
      if (dt_in > dx**2/(2.*max(K_h,1e-8))) 
     &            call halt_stop('K_h too large')
      if (dt_in > dx**2/(2.*max(A_h,1e-8))) 
     &            call halt_stop('A_h too large')
      if (dt_in > dz**2/(2.*max(K_v,1e-8))) 
     &            call halt_stop('K_v too large ')
      if (dt_in > dz**2/(2.*max(A_v,1e-8))) 
     &            call halt_stop('A_v too large ')
      if (my_pe==0) then
        print*,' max. possible K_h/A_h = ',dx**2/dt_in/2.0
        print*,' max. possible K_v/A_v = ',dz**2/dt_in/2.0
      endif
      end subroutine init_numbers



      subroutine init_numbers_diag
c-----------------------------------------------------------------------
c=======================================================================
c      initialize NetCDF snapshot file
c=======================================================================
c-----------------------------------------------------------------------
      use cpflame_module
      use numbers_module
      implicit none
#include "netcdf.inc"
      integer :: ncid,iret,i,j,k,n
      integer :: lon_tdim,lon_udim,z_tdim,z_udim,itimedim
      integer :: lat_tdim,lat_udim, vid
      integer :: dims(4), corner(4), edges(4)
      character (len=80) :: name, unit
      real :: fxa,fxb,fxc
c
c=======================================================================
c       Ekman number on the grid
c=======================================================================
c
      fxa = A_h; fxb = A_v
      if (enable_biharmonic_friction) fxa = max(A_h,Ahbi/dx**2)
      if (enable_vert_biha_friction)  fxb = max(A_v,Avbi/dz**2)
      do j=1,jmt
       fxc = eps + abs(coriolis_t(j))
       Ek_h(j) = fxa/dx**2/fxc
       Ek_v(j) = fxb/dz**2/fxc
      enddo

      if (my_pe==0) then

      call def_grid_cdf('numbers.cdf')
      iret=nf_open('numbers.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,'yt',lat_tdim)
      iret=nf_inq_dimid(ncid,'zt',z_tdim)
      iret=nf_inq_dimid(ncid,'zu',z_udim)
      iret=nf_inq_dimid(ncid,'Time',itimedim)

      vid = ncvdef (ncid,'CFL_h_max', NCFLOAT,1,itimedim,iret)
      name = 'Max. hor. Courant number'; unit = ' '
      call dvcdf(ncid,vid,name,len_trim(name),unit,len_trim(unit),spval)
      vid = ncvdef (ncid,'CFL_v_max', NCFLOAT,1,itimedim,iret)
      name = 'Max. vert. Courant number'; unit = ' '
      call dvcdf(ncid,vid,name,len_trim(name),unit,len_trim(unit),spval)

      vid = ncvdef (ncid,'Pec_h_max', NCFLOAT,1,itimedim,iret)
      name = 'Max. horz. Peclet number'; unit = ' '
      call dvcdf(ncid,vid,name,len_trim(name),unit,len_trim(unit),spval)
      vid = ncvdef (ncid,'Pec_v_max', NCFLOAT,1,itimedim,iret)
      name = 'Max. vert. Peclet number'; unit = ' '
      call dvcdf(ncid,vid,name,len_trim(name),unit,len_trim(unit),spval)

      vid = ncvdef (ncid,'Rey_h_max', NCFLOAT,1,itimedim,iret)
      name = 'Max. horz. Reynolds number'; unit = ' '
      call dvcdf(ncid,vid,name,len_trim(name),unit,len_trim(unit),spval)
      vid = ncvdef (ncid,'Rey_v_max', NCFLOAT,1,itimedim,iret)
      name = 'Max. vert. Reynolds number'; unit = ' '
      call dvcdf(ncid,vid,name,len_trim(name),unit,len_trim(unit),spval)

      vid = ncvdef (ncid,'Ro_max', NCFLOAT,1,itimedim,iret)
      name = 'Max. Rossby number'; unit = ' '
      call dvcdf(ncid,vid,name,len_trim(name),unit,len_trim(unit),spval)
      vid = ncvdef (ncid,'Ri_min', NCFLOAT,1,itimedim,iret)
      name = 'Min. Richardson number'; unit = ' '
      call dvcdf(ncid,vid,name,len_trim(name),unit,len_trim(unit),spval)

#ifdef write_threedim
      dims = (/Lon_tdim,lat_tdim, z_tdim, iTimedim/)
      vid = ncvdef (ncid,'CFL_h', NCFLOAT,4,dims,iret)
      name = 'Horizontal Courant number'; unit = ' '
      call dvcdf(ncid,vid,name,len_trim(name),unit,len_trim(unit),spval)
      vid = ncvdef (ncid,'CFL_v', NCFLOAT,4,dims,iret)
      name = 'Vertical Courant number'; unit = ' '
      call dvcdf(ncid,vid,name,len_trim(name),unit,len_trim(unit),spval)

      vid = ncvdef (ncid,'Pec_h', NCFLOAT,4,dims,iret)
      name = 'Horizontal Peclet number'; unit = ' '
      call dvcdf(ncid,vid,name,len_trim(name),unit,len_trim(unit),spval)
      vid = ncvdef (ncid,'Pec_v', NCFLOAT,4,dims,iret)
      name = 'Vertical Peclet number'; unit = ' '
      call dvcdf(ncid,vid,name,len_trim(name),unit,len_trim(unit),spval)

      vid = ncvdef (ncid,'Rey_h', NCFLOAT,4,dims,iret)
      name = 'Horizontal Reynolds number'; unit = ' '
      call dvcdf(ncid,vid,name,len_trim(name),unit,len_trim(unit),spval)
      vid = ncvdef (ncid,'Rey_v', NCFLOAT,4,dims,iret)
      name = 'Vertical Reynolds number'; unit = ' '
      call dvcdf(ncid,vid,name,len_trim(name),unit,len_trim(unit),spval)

      vid = ncvdef (ncid,'Ro', NCFLOAT,4,dims,iret)
      name = 'Grid Rossby number'; unit = ' '
      call dvcdf(ncid,vid,name,len_trim(name),unit,len_trim(unit),spval)

      vid = ncvdef (ncid,'Nsqr', NCFLOAT,4,dims,iret)
      name = 'Square of stability frequency'; unit = '1/s^2'
      call dvcdf(ncid,vid,name,len_trim(name),unit,len_trim(unit),spval)

      dims = (/Lon_tdim,lat_tdim, z_udim, iTimedim/)

      vid = ncvdef (ncid,'Ri', NCFLOAT,4,dims,iret)
      name = 'Richardson number'; unit = ' '
      call dvcdf(ncid,vid,name,len_trim(name),unit,len_trim(unit),spval)
#endif

      vid = ncvdef (ncid,'f', NCFLOAT,1,lat_tdim,iret)
      name = 'Vert. Coriolis parameter'; unit = '1/s'
      call dvcdf(ncid,vid,name,len_trim(name),unit,len_trim(unit),spval)

      vid = ncvdef (ncid,'Ek_h', NCFLOAT,1,lat_tdim,iret)
      name = 'Hor. Ekman number'; unit = '1/s'
      call dvcdf(ncid,vid,name,len_trim(name),unit,len_trim(unit),spval)

      vid = ncvdef (ncid,'Ek_v', NCFLOAT,1,lat_tdim,iret)
      name = 'Vert. Ekman number'; unit = ' '
      call dvcdf(ncid,vid,name,len_trim(name),unit,len_trim(unit),spval)

      call ncendf(ncid, iret)

      iret=nf_inq_varid(ncid,'f',vid)
      iret= nf_put_vara_double(ncid,vid,1,jmt,coriolis_t)
      iret=nf_inq_varid(ncid,'Ek_h',vid)
      iret= nf_put_vara_double(ncid,vid,1,jmt,Ek_h)
      iret=nf_inq_varid(ncid,'Ek_v',vid)
      iret= nf_put_vara_double(ncid,vid,1,jmt,Ek_v)

      call ncclos (ncid, iret)
      endif
      end subroutine init_numbers_diag




      subroutine diag_numbers
c-----------------------------------------------------------------------
c=======================================================================
c      write to NetCDF snapshot file
c=======================================================================
c-----------------------------------------------------------------------
      use cpflame_module
      use numbers_module
#ifdef enable_tke_closure
      use tke_closure_module
#endif
#ifdef enable_smagorinsky_friction
      use smagorinsky_module
#endif
      implicit none
#include "netcdf.inc"
      integer :: ncid,iret,n,npe, corner(4), edges(4)
      real :: a(imt,js_pe:je_pe,km),fxa,fxb,fxc,uabs
      integer :: itdimid,ilen,varid,itimeid
      integer :: i,j,js,je,k
      type(time_type) :: time
      real :: Nsqrw(imt,js_pe:je_pe,km)
#ifndef enable_tke_closure
      real :: Nsqr(imt,jmt,km)
#endif


      if (snapshot_time_step.or.initial_time==current_time) then
      js=max(2,js_pe); je = min(je_pe,jmt-1)
c
c=======================================================================
c      CFL cirterion
c=======================================================================
c
      cfl_h(:,js:je,:)=0.; cfl_h_max=0
      fxa = dt_in/dx 
      do k=2,km  -1
       do j=js,je
        do i=2,imt-1
         uabs=max(     abs(u(i,j,k,1,tau)) , abs(u(i-1,j,k,1,tau)) )
         uabs=max(uabs,abs(u(i,j,k,2,tau)) , abs(u(i,j-1,k,2,tau)) )
         cfl_h(i,j,k)=fxa*uabs
         cfl_h_max = max( cfl_h_max, cfl_h(i,j,k) )
        enddo
       enddo
      enddo
      call global_max(cfl_h_max)

      cfl_v(:,js:je,:)=0.;cfl_v_max=0
      if (.not. enable_hydrostatic )  then
       fxa = dt_in/dz
       do k=2,km  -1
        do j=js,je
         do i=2,imt-1
          uabs=max(abs(u(i,j,k,3,tau)),abs(u(i,j,k-1,3,tau)) )
          cfl_v(i,j,k)=fxa*uabs
          cfl_v_max = max( cfl_v_max, cfl_v(i,j,k) )
         enddo
        enddo
       enddo
       call global_max(cfl_v_max)
      endif
c
c=======================================================================
c      Peclet number 
c=======================================================================
c
      pec_h(:,js:je,:)=0.; pec_v(:,js:je,:)=0.
      pec_h_max=0; pec_v_max=0
      fxa = dx/(eps+K_h)
      if (enable_biharmonic_diffusion) fxa = dx/max(K_h,Khbi/dx**2)
      fxb = dz/(eps+K_v)
      if (enable_vert_biha_diffusion)  fxb = dz/max(K_v,Kvbi/dz**2)
      do k=2,km-1
       do j=js,je
        do i=2,imt-1
         uabs=max(     abs(u(i,j,k,1,tau)) , abs(u(i-1,j,k,1,tau)) )
         uabs=max(uabs,abs(u(i,j,k,2,tau)) , abs(u(i,j-1,k,2,tau)) )
#ifdef enable_tke_closure
         fxa = dx/max(diffT(i,j,k),eps)
         fxb = dz/max(diffT(i,j,k),eps)
#endif
         pec_h(i,j,k)= fxa*uabs
         pec_h_max   = max( pec_h_max, pec_h(i,j,k) )

         uabs=max(abs(u(i,j,k,3,tau)),abs(u(i,j,k-1,3,tau)) )
         pec_v(i,j,k)= fxb*uabs
         pec_v_max   = max( pec_v_max, pec_v(i,j,k) )
        enddo
       enddo
      enddo
      call global_max(pec_h_max)
      call global_max(pec_v_max)
c
c=======================================================================
c      Reynolds number 
c=======================================================================
c
      rey_h(:,js:je,:)=0.; rey_v(:,js:je,:)=0.
      rey_h_max=0; rey_v_max=0
      fxa = dx/(eps+A_h)
      if (enable_biharmonic_friction) fxa = dx/max(A_h,Ahbi/dx**2)
      fxb = dz/(eps+A_v)
      if (enable_vert_biha_friction)  fxb = dz/max(A_v,Avbi/dz**2)
      do k=2,km-1
       do j=js,je
        do i=2,imt-1
         uabs=max(     abs(u(i,j,k,1,tau)) , abs(u(i-1,j,k,1,tau)) )
         uabs=max(uabs,abs(u(i,j,k,2,tau)) , abs(u(i,j-1,k,2,tau)) )
#ifdef enable_tke_closure
         fxa = dx/max(diffU(i,j,k),eps)
         fxb = dz/max(diffU(i,j,k),eps)
#endif
#ifdef enable_smagorinsky_friction
         fxa = dx/max( diff_T(i,j,k) ,eps)
#endif
         rey_h(i,j,k)= fxa*uabs
         rey_h_max   = max( rey_h_max, rey_h(i,j,k) )

         uabs=max(abs(u(i,j,k,3,tau)),abs(u(i,j,k-1,3,tau)) )
         rey_v(i,j,k)= fxb*uabs
         rey_v_max   = max( rey_v_max, rey_v(i,j,k) )
        enddo
       enddo
      enddo
      call global_max(rey_h_max)
      call global_max(rey_v_max)
c
c=======================================================================
c      output
c=======================================================================
c
      if (my_pe==0) then
         print*,' maximum hor. Courant number  : ',cfl_h_max
         if (.not. enable_hydrostatic ) 
     &   print*,' maximum vert.Courant number  : ',cfl_v_max
         print*,' maximum hor. Peclet number   : ',pec_h_max
         print*,' maximum vert.Peclet number   : ',pec_v_max
         print*,' maximum hor. Reynolds number : ',rey_h_max
         print*,' maximum vert.Reynolds number : ',rey_v_max
         if (cfl_h_max > 1.0 ) 
     &           print*,' WARNING: hor.  CFL criterion violated'
         if (cfl_v_max > 1.0 .and. .not. enable_hydrostatic ) 
     &           print*,' WARNING: vert. CFL criterion violated'
      endif

c
c=======================================================================
c       Rossby number on the grid
c=======================================================================
c
      ro(:,js:je,:)=0.; ro_max=0.
      do k=2,km-1
       do j=js,je
        do i=2,imt-1
         fxa = 1./dx/(eps + abs(coriolis_t(j)))
         ro(i,j,k)=fxa*max(abs(u(i,j,k,1,tau)),abs(u(i,j,k,2,tau)))
         ro_max = max( ro_max, ro(i,j,k) )
        enddo
       enddo
      enddo
      call global_max(ro_max)
c---------------------------------------------------------------------------------
c=======================================================================
c       stability freq.
c=======================================================================
c---------------------------------------------------------------------------------
      Nsqrw=0
      do k=1,km-1
       do j=js,je
        Nsqrw(:,j,k)=-(b(:,j,k+1,tau)-b(:,j,k,tau))/dz
       enddo
      enddo
#ifndef enable_tke_closure
c---------------------------------------------------------------------------------
c=======================================================================
c      Interpolate Nsqr vertically on T grid and bound Nsqr
c=======================================================================
c---------------------------------------------------------------------------------
      do k=2,km-1
       do j=js,je
         Nsqr(:,j,k)=maskT(:,j,k)*(
     &     Nsqrw(:,j,k  )*maskW(:,j,k  )+Nsqrw(:,j,k-1)*maskW(:,j,k-1) )
     &     /(maskW(:,j,k)+maskW(:,j,k-1)+eps)
       enddo
      enddo
      call border_exchg3D(Nsqr,1)
      call setcyclic3D(Nsqr )
#endif
c
c=======================================================================
c       Richardson Number
c=======================================================================
c
      ri(:,js:je,:)=0.; ri_min=1/eps
      do k=1,km-1
       do j=js,je
        do i=2,imt-1
         fxa=0.5*( (u(i  ,j,k+1,1,tau)-u(i  ,j,k,1,tau))/dz )**2
     &      +0.5*( (u(i-1,j,k+1,1,tau)-u(i-1,j,k,1,tau))/dz )**2
         fxb=0.5*( (u(i,j  ,k+1,2,tau)-u(i,j  ,k,2,tau))/dz )**2
     &      +0.5*( (u(i,j-1,k+1,2,tau)-u(i,j-1,k,2,tau))/dz )**2
         Ri(i,j,k)=max(0.0,Nsqrw(i,j,k))/(eps + fxa+fxb)
         if (maskW(i,j,k)==1.0) ri_min = min( ri_min, ri(i,j,k) )
        enddo
       enddo
      enddo
      call global_max(ri_min)
c
c=======================================================================
c      output
c=======================================================================
c
       if (my_pe==0) then
         print*,' maximum grid  Rossby number  : ',ro_max
         print*,' minimal   Richardson number  : ',ri_min
       endif

       do npe=0,n_pes-1
        if (my_pe==npe) then
         iret=nf_open('numbers.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)
         if (my_pe==0) then
          ilen=ilen+1
          time = current_time-initial_time
          fxa = time%days + time%seconds/86400.
          iret= nf_put_vara_double(ncid,itimeid,ilen,1,fxa)
         endif
         iret=nf_inq_varid(ncid,'CFL_h_max',varid)
         iret= nf_put_vara_double(ncid,varid,ilen,1,cfl_h_max)
         iret=nf_inq_varid(ncid,'CFL_v_max',varid)
         iret= nf_put_vara_double(ncid,varid,ilen,1,cfl_v_max)

         iret=nf_inq_varid(ncid,'Pec_h_max',varid)
         iret= nf_put_vara_double(ncid,varid,ilen,1,pec_h_max)
         iret=nf_inq_varid(ncid,'Pec_v_max',varid)
         iret= nf_put_vara_double(ncid,varid,ilen,1,pec_v_max)

         iret=nf_inq_varid(ncid,'Rey_h_max',varid)
         iret= nf_put_vara_double(ncid,varid,ilen,1,Rey_h_max)
         iret=nf_inq_varid(ncid,'Rey_v_max',varid)
         iret= nf_put_vara_double(ncid,varid,ilen,1,Rey_v_max)

         iret=nf_inq_varid(ncid,'Ro_max',varid)
         iret= nf_put_vara_double(ncid,varid,ilen,1,ro_max)
         iret=nf_inq_varid(ncid,'Ri_min',varid)
         iret= nf_put_vara_double(ncid,varid,ilen,1,ri_min)

#ifdef write_threedim
         Corner = (/1,js_pe,1,ilen/); 
         edges  = (/imt,je_pe-js_pe+1,km,1/)
         iret=nf_inq_varid(ncid,'CFL_h',varid)
         a=cfl_h(:,js_pe:je_pe,:)
         where( maskT(:,js_pe:je_pe,:) == 0.) a = spval
         iret= nf_put_vara_double(ncid,varid,corner,edges,a)
         iret=nf_inq_varid(ncid,'CFL_v',varid)
         a=cfl_v(:,js_pe:je_pe,:)
         where( maskT(:,js_pe:je_pe,:) == 0.) a = spval
         iret= nf_put_vara_double(ncid,varid,corner,edges,a)

         iret=nf_inq_varid(ncid,'Pec_h',varid)
         a=Pec_h(:,js_pe:je_pe,:)
         where( maskT(:,js_pe:je_pe,:) == 0.) a = spval
         iret= nf_put_vara_double(ncid,varid,corner,edges,a)
         iret=nf_inq_varid(ncid,'Pec_v',varid)
         a=Pec_v(:,js_pe:je_pe,:)
         where( maskT(:,js_pe:je_pe,:) == 0.) a = spval
         iret= nf_put_vara_double(ncid,varid,corner,edges,a)

         iret=nf_inq_varid(ncid,'Rey_h',varid)
         a=Rey_h(:,js_pe:je_pe,:)
         where( maskT(:,js_pe:je_pe,:) == 0.) a = spval
         iret= nf_put_vara_double(ncid,varid,corner,edges,a)
         iret=nf_inq_varid(ncid,'Rey_v',varid)
         a=Rey_v(:,js_pe:je_pe,:)
         where( maskT(:,js_pe:je_pe,:) == 0.) a = spval
         iret= nf_put_vara_double(ncid,varid,corner,edges,a)

         iret=nf_inq_varid(ncid,'Ro',varid)
         a=ro(:,js_pe:je_pe,:)
         where( maskT(:,js_pe:je_pe,:) == 0.) a = spval
         iret= nf_put_vara_double(ncid,varid,corner,edges,a)
         iret=nf_inq_varid(ncid,'Nsqr',varid)
         a=Nsqr(:,js_pe:je_pe,:)
         where( maskT(:,js_pe:je_pe,:) == 0.) a = spval
         iret= nf_put_vara_double(ncid,varid,corner,edges,a)
         iret=nf_inq_varid(ncid,'Ri',varid)
         a=ri(:,js_pe:je_pe,:)
         where( maskT(:,js_pe:je_pe,:) == 0.) a = spval
         iret= nf_put_vara_double(ncid,varid,corner,edges,a)
#endif
         call ncclos (ncid, iret)
        endif
        call barrier
       enddo
      endif

      end subroutine diag_numbers

#else
      subroutine numbers_dummy
      end
#endif
