#include "options.inc"


c=======================================================================
c       zonal averaging diagnostics
c       linked in the code in driver only
c=======================================================================

#ifdef enable_diag_zonalave

c 
c  Reference buoyancy profile given by N_0 
c 
c  Note that the damping time scales of the resoring zones have to 
c  be adjusted in subroutine init_zonal_averages to match 
c  the setting in  subroutine restoring_zones
c

#define diag_eddy_fluxes
c#define diag_eddy_flux_tendency
c#define diag_energy_budget
c#define diag_potential_vorticity


      module zonal_averages_module
c=======================================================================
c       Module for zonal mean diagnostics
c=======================================================================
      implicit none
      real, allocatable, dimension(:,:) :: ulen,vlen,wlen,blen
      real, allocatable, dimension(:,:) :: um,vm,wm,bm,pm
      real, allocatable, dimension(:,:,:)   :: vb2,wb2 ,tm
      real, allocatable, dimension(:,:,:,:) :: vt2,wt2 
      real, allocatable, dimension(:,:) :: eke,bws,epe,T2,T4u,T4v,T4
      real, allocatable, dimension(:,:) :: mke,mpe,bwq,rest,restm
      real, allocatable, dimension(:,:) :: eps_b,eps_u,eps_v
      real, allocatable, dimension(:,:) :: eps_kh,eps_kv
      real, allocatable, dimension(:,:) :: vps,wps,veke,weke,vepe,wepe
      real, allocatable, dimension(:,:) :: vv,vvb,ageo,vw,ageo2,uv,vvu
      real, allocatable, dimension(:,:) :: agev1,agev2
      real, allocatable, dimension(:,:) :: mpv_qg,mpv_er,mrv_er
      real, allocatable, dimension(:,:,:) :: vq_qg,vq_er,wq_er,wq_qg
      real, allocatable, dimension(:,:) :: vvq_qg,vvq_er

      integer :: nr_var = 5  ! highest moment to be calculated
      integer :: nr_tr  = 0  ! numbers of tracers to be considered
      real, allocatable, dimension(:) :: t_rest
      real,parameter :: epsln     = 1.0e-20
      end module zonal_averages_module



      subroutine zonal_averages
c=======================================================================
c      Verteiler function
c=======================================================================
      use cpflame_module
      use zonal_averages_module
      implicit none
      if (snapshot_time_step.or.initial_time==current_time) then
       call zonal_averages_means
       call zonal_averages_eddy_fluxes
       call zonal_averages_energy
       call zonal_averages_flux_tendency
       call zonal_averages_pv_fluxes
       call zonal_averages_write
      endif
      end subroutine zonal_averages



      subroutine init_zonal_averages
c=======================================================================
c      Initialize zonal mean diagnostics
c=======================================================================
      use cpflame_module
#ifdef enable_diag_tracer
      use tracer_module
#endif
      use zonal_averages_module
      implicit none
#include "netcdf.inc"
      integer :: i,j,k,js,je,n
      integer :: ncid,iret
      integer :: z_tdim,z_udim,itimedim,varid
      integer :: lat_tdim,lat_udim, dims(4)
      character :: name*64, unit*32
c      real          :: fac=1e-5*0.2/rho_0*g  ! hor. shear of initial condition

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

#ifdef enable_diag_tracer
      nr_tr = nt
#endif

      if (my_pe==0) then
       print*,' found ',nr_tr,' tracer '
       print*,''
      endif

      allocate( um(jmt,km) ); um=0.0
      allocate( vm(jmt,km) ); vm=0.0
      allocate( wm(jmt,km) ); wm=0.0
      allocate( bm(jmt,km) ); bm=0.0
      allocate( pm(jmt,km) ); pm=0.0
      allocate( tm(jmt,km,nr_tr) ); tm=0.0

#ifdef diag_eddy_fluxes
      allocate( vb2(jmt,km,1:nr_var) ); vb2=0.0
      allocate( wb2(jmt,km,1:nr_var) ); wb2=0.0
      allocate( vt2(jmt,km,nr_tr,1:nr_var) ); vt2=0.0
      allocate( wt2(jmt,km,nr_tr,1:nr_var) ); wt2=0.0
#endif

#ifdef diag_energy_budget
      allocate( eke(jmt,km) ); eke=0.0
      allocate( epe(jmt,km) ); epe=0.0
      allocate( mke(jmt,km) ); mke=0.0
      allocate( mpe(jmt,km) ); mpe=0.0
      allocate( bws(jmt,km) ); bws=0.0
      allocate( bwq(jmt,km) ); bwq=0.0
      allocate( T2 (jmt,km) ); T2=0.0
      allocate( T4u(jmt,km) ); T4u=0.0
      allocate( T4v(jmt,km) ); T4v=0.0
      allocate( T4(jmt,km) );  T4=0.0

      allocate( eps_b(jmt,km) ); eps_b=0.0
      allocate( eps_u(jmt,km) ); eps_u=0.0
      allocate( eps_v(jmt,km) ); eps_v=0.0
      allocate( eps_kh(jmt,km)); eps_kh=0.0
      allocate( eps_kv(jmt,km)); eps_kv=0.0
      allocate( rest(jmt,km) );  rest=0.0
      allocate( restm(jmt,km) ); restm=0.0

      allocate( vps(jmt,km) ); vps=0.0
      allocate( wps(jmt,km) ); wps=0.0
      allocate( veke(jmt,km) ); veke=0.0
      allocate( weke(jmt,km) ); weke=0.0
      allocate( vepe(jmt,km) ); vepe=0.0
      allocate( wepe(jmt,km) ); wepe=0.0
#endif

#ifdef diag_eddy_flux_tendency
      allocate( vv(jmt,km)  ); vv  =0.0
      allocate( vvb(jmt,km) ); vvb =0.0
      allocate( ageo(jmt,km)); ageo=0.0
      allocate( ageo2(jmt,km));ageo2=0.0
      allocate( vw(jmt,km)  ); vw  =0.0
      allocate( uv(jmt,km)  ); uv  =0.0
      allocate( vvu(jmt,km)  ); vvu  =0.0
      allocate( agev1(jmt,km));agev1=0.0
      allocate( agev2(jmt,km));agev2=0.0
#endif

#ifdef diag_potential_vorticity
      allocate( mpv_qg(jmt,km)  ); mpv_qg=0.0
      allocate( mpv_er(jmt,km)  ); mpv_er=0.0
      allocate( mrv_er(jmt,km)  ); mrv_er=0.0
      allocate( vq_qg(jmt,km,1:nr_var)  );  vq_qg=0.0
      allocate( vq_er(jmt,km,1:nr_var)  );  vq_er=0.0
      allocate( wq_qg(jmt,km,1:nr_var)  );  wq_qg=0.0
      allocate( wq_er(jmt,km,1:nr_var)  );  wq_er=0.0
      allocate( vvq_er(jmt,km)  ); vvq_er=0.0
      allocate( vvq_qg(jmt,km)  ); vvq_qg=0.0
#endif

      allocate( ulen(jmt,km) ); ulen=0.0
      allocate( vlen(jmt,km) ); vlen=0.0
      allocate( wlen(jmt,km) ); wlen=0.0
      allocate( blen(jmt,km) ); blen=0.0

      js=max(2,js_pe); je = min(je_pe,jmt-1)
      do j=js,je
       do i=2,imt-1
        ulen(j,:)=ulen(j,:)+dx*maskU(i,j,:)
        vlen(j,:)=vlen(j,:)+dx*maskV(i,j,:)
        wlen(j,:)=wlen(j,:)+dx*maskW(i,j,:)
        blen(j,:)=blen(j,:)+dx*maskT(i,j,:)
       enddo 
      enddo 
c       Restoring time scale for buffer zone
      allocate( t_rest(jmt)); t_rest=0.
      do j=2,3+1
        t_rest(j)= 1./(5*85400)/(j-1.)
      enddo
      do j=jmt-1,jmt-3,-1
        t_rest(j)= 1./(5*85400)/(-1.*(j-jmt))
      enddo

      if (my_pe==0) then 

      call def_grid_cdf('zonal_ave.cdf')
      iret=nf_open('zonal_ave.cdf',NF_WRITE,ncid)
      iret=nf_set_fill(ncid, NF_NOFILL, iret)
      call ncredf(ncid, iret)
      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)
      iret=nf_inq_dimid(ncid,'Time',itimedim)

      dims = (/lat_tdim, z_tdim, iTimedim,1/)
      varid = ncvdef (ncid,'um', NCFLOAT,3,dims,iret)
      name = 'Zonal mean zonal velocity'; unit = 'm/s'
      call dvcdf(ncid,varid,name,64,unit,32,spval)

      dims = (/lat_udim, z_tdim, iTimedim,1/)
      varid = ncvdef (ncid,'vm', NCFLOAT,3,dims,iret)
      name = 'Zonal mean meridional velocity'; unit = 'm/s'
      call dvcdf(ncid,varid,name,64,unit,32,spval)

      dims = (/lat_tdim, z_udim, iTimedim,1/)
      varid = ncvdef (ncid,'wm', NCFLOAT,3,dims,iret)
      name = 'Zonal mean vertical velocity'; unit = 'm/s'
      call dvcdf(ncid,varid,name,64,unit,32,spval)

      dims = (/lat_tdim, z_tdim, iTimedim,1/)
      varid = ncvdef (ncid,'bm', NCFLOAT,3,dims,iret)
      name = 'Zonal mean buoyancy'; unit = 'm/s**2'
      call dvcdf(ncid,varid,name,64,unit,32,spval)

      dims = (/lat_tdim, z_tdim, iTimedim,1/)
      varid = ncvdef (ncid,'pm', NCFLOAT,3,dims,iret)
      name = 'Zonal mean pressure'; unit = 'm**2/s**2'
      call dvcdf(ncid,varid,name,64,unit,32,spval)

      do n=1,nr_tr
       dims = (/lat_tdim, z_tdim, iTimedim,1/)
       write(name, '("trm",i2)') n; call replace_space_zero(name)
       varid = ncvdef (ncid,name, NCFLOAT,3,dims,iret)
       name = 'zonal mean tracer'; unit = ' '
       call dvcdf(ncid,varid,name,64,unit,32,spval)
      enddo

#ifdef diag_eddy_fluxes
      do k=1,nr_var
       write(unit, '("m**",i2,"/s**",i2)')k+1,2*k+1
       dims = (/lat_udim, z_tdim, iTimedim,1/)
       write(name, '("vb_",i2)') k; call replace_space_zero(name)
       varid = ncvdef (ncid,name, NCFLOAT,3,dims,iret)
       write(name, '("Meridional buoyancy ",i2,". moment flux")')k
       call dvcdf(ncid,varid,name,64,unit,32,spval)

       dims = (/lat_tdim, z_udim, iTimedim,1/)
       write(name, '("wb_",i2)') k; call replace_space_zero(name)
       varid = ncvdef (ncid,name, NCFLOAT,3,dims,iret)
       write(name, '("Vertical buoyancy ",i2,". moment flux")')k
       call dvcdf(ncid,varid,name,64,unit,32,spval)

       do n=1,nr_tr
        unit = 'm/s'
        dims = (/lat_udim, z_tdim, iTimedim,1/)
        write(name, '("vt",i2,"_",i2)')n,k;call replace_space_zero(name)
        varid = ncvdef (ncid,name, NCFLOAT,3,dims,iret)
        write(name, '("Meridional ",i2,". moment tracer flux")')k
        call dvcdf(ncid,varid,name,64,unit,32,spval)

        dims = (/lat_tdim, z_udim, iTimedim,1/)
        write(name, '("wt",i2,"_",i2)')n,k;call replace_space_zero(name)
        varid = ncvdef (ncid,name, NCFLOAT,3,dims,iret)
        write(name, '("Vertical ",i2,". moment tracer flux")')k
        call dvcdf(ncid,varid,name,64,unit,32,spval)
       enddo
      enddo
#endif

#ifdef diag_energy_budget
      dims = (/lat_tdim, z_tdim, iTimedim,1/)
      varid = ncvdef (ncid,'eke', NCFLOAT,3,dims,iret)
      name = 'Eddy kinetic energy'; unit = 'm**2/s**2'
      call dvcdf(ncid,varid,name,64,unit,32,spval)

      dims = (/lat_tdim, z_tdim, iTimedim,1/)
      varid = ncvdef (ncid,'mke', NCFLOAT,3,dims,iret)
      name = 'Mean kinetic energy'; unit = 'm**2/s**2'
      call dvcdf(ncid,varid,name,64,unit,32,spval)

      dims = (/lat_tdim, z_tdim, iTimedim,1/)
      varid = ncvdef (ncid,'epe', NCFLOAT,3,dims,iret)
      name = 'Eddy potential energy'; unit = 'm**2/s**2'
      call dvcdf(ncid,varid,name,64,unit,32,spval)

      dims = (/lat_tdim, z_tdim, iTimedim,1/)
      varid = ncvdef (ncid,'mpe', NCFLOAT,3,dims,iret)
      name = 'Mean potential energy'; unit = 'm**2/s**2'
      call dvcdf(ncid,varid,name,64,unit,32,spval)

      dims = (/lat_tdim, z_tdim, iTimedim,1/)
      varid = ncvdef (ncid,'bws', NCFLOAT,3,dims,iret)
      name = 'EKE production (EPE->EKE)'; unit = 'm**2/s**3'
      call dvcdf(ncid,varid,name,64,unit,32,spval)

      dims = (/lat_tdim, z_tdim, iTimedim,1/)
      varid = ncvdef (ncid,'bwq', NCFLOAT,3,dims,iret)
      name = 'MPE production (MPE->MKE)'; unit = 'm**2/s**3'
      call dvcdf(ncid,varid,name,64,unit,32,spval)

      dims = (/lat_tdim, z_tdim, iTimedim,1/)
      varid = ncvdef (ncid,'T2', NCFLOAT,3,dims,iret)
      name = 'MPE production (EPE->MPE)'; unit = 'm**2/s**3'
      call dvcdf(ncid,varid,name,64,unit,32,spval)

      dims = (/lat_tdim, z_tdim, iTimedim,1/)
      varid = ncvdef (ncid,'T4', NCFLOAT,3,dims,iret)
      name = 'MKE production (EKE->MKE)'; unit = 'm**2/s**3'
      call dvcdf(ncid,varid,name,64,unit,32,spval)

      dims = (/lat_tdim, z_tdim, iTimedim,1/)
      varid = ncvdef (ncid,'eps_b', NCFLOAT,3,dims,iret)
      name = 'EPE dissipation'; unit = 'm**2/s**3'
      call dvcdf(ncid,varid,name,64,unit,32,spval)

      dims = (/lat_tdim, z_tdim, iTimedim,1/)
      varid = ncvdef (ncid,'eps_kh', NCFLOAT,3,dims,iret)
      name = 'Horizontal EKE dissipation'; unit = 'm**2/s**3'
      call dvcdf(ncid,varid,name,64,unit,32,spval)

      dims = (/lat_tdim, z_tdim, iTimedim,1/)
      varid = ncvdef (ncid,'eps_kv', NCFLOAT,3,dims,iret)
      name = 'Vertical EKE dissipation'; unit = 'm**2/s**3'
      call dvcdf(ncid,varid,name,64,unit,32,spval)

      dims = (/lat_tdim, z_tdim, iTimedim,1/)
      varid = ncvdef (ncid,'rest', NCFLOAT,3,dims,iret)
      name = 'EPE dissipation by nudging'; unit = 'm**2/s**3'
      call dvcdf(ncid,varid,name,64,unit,32,spval)

      dims = (/lat_tdim, z_tdim, iTimedim,1/)
      varid = ncvdef (ncid,'restm', NCFLOAT,3,dims,iret)
      name = 'MPE dissipation by nudging'; unit = 'm**2/s**3'
      call dvcdf(ncid,varid,name,64,unit,32,spval)

      dims = (/lat_udim, z_tdim, iTimedim,1/)
      varid = ncvdef (ncid,'vps', NCFLOAT,3,dims,iret)
      name = 'Meridional wave radiation'; unit = 'm**3/s**3'
      call dvcdf(ncid,varid,name,64,unit,32,spval)

      dims = (/lat_tdim, z_udim, iTimedim,1/)
      varid = ncvdef (ncid,'wps', NCFLOAT,3,dims,iret)
      name = 'Vertical wave radiation'; unit = 'm**3/s**3'
      call dvcdf(ncid,varid,name,64,unit,32,spval)

      dims = (/lat_udim, z_tdim, iTimedim,1/)
      varid = ncvdef (ncid,'veke', NCFLOAT,3,dims,iret)
      name = 'Meridional EKE advection'; unit = 'm**3/s**3'
      call dvcdf(ncid,varid,name,64,unit,32,spval)

      dims = (/lat_tdim, z_udim, iTimedim,1/)
      varid = ncvdef (ncid,'weke', NCFLOAT,3,dims,iret)
      name = 'Vertical EKE advection'; unit = 'm**3/s**3'
      call dvcdf(ncid,varid,name,64,unit,32,spval)

      dims = (/lat_udim, z_tdim, iTimedim,1/)
      varid = ncvdef (ncid,'vepe', NCFLOAT,3,dims,iret)
      name = 'Meridional EPE advection'; unit = 'm**3/s**3'
      call dvcdf(ncid,varid,name,64,unit,32,spval)

      dims = (/lat_tdim, z_udim, iTimedim,1/)
      varid = ncvdef (ncid,'wepe', NCFLOAT,3,dims,iret)
      name = 'Vertical EPE advection'; unit = 'm**3/s**3'
      call dvcdf(ncid,varid,name,64,unit,32,spval)
#endif

#ifdef diag_eddy_flux_tendency
      dims = (/lat_udim, z_tdim, iTimedim,1/)
      varid = ncvdef (ncid,'vv', NCFLOAT,3,dims,iret)
      name = 'Meridional velocity variance'; unit = 'm^2/s^2'
      call dvcdf(ncid,varid,name,64,unit,32,spval)

      dims = (/lat_udim, z_tdim, iTimedim,1/)
      varid = ncvdef (ncid,'vvb', NCFLOAT,3,dims,iret)
      name = 'Triple variance flux'; unit = 'm^3/s^4'
      call dvcdf(ncid,varid,name,64,unit,32,spval)

      dims = (/lat_udim, z_tdim, iTimedim,1/)
      varid = ncvdef (ncid,'ageo', NCFLOAT,3,dims,iret)
      name = 'Ageostrophic terms'; unit = 'm^2/s^4'
      call dvcdf(ncid,varid,name,64,unit,32,spval)

      dims = (/lat_udim, z_tdim, iTimedim,1/)
      varid = ncvdef (ncid,'ageo2', NCFLOAT,3,dims,iret)
      name = 'Ageostrophic (beta) term'; unit = 'm^2/s^4'
      call dvcdf(ncid,varid,name,64,unit,32,spval)

      dims = (/lat_tdim, z_udim, iTimedim,1/)
      varid = ncvdef (ncid,'vw', NCFLOAT,3,dims,iret)
      name = 'Velocity correlation'; unit = 'm^2/s^4'
      call dvcdf(ncid,varid,name,64,unit,32,spval)

      dims = (/lat_udim, z_tdim, iTimedim,1/)
      varid = ncvdef (ncid,'uv', NCFLOAT,3,dims,iret)
      name = 'Velocity correlation'; unit = 'm^2/s^2'
      call dvcdf(ncid,varid,name,64,unit,32,spval)

      dims = (/lat_udim, z_tdim, iTimedim,1/)
      varid = ncvdef (ncid,'vvu', NCFLOAT,3,dims,iret)
      name = 'Flux of velocity correlation'; unit = 'm^3/s^3'
      call dvcdf(ncid,varid,name,64,unit,32,spval)

      dims = (/lat_udim, z_tdim, iTimedim,1/)
      varid = ncvdef (ncid,'agev1', NCFLOAT,3,dims,iret)
      name = 'Ageostrophic term 1'; unit = 'm^2/s^3'
      call dvcdf(ncid,varid,name,64,unit,32,spval)

      dims = (/lat_udim, z_tdim, iTimedim,1/)
      varid = ncvdef (ncid,'agev2', NCFLOAT,3,dims,iret)
      name = 'Ageostrophic term 2'; unit = 'm^2/s^3'
      call dvcdf(ncid,varid,name,64,unit,32,spval)

#endif

#ifdef diag_potential_vorticity
      dims = (/lat_tdim, z_tdim, iTimedim,1/)
      varid = ncvdef (ncid,'mpv_qg', NCFLOAT,3,dims,iret)
      name = 'Potential vorticity (QG)'; unit = '1/s'
      call dvcdf(ncid,varid,name,64,unit,32,spval)

      dims = (/lat_udim, z_tdim, iTimedim,1/)
      varid = ncvdef (ncid,'vvq_qg', NCFLOAT,3,dims,iret)
      name = 'Triple var. flux of PV (QG)'; unit = 'm/s^2'
      call dvcdf(ncid,varid,name,64,unit,32,spval)

      dims = (/lat_tdim, z_tdim, iTimedim,1/)
      varid = ncvdef (ncid,'mpv_er', NCFLOAT,3,dims,iret)
      name = 'Potential vorticity (Ertel)'; unit = '1/s'
      call dvcdf(ncid,varid,name,64,unit,32,spval)

      dims = (/lat_udim, z_tdim, iTimedim,1/)
      varid = ncvdef (ncid,'mrv_er', NCFLOAT,3,dims,iret)
      name = 'Relative vorticity (Ertel)'; unit = '1/s'
      call dvcdf(ncid,varid,name,64,unit,32,spval)

      dims = (/lat_udim, z_tdim, iTimedim,1/)
      varid = ncvdef (ncid,'vvq_er', NCFLOAT,3,dims,iret)
      name = 'Triple var. flux of PV (Ertel)'; unit = 'm^2/s^3'
      call dvcdf(ncid,varid,name,64,unit,32,spval)

      do k=1,nr_var
       write(unit, '("m**",i2,"/s**",i2)')k,k+1

       dims = (/lat_udim, z_tdim, iTimedim,1/)
       write(name, '("vq_qg_",i2)') k; call replace_space_zero(name)
       varid = ncvdef (ncid,name, NCFLOAT,3,dims,iret)
       write(name, '("Meridional PV (QG) ",i2,". moment flux")')k
       call dvcdf(ncid,varid,name,64,unit,32,spval)

       dims = (/lat_udim, z_tdim, iTimedim,1/)
       write(name, '("vq_er_",i2)') k; call replace_space_zero(name)
       varid = ncvdef (ncid,name, NCFLOAT,3,dims,iret)
       write(name, '("Meridional PV (Ertel) ",i2,". moment flux")')k
       call dvcdf(ncid,varid,name,64,unit,32,spval)

       dims = (/lat_tdim, z_udim, iTimedim,1/)
       write(name, '("wq_qg_",i2)') k; call replace_space_zero(name)
       varid = ncvdef (ncid,name, NCFLOAT,3,dims,iret)
       write(name, '("Vertical PV (QG) ",i2,". moment flux")')k
       call dvcdf(ncid,varid,name,64,unit,32,spval)

       dims = (/lat_tdim, z_udim, iTimedim,1/)
       write(name, '("wq_er_",i2)') k; call replace_space_zero(name)
       varid = ncvdef (ncid,name, NCFLOAT,3,dims,iret)
       write(name, '("Vertical PV (Ertel) ",i2,". moment flux")')k
       call dvcdf(ncid,varid,name,64,unit,32,spval)

      enddo

#endif
      varid = ncvdef (ncid,'b_r', NCFLOAT,1,z_tdim,iret)
      name = 'Reference buoyancy profile'; unit = 'm/s^2'
      call dvcdf(ncid,varid,name,64,unit,32,spval)

      call ncendf(ncid, iret)
      iret= nf_put_vara_double(ncid,varid,1,km,b_r)

      call ncclos (ncid, iret)
      endif ! my_pe ==0
      end subroutine init_zonal_averages



      subroutine zonal_averages_write
c=======================================================================
c       compute zonal mean diagnostics
c=======================================================================
      use cpflame_module
      use zonal_averages_module
      implicit none
#include "netcdf.inc"
      integer :: i,j,k,m
      integer :: ncid,iret,n,npe, corner(4), edges(4)
      integer :: itdimid,ilen,itimeid
      integer :: umid,vmid,wmid,bmid,trmid(nr_tr),varid
      integer :: vbid,wbid,vtid(nr_tr),wtid(nr_tr)
      real :: fxa
      type(time_type) :: time
      character :: name*64, unit*32
c-----------------------------------------------------------------------
c      write to netcdf file
c-----------------------------------------------------------------------
      do npe=0,n_pes
       if (my_pe==npe) then

         iret=nf_open('zonal_ave.cdf',NF_WRITE,ncid)
         iret=nf_set_fill(ncid, NF_NOFILL, iret)
         do n=1,nr_tr
          write(name, '("trm",i2)') n; call replace_space_zero(name)
          iret=nf_inq_varid(ncid,name,trmid(n))
         enddo

         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

         Corner = (/js_pe,1,ilen,1/); 
         edges  = (/je_pe-js_pe+1,km,1,1/)
         iret=nf_inq_varid(ncid,'um',varid)
         iret= nf_put_vara_double(ncid,varid,corner,edges, 
     &          um(js_pe:je_pe,:))
         iret=nf_inq_varid(ncid,'vm',varid)
         iret= nf_put_vara_double(ncid,varid,corner,edges,
     &          vm(js_pe:je_pe,:))
         iret=nf_inq_varid(ncid,'wm',varid)
         iret= nf_put_vara_double(ncid,varid,corner,edges,
     &          wm(js_pe:je_pe,:))
         iret=nf_inq_varid(ncid,'bm',varid)
         iret= nf_put_vara_double(ncid,varid,corner,edges,
     &          bm(js_pe:je_pe,:))
         iret=nf_inq_varid(ncid,'pm',varid)
         iret= nf_put_vara_double(ncid,varid,corner,edges,
     &          pm(js_pe:je_pe,:))
         do n=1,nr_tr
          iret= nf_put_vara_double(ncid,trmid(n),corner,edges,
     &          tm(js_pe:je_pe,:,n))
         enddo

#ifdef diag_eddy_fluxes
         do m=1,nr_var
          write(name, '("vb_",i2)')m; call replace_space_zero(name)
          iret=nf_inq_varid(ncid,name,varid)
          iret= nf_put_vara_double(ncid,varid,corner,edges,
     &          vb2(js_pe:je_pe,:,m))
          write(name, '("wb_",i2)')m; call replace_space_zero(name)
          iret=nf_inq_varid(ncid,name,varid)
          iret= nf_put_vara_double(ncid,varid,corner,edges,
     &          wb2(js_pe:je_pe,:,m))
          do n=1,nr_tr
           write(name, '("vt",i2,"_",i2)')n,m;
           call replace_space_zero(name)
           iret=nf_inq_varid(ncid,name,varid)
           iret= nf_put_vara_double(ncid,varid,corner,edges,
     &          vt2(js_pe:je_pe,:,n,m))
           write(name, '("wt",i2,"_",i2)')n,m;
           call replace_space_zero(name)
           iret=nf_inq_varid(ncid,name,varid)
           iret= nf_put_vara_double(ncid,varid,corner,edges,
     &          wt2(js_pe:je_pe,:,n,m))
          enddo
         enddo
#endif

#ifdef diag_energy_budget
         iret=nf_inq_varid(ncid,'eke',varid)
         iret=nf_put_vara_double(ncid,varid,corner,edges,
     &          eke(js_pe:je_pe,:))
         iret=nf_inq_varid(ncid,'mke',varid)
         iret=nf_put_vara_double(ncid,varid,corner,edges,
     &          mke(js_pe:je_pe,:))
         iret=nf_inq_varid(ncid,'epe',varid)
         iret=nf_put_vara_double(ncid,varid,corner,edges,
     &          epe(js_pe:je_pe,:))
         iret=nf_inq_varid(ncid,'mpe',varid)
         iret=nf_put_vara_double(ncid,varid,corner,edges,
     &          mpe(js_pe:je_pe,:))
         iret=nf_inq_varid(ncid,'bws',varid)
         iret=nf_put_vara_double(ncid,varid,corner,edges,
     &          bws(js_pe:je_pe,:))
         iret=nf_inq_varid(ncid,'bwq',varid)
         iret=nf_put_vara_double(ncid,varid,corner,edges,
     &          bwq(js_pe:je_pe,:))
         iret=nf_inq_varid(ncid,'T2',varid)
         iret=nf_put_vara_double(ncid,varid,corner,edges,
     &          T2(js_pe:je_pe,:))
         iret=nf_inq_varid(ncid,'T4',varid)
         iret=nf_put_vara_double(ncid,varid,corner,edges,
     &          T4u(js_pe:je_pe,:))
         iret=nf_inq_varid(ncid,'T4v',varid)
         iret=nf_put_vara_double(ncid,varid,corner,edges,
     &          T4v(js_pe:je_pe,:))
         iret=nf_inq_varid(ncid,'eps_b',varid)
         iret=nf_put_vara_double(ncid,varid,corner,edges,
     &          eps_b(js_pe:je_pe,:))
         iret=nf_inq_varid(ncid,'eps_kh',varid)
         iret=nf_put_vara_double(ncid,varid,corner,edges,
     &          eps_kh(js_pe:je_pe,:))
         iret=nf_inq_varid(ncid,'eps_kv',varid)
         iret=nf_put_vara_double(ncid,varid,corner,edges,
     &          eps_kv(js_pe:je_pe,:))
         iret=nf_inq_varid(ncid,'rest',varid)
         iret=nf_put_vara_double(ncid,varid,corner,edges,
     &          rest(js_pe:je_pe,:))
         iret=nf_inq_varid(ncid,'restm',varid)
         iret=nf_put_vara_double(ncid,varid,corner,edges,
     &          restm(js_pe:je_pe,:))
         iret=nf_inq_varid(ncid,'vps',varid)
         iret=nf_put_vara_double(ncid,varid,corner,edges,
     &          vps(js_pe:je_pe,:))
         iret=nf_inq_varid(ncid,'wps',varid)
         iret=nf_put_vara_double(ncid,varid,corner,edges,
     &          wps(js_pe:je_pe,:))
         iret=nf_inq_varid(ncid,'veke',varid)
         iret=nf_put_vara_double(ncid,varid,corner,edges,
     &          veke(js_pe:je_pe,:))
         iret=nf_inq_varid(ncid,'weke',varid)
         iret=nf_put_vara_double(ncid,varid,corner,edges,
     &          weke(js_pe:je_pe,:))
         iret=nf_inq_varid(ncid,'vepe',varid)
         iret=nf_put_vara_double(ncid,varid,corner,edges,
     &          vepe(js_pe:je_pe,:))
         iret=nf_inq_varid(ncid,'wepe',varid)
         iret=nf_put_vara_double(ncid,varid,corner,edges,
     &          wepe(js_pe:je_pe,:))
#endif

#ifdef diag_eddy_flux_tendency
         iret=nf_inq_varid(ncid,'vv',varid)
         iret=nf_put_vara_double(ncid,varid,corner,edges,
     &          vv(js_pe:je_pe,:))
         iret=nf_inq_varid(ncid,'vvb',varid)
         iret=nf_put_vara_double(ncid,varid,corner,edges,
     &          vvb(js_pe:je_pe,:))
         iret=nf_inq_varid(ncid,'vw',varid)
         iret=nf_put_vara_double(ncid,varid,corner,edges,
     &          vw(js_pe:je_pe,:))
         iret=nf_inq_varid(ncid,'ageo',varid)
         iret=nf_put_vara_double(ncid,varid,corner,edges,
     &          ageo(js_pe:je_pe,:))
         iret=nf_inq_varid(ncid,'ageo2',varid)
         iret=nf_put_vara_double(ncid,varid,corner,edges,
     &          ageo2(js_pe:je_pe,:))
         iret=nf_inq_varid(ncid,'uv',varid)
         iret=nf_put_vara_double(ncid,varid,corner,edges,
     &          uv(js_pe:je_pe,:))
         iret=nf_inq_varid(ncid,'vvu',varid)
         iret=nf_put_vara_double(ncid,varid,corner,edges,
     &          vvu(js_pe:je_pe,:))
         iret=nf_inq_varid(ncid,'agev1',varid)
         iret=nf_put_vara_double(ncid,varid,corner,edges,
     &          agev1(js_pe:je_pe,:))
         iret=nf_inq_varid(ncid,'agev2',varid)
         iret=nf_put_vara_double(ncid,varid,corner,edges,
     &          agev2(js_pe:je_pe,:))
#endif

#ifdef diag_potential_vorticity
         iret=nf_inq_varid(ncid,'mpv_qg',varid)
         iret=nf_put_vara_double(ncid,varid,corner,edges,
     &          mpv_qg(js_pe:je_pe,:))

         iret=nf_inq_varid(ncid,'mpv_er',varid)
         iret=nf_put_vara_double(ncid,varid,corner,edges,
     &          mpv_er(js_pe:je_pe,:))

         iret=nf_inq_varid(ncid,'mrv_er',varid)
         iret=nf_put_vara_double(ncid,varid,corner,edges,
     &          mrv_er(js_pe:je_pe,:))

         iret=nf_inq_varid(ncid,'vvq_er',varid)
         iret=nf_put_vara_double(ncid,varid,corner,edges,
     &          vvq_er(js_pe:je_pe,:))

         iret=nf_inq_varid(ncid,'vvq_qg',varid)
         iret=nf_put_vara_double(ncid,varid,corner,edges,
     &          vvq_qg(js_pe:je_pe,:))

         do m=1,nr_var

          write(name, '("vq_qg_",i2)')m; call replace_space_zero(name)
          iret=nf_inq_varid(ncid,name,varid)
          iret=nf_put_vara_double(ncid,varid,corner,edges,
     &          vq_qg(js_pe:je_pe,:,m))

          write(name, '("wq_qg_",i2)')m; call replace_space_zero(name)
          iret=nf_inq_varid(ncid,name,varid)
          iret=nf_put_vara_double(ncid,varid,corner,edges,
     &          wq_qg(js_pe:je_pe,:,m))

          write(name, '("vq_er_",i2)')m; call replace_space_zero(name)
          iret=nf_inq_varid(ncid,name,varid)
          iret=nf_put_vara_double(ncid,varid,corner,edges,
     &           vq_er(js_pe:je_pe,:,m))

          write(name, '("wq_er_",i2)')m; call replace_space_zero(name)
          iret=nf_inq_varid(ncid,name,varid)
          iret=nf_put_vara_double(ncid,varid,corner,edges,
     &           wq_er(js_pe:je_pe,:,m))

         enddo
#endif

         call ncclos (ncid, iret)
       endif
       call barrier
      enddo
      end subroutine zonal_averages_write





      subroutine zonal_averages_means
c=======================================================================
c       get zonal means
c=======================================================================
      use cpflame_module
      use zonal_averages_module
#ifdef enable_diag_tracer
      use tracer_module
#endif
      implicit none
      integer :: i,j,k,js,je,n

      js=max(2,js_pe); je = min(je_pe,jmt-1)
      um=0.0;vm=0.0;wm=0.0;bm=0.0;tm=0.0
      do j=js,je 
        do i=2,imt-1
         um(j,:)=um(j,:)+u(i,j,:,1,tau)*dx*maskU(i,j,:) 
         vm(j,:)=vm(j,:)+u(i,j,:,2,tau)*dx*maskV(i,j,:) 
         wm(j,:)=wm(j,:)+u(i,j,:,3,tau)*dx*maskW(i,j,:) 
         bm(j,:)=bm(j,:)+b(i,j,:,  tau)*dx*maskT(i,j,:) 
         pm(j,:)=pm(j,:)+p_full(i,j,:,tau)*dx*maskT(i,j,:) 
#ifdef enable_diag_tracer
         do n=1,nr_tr
          tm(j,:,n)=tm(j,:,n)+tr(i,j,:,tau,n)*dx*maskT(i,j,:) 
         enddo
#endif
        enddo 
      enddo 
      where (ulen/=0); um=um/ulen; else where; um=spval ; end where
      where (vlen/=0); vm=vm/vlen; else where; vm=spval ; end where
      where (wlen/=0); wm=wm/wlen; else where; wm=spval ; end where
      where (blen/=0); bm=bm/blen; else where; bm=spval ; end where
      where (blen/=0); pm=pm/blen; else where; pm=spval ; end where
      call border_exchg_merid(um,1)
      call border_exchg_merid(vm,1)
      call border_exchg_merid(wm,1)
      call border_exchg_merid(bm,1)
      call border_exchg_merid(pm,1)
      do n=1,nr_tr
        where (blen/=0); tm(:,:,n)=tm(:,:,n)/blen
        else where; tm(:,:,n)=spval;end where
        call border_exchg_merid(tm(:,:,n),1)
      enddo
      end subroutine zonal_averages_means





      subroutine zonal_averages_eddy_fluxes
#ifdef diag_eddy_fluxes
c=======================================================================
c       compute zonal mean eddy fluxes
c=======================================================================
      use cpflame_module
#ifdef enable_diag_tracer
      use tracer_module
#endif
      use zonal_averages_module
      implicit none
      integer :: i,j,k,js,je,m,n

      js=max(2,js_pe); je = min(je_pe,jmt-1)
      vb2=0.0;wb2=0.0;vt2=0.0;wt2=0.0
      do m=1,nr_var
        do j=js,je 
         do k=1,km
          do i=2,imt-1
           vb2(j,k,m)=vb2(j,k,m)+u(i,j,k,2,tau)*dx*maskV(i,j,k) 
     &     *.5/m*( (b(i,j  ,k,tau)-bm(j  ,k))**m
     &           + (b(i,j+1,k,tau)-bm(j+1,k))**m  )
          enddo
         enddo
         do k=1,km-1
          do i=2,imt-1
           wb2(j,k,m)=wb2(j,k,m)+u(i,j,k,3,tau)*dx*maskW(i,j,k) 
     &      *.5/m*( (b(i,j,k  ,tau)-bm(j,k  ))**m
     &            + (b(i,j,k+1,tau)-bm(j,k+1))**m )
          enddo
         enddo
       enddo
       where (vlen/=0); vb2(:,:,m)=vb2(:,:,m)/vlen; 
       else where; vb2(:,:,m)=spval; end where
       where (wlen/=0); wb2(:,:,m)=wb2(:,:,m)/wlen; 
       else where; wb2(:,:,m)=spval; end where

#ifdef enable_diag_tracer
       do n=1,nr_tr
         do j=js,je 
          do k=1,km
           do i=2,imt-1
            vt2(j,k,n,m)=vt2(j,k,n,m)+u(i,j,k,2,tau)*dx*maskV(i,j,k) 
     &      *.5/m*( (tr(i,j  ,k,tau,n)-tm(j  ,k,n))**m 
     &            + (tr(i,j+1,k,tau,n)-tm(j+1,k,n))**m )
           enddo
          enddo
          do k=1,km-1
           do i=2,imt-1
            wt2(j,k,n,m)=wt2(j,k,n,m)+u(i,j,k,3,tau)*dx*maskW(i,j,k) 
     &      *.5/m*( (tr(i,j,k  ,tau,n)-tm(j,k  ,n))**m
     &            + (tr(i,j,k+1,tau,n)-tm(j,k+1,n))**m )
           enddo
          enddo
         enddo ! j
         where (vlen/=0); vt2(:,:,n,m)=vt2(:,:,n,m)/vlen; 
         else where; vt2(:,:,n,m)=spval; end where
         where (wlen/=0); wt2(:,:,n,m)=wt2(:,:,n,m)/wlen; 
         else where; wt2(:,:,n,m)=spval; end where
       enddo ! n
#endif
      enddo ! m
#endif
      end subroutine zonal_averages_eddy_fluxes





      subroutine zonal_averages_pv_fluxes
#ifdef diag_potential_vorticity
c=======================================================================
c       compute zonal mean potential vorticity
c=======================================================================
      use cpflame_module
      use zonal_averages_module
      implicit none
      integer :: i,j,k,js,je,m
      real   :: pv_qg(imt,jmt,km),f0,y0,ps(imt,jmt,km)
      real   :: rv_qg(imt,jmt,km),pv_er(imt,jmt,km)
      real   :: rv(imt,jmt,km),fxa

      f0 = coriolis_t(jmt/2); y0 = yt(jmt/2)
      js=max(2,js_pe); je = min(je_pe,jmt-1)
c
c       relative potential vorticity
c
      rv=0.0
      do j=js,je 
       do k=2,km-1
        do i=2,imt-1
          rv(i,j,k)=
     &       (u(i+1,j,k,2,tau)-u(i,j,k,2,tau) )/dx*maskU(i,j,k)
     &      -(u(i,j+1,k,1,tau)-u(i,j,k,1,tau) )/dx*maskV(i,j,k)
        enddo 
       enddo 
      enddo 
      call border_exchg3D(rv,1)
c
c       Ertel potential vorticity  Q N0^2 = b_z (v_x-u_y + f) + u_z b_y - v_z b_x
c
      pv_er=0.
      do j=js,je 
       do k=2,km-1
        do i=2,imt-1
         fxa= (rv(i,j,k)+rv(i-1,j,k)+rv(i,j-1,k)+rv(i-1,j-1,k))/4.0
         pv_er(i,j,k)=(fxa+coriolis_t(j)  )*
     &    ((b(i,j,k+1,tau)-b(i,j,k,tau))/dz*maskW(i,j,k)
     &    +(b(i,j,k,tau)-b(i,j,k-1,tau))/dz*maskW(i,j,k-1))/
     &     (maskW(i,j,k)+maskW(i,j,k-1)+epsln) 
         pv_er(i,j,k)=-N_0**(-2)*pv_er(i,j,k)
        enddo 
       enddo 
      enddo 
      call border_exchg3D(pv_er,1); call setcyclic3D(pv_er)
c
c       QG streamfunction
c
      psi=0.0
      do k=1,km
       do j=js,je 
        ps(:,j,k)=(p_full(:,j,k,tau)-p_r(k))/f0
       enddo
      enddo
      call border_exchg3D(ps,1); call setcyclic3D(ps)
c
c       QG relative potential vorticity
c
      rv_qg=0.
      do j=js,je 
       do k=2,km-1
        do i=2,imt-1
         rv_qg(i,j,k)=
     &    ((ps(i,j+1,k)-ps(i,j,k))/dx*maskV(i,j  ,k)
     &    -(ps(i,j,k)-ps(i,j-1,k))/dx*maskV(i,j-1,k))/dx
     &   +((ps(i+1,j,k)-ps(i,j,k))/dx*maskU(i-1,j,k)
     &    -(ps(i,j,k)-ps(i-1,j,k))/dx*maskU(i-1,j,k))/dx
        enddo 
       enddo 
      enddo 
      call border_exchg3D(rv_qg,1); call setcyclic3D(rv_qg)
c
c       add QG planetary and stretching term to relative vorticity
c
      pv_qg=0.
      do j=js,je 
       pv_qg(:,j,:)=rv_qg(:,j,:)+beta*(yt(j)-y0)
       do k=2,km-1
         pv_qg(:,j,k)=pv_qg(:,j,k)+f0**2/N_0**2*
     &    ((ps(:,j,k+1)-ps(:,j,k))/dz*maskW(:,j,k)
     &    -(ps(:,j,k)-ps(:,j,k-1))/dz*maskW(:,j,k-1))/dz
       enddo 
       k=2
       pv_qg(:,j,k)=pv_qg(:,j,k)+f0**2/N_0**2*
     &    ((ps(:,j,k+1)-ps(:,j,k))/dz*maskW(:,j,k)
     &    -(b(:,j,k,tau)-bm(j,k))*maskT(i,j,k) )/dz
       k=km-1
       pv_qg(:,j,k)=pv_qg(:,j,k)+f0**2/N_0**2*
     &    ((b(:,j,k,tau)-bm(j,k))*maskT(:,j,k)
     &    -(ps(:,j,k)-ps(:,j,k-1))/dz*maskW(:,j,k-1))/dz
      enddo 
      call border_exchg3D(pv_qg,1); call setcyclic3D(pv_qg)
c
c       mean QG and Ertel PV
c
      mpv_qg=0.
      do j=js,je 
        do i=2,imt-1
         mpv_qg(j,:)=mpv_qg(j,:)+pv_qg(i,j,:)*dx*maskT(i,j,:) 
         mpv_er(j,:)=mpv_er(j,:)+pv_er(i,j,:)*dx*maskT(i,j,:) 
         mrv_er(j,:)=mrv_er(j,:)+rv(i,j,:)*dx*maskT(i,j,:) 
        enddo 
      enddo 
      where (blen/=0); mpv_qg=mpv_qg/blen; 
      else where; mpv_qg=spval ; end where
      where (blen/=0); mpv_er=mpv_er/blen; 
      else where; mpv_er=spval ; end where
      where (vlen/=0); mrv_er=mrv_er/vlen; 
      else where; mrv_er=spval ; end where
      call border_exchg_merid(mpv_qg,1)
      call border_exchg_merid(mpv_er,1)
      call border_exchg_merid(mrv_er,1)
c

c       meridional and vertical fluxes of PV
c
      vq_qg=0.;vq_er=0.; wq_qg=0.;wq_er=0.; 
      do m=1,nr_var
        do j=js,je 
         do k=1,km
          do i=2,imt-1
           vq_qg(j,k,m)=vq_qg(j,k,m)+u(i,j,k,2,tau)*dx*maskV(i,j,k) 
     &     *.5/m*( (pv_qg(i,j  ,k)-mpv_qg(j  ,k))**m
     &           + (pv_qg(i,j+1,k)-mpv_qg(j+1,k))**m  )
           vq_er(j,k,m)=vq_er(j,k,m)+u(i,j,k,2,tau)*dx*maskV(i,j,k) 
     &     *.5/m*( (pv_er(i,j  ,k)-mpv_er(j  ,k))**m
     &           + (pv_er(i,j+1,k)-mpv_er(j+1,k))**m  )
          enddo
         enddo
         do k=1,km-1
          do i=2,imt-1
           wq_qg(j,k,m)=wq_qg(j,k,m)+u(i,j,k,3,tau)*dx*maskW(i,j,k) 
     &      *.5/m*( (pv_qg(i,j,k  )-mpv_qg(j,k  ))**m
     &            + (pv_qg(i,j,k+1)-mpv_qg(j,k+1))**m )
           wq_er(j,k,m)=wq_er(j,k,m)+u(i,j,k,3,tau)*dx*maskW(i,j,k) 
     &      *.5/m*( (pv_er(i,j,k  )-mpv_er(j,k  ))**m
     &            + (pv_er(i,j,k+1)-mpv_er(j,k+1))**m )
          enddo
         enddo
       enddo
       where (vlen/=0); vq_qg(:,:,m)=vq_qg(:,:,m)/vlen; 
       else where; vq_qg(:,:,m)=spval; end where
       where (wlen/=0); wq_qg(:,:,m)=wq_qg(:,:,m)/wlen; 
       else where; wq_qg(:,:,m)=spval; end where
       where (vlen/=0); vq_er(:,:,m)=vq_er(:,:,m)/vlen; 
       else where; vq_er(:,:,m)=spval; end where
       where (wlen/=0); wq_er(:,:,m)=wq_er(:,:,m)/wlen; 
       else where; wq_er(:,:,m)=spval; end where
      enddo
c
c      triple flux of PV (for eddy pv flux tendency)
c
      vvq_qg=0.;vvq_er=0.
      do j=js,je 
       do k=1,km
        do i=2,imt-1
         vvq_qg(j,k)=vvq_qg(j,k)
     &     +(u(i,j,k,2,tau)-vm(j,k))**2*dx*maskV(i,j,k) 
     &     *((pv_qg(i,j  ,k)-mpv_qg(j  ,k))
     &      +(pv_qg(i,j+1,k)-mpv_qg(j+1,k)))/2.
         vvq_er(j,k)=vvq_er(j,k)
     &     +(u(i,j,k,2,tau)-vm(j,k))**2*dx*maskV(i,j,k) 
     &     *((pv_er(i,j  ,k)-mpv_er(j  ,k))
     &      +(pv_er(i,j+1,k)-mpv_er(j+1,k)))/2.
        enddo
       enddo
      enddo
      where (vlen/=0); vvq_qg(:,:)=vvq_qg(:,:)/vlen; 
      else where; vvq_qg(:,:)=spval; end where
      where (vlen/=0); vvq_er(:,:)=vvq_er(:,:)/vlen; 
      else where; vvq_er(:,:)=spval; end where
#endif
      end subroutine zonal_averages_pv_fluxes



      subroutine zonal_averages_flux_tendency
#ifdef diag_eddy_flux_tendency
c=======================================================================
c        eddy flux balance :  (v'b')_t = - v'^2 b_y + ...
c=======================================================================
      use cpflame_module
      use zonal_averages_module
      implicit none
      integer :: i,j,k,js,je
      real :: fxa,a(imt,jmt,km)

      js=max(2,js_pe); je = min(je_pe,jmt-1)
      vv=0.;vvb=0.;vw=0.0; ageo=0.0; ageo2=0.
      do j=js,je 
        do k=2,km-1
         do i=2,imt-1
          vv(j,k)=vv(j,k)+maskV(i,j,k)*dx*(u(i,j,k,2,tau)-vm(j,k))**2
          vvb(j,k)=vvb(j,k)+maskV(i,j,k)*dx*
     &           (u(i,j,k,2,tau)-vm(j,k))**2
     &        *0.5*( b(i,j,k,tau)-bm(j,k)+b(i,j+1,k,tau)-bm(j+1,k))
          vw(j,k)=vw(j,k)+maskW(i,j,k)* N_0**2*dx*
     &        0.5*( (u(i,j  ,k,2,tau)-vm(j  ,k))*maskV(i,j  ,k)
     &             +(u(i,j-1,k,2,tau)-vm(j-1,k))*maskV(i,j-1,k))
     &           *(u(i,j,k,3,tau)-wm(j,k))
          ageo(j,k)=ageo(j,k)+maskV(i,j,k)*dx*
     &        0.5*( b(i,j,k,tau)-bm(j,k)+b(i,j+1,k,tau)-bm(j+1,k))
     &        *( 
     &           (p_full(i,j+1,k,tau)-pm(j+1,k)
     &           -p_full(i,j  ,k,tau)+pm(j  ,k))/dx
     &    +(coriolis_t(j    )*(u(i-1,j  ,k,1,tau)-um(j,k)
     &                      +u(i  ,j  ,k,1,tau)-um(j,k)) 
     &     +coriolis_t(j+1  )*(u(i-1,j+1,k,1,tau)-um(j+1,k)
     &                      +u(i  ,j+1,k,1,tau)-um(j+1,k)))/4.0
     &          )
          ageo2(j,k)=ageo2(j,k)+maskV(i,j,k)*dx*
     &        0.5*( b(i,j,k,tau)-bm(j,k)+b(i,j+1,k,tau)-bm(j+1,k))
     &        *( 
     &    +(beta*(yt(j  )-yt(jmt/2))*(u(i-1,j  ,k,1,tau)-um(j,k)
     &                      +u(i  ,j  ,k,1,tau)-um(j,k)) 
     &     +beta*(yt(j+1)-yt(jmt/2))*(u(i-1,j+1,k,1,tau)-um(j+1,k)
     &                      +u(i  ,j+1,k,1,tau)-um(j+1,k)))/4.0
     &          )
         enddo 
        enddo 
      enddo 
      where (vlen/=0); vv=vv/vlen; else where; vv=spval; end where
      where (vlen/=0); vvb=vvb/vlen; else where; vvb=spval; end where
      where (wlen/=0); vw=vw/wlen; else where; vw=spval; end where
      where (vlen/=0); ageo=ageo/vlen;else where;ageo=spval; end where
      where (vlen/=0);ageo2=ageo2/vlen;else where;ageo2=spval;end where

c
c        tendency of cross correlation u'v' 
c
      uv=0.;vvu=0.
      do j=js,je 
       do k=2,km-1
        do i=2,imt-1
         fxa=0.5*(u(i,j,k,1,tau)-um(j,k)+u(i,j+1,k,1,tau)-um(j+1,k) )
     &      *0.5*(u(i,j,k,2,tau)-vm(j,k)+u(i+1,j,k,2,tau)-vm(j,k) )
         uv(j,k)=uv(j,k)+maskV(i,j,k)*dx*fxa
         vvu(j,k)=vvu(j,k)+maskV(i,j,k)*dx*fxa
     &      *0.5*(u(i,j,k,2,tau)+u(i+1,j,k,2,tau))
        enddo 
       enddo 
      enddo 
      where (vlen/=0);uv=uv/vlen;else where;uv=spval;end where
      where (vlen/=0);vvu=vvu/vlen;else where;vvu=spval;end where

      do j=js,je 
       do k=2,km-1
        do i=2,imt-1
          a(i,j,k)= maskU(i,j,k)*( 
     &     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
     &              -(p_full(i+1,j,k,tau)-p_full(i,j,k,tau))/dx )
        enddo 
       enddo 
      enddo 
      call border_exchg3D(a,1); call setcyclic3D(a)
      agev1=0.
      do j=js,je 
       do k=2,km-1
        do i=2,imt-1
         agev1(j,k)=agev1(j,k)+maskV(i,j,k)*dx*
     &    0.5*(a(i,j,k)+a(i,j+1,k))*
     &    0.5*(u(i,j,k,2,tau)-vm(j,k)+u(i+1,j,k,2,tau)-vm(j,k))
        enddo 
       enddo 
      enddo 
      where (vlen/=0);agev1=agev1/vlen;else where;agev1=spval;end where

      do j=js,je 
       do k=2,km-1
        do i=2,imt-1
          a(i,j,k)= maskV(i,j,k)*( 
     &    -(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
     &    -(p_full(i,j+1,k,tau)-p_full(i,j,k,tau))/dx) 
        enddo 
       enddo 
      enddo 
      call border_exchg3D(a,1); call setcyclic3D(a)
      agev2=0.
      do j=js,je 
       do k=2,km-1
        do i=2,imt-1
         agev2(j,k)=agev2(j,k)+maskU(i,j,k)*dx*
     &    0.5*(a(i,j,k)+a(i+1,j,k))*
     &    0.5*(u(i,j,k,1,tau)-um(j,k)+u(i,j+1,k,1,tau)-um(j+1,k))
        enddo 
       enddo 
      enddo 
      where (vlen/=0);agev2=agev2/vlen;else where;agev2=spval;end where

#endif
      end subroutine zonal_averages_flux_tendency







      subroutine zonal_averages_energy
#ifdef diag_energy_budget
c=======================================================================
c       compute zonal mean energy budget
c=======================================================================
      use cpflame_module
      use zonal_averages_module
      implicit none

      integer :: i,j,k,js,je,m,n
      real   :: fxa,a(jmt,km),bs(imt,jmt,km),bms(jmt,km)
      real   :: bss(imt,jmt,km),del2(imt,jmt,km),us(imt,jmt,km)

      js=max(2,js_pe); je = min(je_pe,jmt-1)
      bs=0.0;bms=0.0;bss=0.  ! buoyancy deviation from reference state
      do j=js,je 
        do k=2,km-1
         do i=2,imt-1
           bs(i,j,k)=(b(i,j,k,tau)-b_r(k))*maskT(i,j,k)
         enddo 
        enddo 
      enddo 
      call border_exchg3D(bs,1); call setcyclic3D(bs)
      do j=js,je 
        do i=2,imt-1
         bms(j,:)=bms(j,:)+bs(i,j,:)*dx*maskT(i,j,:) 
        enddo 
      enddo 
      where (blen/=0) bms=bms/blen
      call border_exchg_merid(bms,1)
      do j=js,je 
        do k=2,km-1
         do i=2,imt-1
           bss(i,j,k)=(bs(i,j,k)-bms(j,k))*maskT(i,j,k)
         enddo 
        enddo 
      enddo 
      call border_exchg3D(bss,1); call setcyclic3D(bss)
c-----------------------------------------------------------------------
c       Mean energy
c-----------------------------------------------------------------------
      where (vlen==0); vm=0; end where
      where (wlen==0); wm=0; end where
      call border_exchg_merid(vm,1)
      call border_exchg_merid(wm,1)

      do j=js,je 
        do k=2,km-1
         mke(j,k)=um(j,k)**2/2+(0.5*(vm(j,k)+vm(j-1,k)))**2/2.
         mpe(j,k)=bms(j,k)**2/2/N_0**2
         bwq(j,k)= 0.25*wm(j,k  )*(bms(j,k  )+bms(j,k+1))
     &            +0.25*wm(j,k-1)*(bms(j,k-1)+bms(j,k  ))
        enddo
      enddo
      where (blen==0); mke=spval; end where
      where (blen==0); mpe=spval; end where
      where (blen==0); bwq=spval; end where

      where (vlen==0); vm=spval; end where
      where (wlen==0); wm=spval; end where
      call border_exchg_merid(vm,1)
      call border_exchg_merid(wm,1)
c-----------------------------------------------------------------------
c        Eddy energy
c-----------------------------------------------------------------------
      eke=0.0;bws=0.0;epe=0.0;
      do j=js,je 
        do k=2,km-1
         do i=2,imt-1
          eke(j,k)=eke(j,k)+dx*( 
     &      ( 0.5*((u(i  ,j,k,1,tau)-um(j  ,k))*maskU(i  ,j,k)
     &            +(u(i-1,j,k,1,tau)-um(j  ,k))*maskU(i-1,j,k) ) )**2/2
     &     +( 0.5*((u(i,j-1,k,2,tau)-vm(j-1,k))*maskV(i,j-1,k)
     &            +(u(i,j-1,k,2,tau)-vm(j  ,k))*maskV(i,j  ,k) ) )**2/2
     &                          )
          epe(j,k)=epe(j,k)+dx*bss(i,j,k)**2/2/N_0**2

          bws(j,k)=bws(j,k)+dx*( 
     &      (u(i,j,k  ,3,tau)-wm(j,k  ))*0.25*maskW(i,j,k  )*
     &         (bss(i,j,k  )+bss(i,j,k+1))
     &     +(u(i,j,k-1,3,tau)-wm(j,k-1))*0.25*maskW(i,j,k-1)*
     &         (bss(i,j,k-1)+bss(i,j,k  ))
     &                          )
         enddo 
        enddo 
      enddo 
      where (blen/=0); eke=eke/blen; else where; eke=spval; end where
      where (blen/=0); epe=epe/blen; else where; epe=spval; end where
      where (blen/=0); bws=bws/blen; else where; bws=spval; end where
c-----------------------------------------------------------------------
c        Baroclinic instability term
c-----------------------------------------------------------------------
      a=0.0 ! v'b' bar b_y /N_0**2 
      do j=js,je 
        do k=2,km-1
         do i=2,imt-1
          a(j,k)=a(j,k)+dx/N_0**2*( 
     &          (u(i,j,k,2,tau)-vm(j,k))*
     &      0.5*(bss(i,j,k)+bss(i,j+1,k))*(bms(j+1,k)-bms(j,k))/dx 
     &                          )*maskV(i,j,k)
         enddo 
        enddo 
      enddo 
      where (vlen/=0) a=a/vlen
      call border_exchg_merid(a,1)
      do j=js,je 
        do k=2,km-1
         T2(j,k)=0.5*( a(j-1,k)+a(j,k))
        enddo 
      enddo 

      a=0.0 ! w'b' bar b_z /N_0**2 
      do j=js,je 
        do k=2,km-1
         do i=2,imt-1
          a(j,k)=a(j,k)+dx/N_0**2*( 
     &          (u(i,j,k,3,tau)-wm(j,k))*
     &      0.5*(bss(i,j,k)+bss(i,j,k+1))*( bms(j,k+1)-bms(j,k))/dz 
     &                          )*maskW(i,j,k)
         enddo 
        enddo 
      enddo 
      where (wlen/=0) a=a/wlen
      do j=js,je 
        do k=2,km-1
         T2(j,k)=T2(j,k)+0.5*(a(j,k-1)+a(j,k))
        enddo 
      enddo 
      where (blen==0); T2 =spval; end where

c-----------------------------------------------------------------------
c       Shear instability terms
c-----------------------------------------------------------------------
      a=0.0; 
      do j=js,je 
        do k=2,km-1
         do i=2,imt-1 ! bar(u' v') bar u_y    
          a(j,k)=a(j,k)+dx*( 
     &        0.5*( u(i,j,k,1,tau)-um(j,k) + u(i,j+1,k,1,tau)-um(j+1,k))
     &       *0.5*( u(i,j,k,2,tau)-vm(j,k) + u(i+1,j,k,2,tau)-vm(j,k))
     &       *(um(j+1,k)-um(j,k))/dx
     &                          )*maskU(i,j,k)*maskU(i,j+1,k)
         enddo 
        enddo 
      enddo 
      where (ulen/=0) a=a/ulen
      call border_exchg_merid(a,1)
      do j=js,je 
        do k=2,km-1
         T4u(j,k)=0.5*( a(j-1,k)+a(j,k))
        enddo 
      enddo 

      a=0.0; 
      do j=js,je 
        do k=2,km-1
         do i=2,imt-1 ! bar(u' w') bar u_z
          a(j,k)=a(j,k)+dx*( 
     &        0.5*( u(i,j,k,1,tau)-um(j,k) + u(i,j,k+1,1,tau)-um(j,k+1))
     &       *0.5*( u(i,j,k,3,tau)-wm(j,k) + u(i+1,j,k,3,tau)-wm(j,k))
     &       *(um(j,k+1)-um(j,k))/dz
     &                          )*maskU(i,j,k)*maskU(i,j,k+1)
         enddo 
        enddo 
      enddo 
      where (ulen/=0) a=a/ulen
      do j=js,je 
        do k=2,km-1
         T4u(j,k)=T4u(j,k)+0.5*( a(j,k)+a(j,k-1))
        enddo 
      enddo 
c       where (blen==0); T4u=spval; end where

      a=0.0; 
      do j=js,je 
        do k=2,km-1
         do i=2,imt-1 ! bar(v' v') bar v_y
          a(j,k)=a(j,k)+dx*( 
     &        0.5*( u(i,j,k,2,tau)-vm(j,k) + u(i,j-1,k,2,tau)-vm(j-1,k))
     &       *0.5*( u(i,j,k,2,tau)-vm(j,k) + u(i,j-1,k,2,tau)-vm(j-1,k))
     &       *(vm(j,k)-vm(j-1,k))/dx
     &                          )*maskV(i,j,k)*maskV(i,j-1,k)
         enddo 
        enddo 
      enddo 
      where (vlen/=0) a=a/vlen
      do j=js,je 
        do k=2,km-1
         T4v(j,k)=a(j,k)
        enddo 
      enddo 

      a=0.0; 
      do j=js,je 
        do k=2,km-1
         do i=2,imt-1 ! bar(v' w') bar v_z
          a(j,k)=a(j,k)+dx*( 
     &        0.5*( u(i,j,k,2,tau)-vm(j,k) + u(i,j,k+1,2,tau)-vm(j,k+1))
     &       *0.5*( u(i,j,k,3,tau)-wm(j,k) + u(i,j+1,k,3,tau)-wm(j+1,k))
     &       *(vm(j,k+1)-vm(j,k))/dz
     &                          )*maskV(i,j,k)*maskV(i,j,k+1)
         enddo 
        enddo 
      enddo 
      where (vlen/=0) a=a/vlen
      call border_exchg_merid(a,1)
      do j=js,je 
        do k=2,km-1
         T4v(j,k)=T4v(j,k)+(a(j,k)+a(j,k-1)+a(j-1,k)+a(j-1,k-1))/4.0
        enddo 
      enddo 
      T4u=T4u+T4v
      where (blen==0); T4u=spval; end where
      where (blen==0); T4v=spval; end where
c-----------------------------------------------------------------------
c        Dissipation of EKE:  M'_u**2   M'_v**2
c-----------------------------------------------------------------------
      eps_kh = 0.
      if (enable_biharmonic_friction) then
        do j=js,je 
         do k=2,km-1
          do i=2,imt-1
            us(i,j,k)=(u(i,j,k,1,tau)-um(j,k))*maskU(i,j,k)
          enddo 
         enddo 
        enddo 
        call border_exchg3D(us,1); call setcyclic3D(us)
        call biha_diss_u(us,del2)
        eps_u=0.0;
        do j=js,je 
         do k=2,km-1
          do i=2,imt-1
           eps_u(j,k)=eps_u(j,k)+dx*del2(i,j,k)**2*maskU(i,j,k)
          enddo 
         enddo 
        enddo 
        do j=js,je 
         do k=2,km-1
          do i=2,imt-1
            us(i,j,k)=(u(i,j,k,2,tau)-vm(j,k))*maskV(i,j,k)
          enddo 
         enddo 
        enddo 
        call border_exchg3D(us,1); call setcyclic3D(us)
        call biha_diss_v(us,del2)
        eps_v=0.0;
        do j=js,je 
         do k=2,km-1
          do i=2,imt-1
           eps_v(j,k)=eps_v(j,k)+dx*del2(i,j,k)**2*maskV(i,j,k)
          enddo 
         enddo 
        enddo 
        call border_exchg_merid(eps_v,1)
        eps_kh=0
        do j=js,je 
         do k=2,km-1
          eps_kh(j,k)=eps_u(j,k)+0.5*(eps_v(j,k)+eps_v(j-1,k))
         enddo 
        enddo 
      endif

c        add here dissipation by harmonic friction

      where (blen/=0); eps_kh=eps_kh/blen; 
      else where; eps_kh=spval; end where
c-----------------------------------------------------------------------
c        Dissipation of EKE:  A_v (u'_z)**2  +A_v (v'_z)**2 
c-----------------------------------------------------------------------
      do j=js,je 
        do k=2,km-1
         do i=2,imt-1
           us(i,j,k)=(u(i,j,k,1,tau)-um(j,k))*maskU(i,j,k)
         enddo 
        enddo 
      enddo 
      call border_exchg3D(us,1); call setcyclic3D(us)
      eps_u=0.0;
      do j=js,je 
        do k=2,km-1
         do i=2,imt-1
          eps_u(j,k)=eps_u(j,k)+dx*A_v*(
     &      0.5*(us(i,j,k+1)-us(i,j,k  ))/dz*maskW(i,j,k  ) 
     &     +0.5*(us(i,j,k  )-us(i,j,k-1))/dz*maskW(i,j,k-1) 
     &           )**2         
         enddo 
        enddo 
      enddo 
      do j=js,je 
        do k=2,km-1
         do i=2,imt-1
           us(i,j,k)=(u(i,j,k,2,tau)-vm(j,k))*maskV(i,j,k)
         enddo 
        enddo 
      enddo 
      call border_exchg3D(us,1); call setcyclic3D(us)
      eps_v=0.0;
      do j=js,je 
        do k=2,km-1
         do i=2,imt-1
          eps_v(j,k)=eps_v(j,k)+dx*A_v*(
     &      0.5*(us(i,j,k+1)-us(i,j,k  ))/dz*maskW(i,j,k  ) 
     &     +0.5*(us(i,j,k  )-us(i,j,k-1))/dz*maskW(i,j,k-1) 
     &           )**2         
         enddo 
        enddo 
      enddo 
      call border_exchg_merid(eps_v,1)
      eps_kv=0
      do j=js,je 
        do k=2,km-1
         eps_kv(j,k)=eps_u(j,k)+0.5*(eps_v(j,k)+eps_v(j-1,k))
        enddo 
      enddo 
c
c        Dissipation of EKE by bottom stress
c
      if (enable_bottom_stress) then
        eps_u=0.;eps_v=0.
        do j=js,je 
         do i=2,imt-1
          k=max(1,kmu(i,j))
          eps_u(j,k)=eps_u(j,k)
     &          +dx*cdbot*(u(i,j,k,1,tau)-um(j,k))**2*masku(i,j,k)
          k=max(1,kmv(i,j))
          eps_v(j,k)=eps_v(j,k)
     &          +dx*cdbot*(u(i,j,k,2,tau)-vm(j,k))**2*maskV(i,j,k)
         enddo
        enddo
        do j=js,je 
         do k=2,km-1
          eps_kv(j,k)=eps_kv(j,k)
     &             +eps_u(j,k)+0.5*(eps_v(j,k)+eps_v(j-1,k))
         enddo 
        enddo 
      endif
c
c        Dissipation of EKE by interior stress
c
      if (enable_interior_stress) then
        eps_u=0.;eps_v=0.
        do k=2,km-1
         do j=js,je 
          do i=2,imt-1
           eps_u(j,k)=eps_u(j,k)
     &           +dx*cdint*(u(i,j,k,1,tau)-um(j,k))**2*masku(i,j,k)
           eps_v(j,k)=eps_v(j,k)
     &           +dx*cdint*(u(i,j,k,2,tau)-vm(j,k))**2*maskV(i,j,k)
          enddo
         enddo
        enddo
        do k=2,km-1
         do j=js,je 
           eps_kv(j,k)=eps_kv(j,k)+
     &              eps_u(j,k)+0.5*(eps_v(j,k)+eps_v(j-1,k))
         enddo 
        enddo 
      endif

      where (blen/=0); eps_kv=eps_kv/blen; 
      else where; eps_kv=spval; end where
c-----------------------------------------------------------------------
c        Dissipation of EPE:  K_v (b'_z)**2 /N**2 +  M'_b**2
c-----------------------------------------------------------------------
      eps_b=0.0;
      do j=js,je 
        do k=2,km-1
         do i=2,imt-1
          eps_b(j,k)=eps_b(j,k)+dx/N_0**2*( 
     &     +0.5*(K_v+K_b(i,j,k  ))
     &       *(bss(i,j,k+1)-bss(i,j,k))/dz
     &       *(b(i,j,k+1,tau)-b(i,j,k,tau))/dz*maskW(i,j,k) 
     &     +0.5*(K_v+K_b(i,j,k-1))
     &       *(bss(i,j,k)-bss(i,j,k-1))/dz
     &       *(b(i,j,k,tau)-b(i,j,k-1,tau))/dz*maskW(i,j,k-1) 
c     &       (K_v+0.5*(K_b(i,j,k)+K_b(i,j,k-1)) )
c     &        *( (b(i,j,k+1,tau)-b(i,j,k-1,tau))/(2*dz) )**2
     &                          )
         enddo 
        enddo 
      enddo 
      if (enable_biharmonic_diffusion) then
        call biha_diss_b(bss,del2)
        do j=js,je 
         do k=2,km-1
          do i=2,imt-1
           eps_b(j,k)=eps_b(j,k)+dx/N_0**2*del2(i,j,k)**2*maskT(i,j,k)
          enddo 
         enddo 
        enddo 
      endif
      where (blen/=0); eps_b=eps_b/blen; 
      else where; eps_b=spval; end where
c-----------------------------------------------------------------------
c        Dissipation of EPE due to buffer zones
c-----------------------------------------------------------------------
      rest =0.0;
      do j=js,je 
        do k=2,km-1
         do i=2,imt-1
          rest (j,k)=rest (j,k)+dx/N_0**2*( 
     &       bss(i,j,k)*b(i,j,k,tau)*t_rest(j)*maskT(i,j,k)
     &                          )
         enddo 
        enddo 
      enddo 
      where (blen/=0); rest =rest /blen; 
      else where; rest =spval; end where

      restm =0.0;
      do j=js,je 
        do k=2,km-1
         do i=2,imt-1
          restm(j,k)=restm(j,k)+dx/N_0**2*( 
     &       bms(j,k)*b(i,j,k,tau)*t_rest(j)*maskT(i,j,k)
     &                          )
         enddo 
        enddo 
      enddo 
      where (blen/=0); restm=restm/blen; 
      else where; restm=spval; end where
c-----------------------------------------------------------------------
c        EKE radiation  vp wp
c-----------------------------------------------------------------------
      vps=0.0; wps=0.0;
      do j=js,je 
        do k=2,km-1
         do i=2,imt-1
          vps(j,k)=vps(j,k)+dx*( 
     &     ( u(i,j,k,2,tau)-vm(j,k) )*
     &       0.5*(p_full(i,j,k,tau)+p_full(i,j+1,k,tau))*maskV(i,j,k)
     &                          )
          wps(j,k)=wps(j,k)+dx*( 
     &     ( u(i,j,k,3,tau)-wm(j,k) )*
     &       0.5*(p_full(i,j,k,tau)+p_full(i,j,k+1,tau))*maskW(i,j,k)
     &                          )
         enddo 
        enddo 
      enddo 
      where (vlen/=0); vps  =vps  /vlen; 
      else where; vps  =spval; end where
      where (wlen/=0); wps  =wps  /wlen; 
      else where; wps  =spval; end where
c-----------------------------------------------------------------------
c        EKE advection
c-----------------------------------------------------------------------
      del2=0 
      do j=js,je 
        do k=2,km-1
         do i=2,imt-1
           del2(i,j,k)=
     &      ( 0.5*((u(i  ,j,k,1,tau)-um(j  ,k))*maskU(i  ,j,k)
     &            +(u(i-1,j,k,1,tau)-um(j  ,k))*maskU(i-1,j,k) ) )**2/2
     &     +( 0.5*((u(i,j-1,k,2,tau)-vm(j-1,k))*maskV(i,j-1,k)
     &            +(u(i,j-1,k,2,tau)-vm(j  ,k))*maskV(i,j  ,k) ) )**2/2
         enddo 
        enddo 
      enddo 
      call border_exchg3D(del2,1); call setcyclic3D(del2)

      veke=0.0; weke=0.0;
      do j=js,je 
        do k=2,km-1
         do i=2,imt-1
          veke(j,k)=veke(j,k)+dx*( 
     &     u(i,j,k,2,tau)*maskV(i,j,k)*
     &       0.5*(del2(i,j,k)+del2(i,j+1,k))*maskV(i,j,k)
     &                          )
          weke(j,k)=weke(j,k)+dx*( 
     &     u(i,j,k,3,tau)*maskW(i,j,k)*
     &       0.5*(del2(i,j,k)+del2(i,j,k+1))*maskW(i,j,k)
     &                          )
         enddo 
        enddo 
      enddo 
      where (vlen/=0); veke =veke /vlen; 
      else where; veke =spval; end where
      where (wlen/=0); weke =weke /wlen; 
      else where; weke =spval; end where
c-----------------------------------------------------------------------
c        EPE advection
c-----------------------------------------------------------------------
      del2=0 
      do j=js,je 
        do k=2,km-1
         do i=2,imt-1
           del2(i,j,k)= bss(i,j,k)**2/2/N_0**2
         enddo 
        enddo 
      enddo 
      call border_exchg3D(del2,1); call setcyclic3D(del2)

      vepe=0.0; wepe=0.0;
      do j=js,je 
        do k=2,km-1
         do i=2,imt-1
          vepe(j,k)=vepe(j,k)+dx*( 
     &     u(i,j,k,2,tau)*maskV(i,j,k)*
     &       0.5*(del2(i,j,k)+del2(i,j+1,k))*maskV(i,j,k)
     &                          )
          wepe(j,k)=wepe(j,k)+dx*( 
     &     u(i,j,k,3,tau)*maskW(i,j,k)*
     &       0.5*(del2(i,j,k)+del2(i,j,k+1))*maskW(i,j,k)
     &                          )
         enddo 
        enddo 
      enddo 
      where (vlen/=0); vepe =vepe /vlen; 
      else where; vepe =spval; end where
      where (wlen/=0); wepe =wepe /wlen; 
      else where; wepe =spval; end where
#endif
      end subroutine zonal_averages_energy








#ifdef diag_energy_budget

      subroutine biha_diss_b(var,del2)
      use cpflame_module
      implicit none
      integer :: i,j,k,js,je
      real :: var(imt,jmt,km) 
      real :: diff_fe(imt,jmt,km), diff_fn(imt,jmt,km) 
      real :: del2(imt,jmt,km),diff_tx,diff_ty,diffx
      DIFF_Tx(i,j,k)=(diff_fe(i,j,k)-diff_fe(i-1,j,k))/dx
      DIFF_Ty(i,j,k)=(diff_fn(i,j,k)-diff_fn(i,j-1,k))/dx

      js=max(2,js_pe); je = min(je_pe,jmt-1)
      diffx = sqrt(abs(Khbi)); diff_fe=0.0;diff_fn=0.0
      do k=1,km
       do j=js,je
        do i=1,imt-1
         diff_fe(i,j,k)=diffx*(var(i+1,j,k)-var(i,j,k))/dx
     &                  *maskU(i,j,k)
        enddo
       enddo
      enddo
      call border_exchg3D(diff_fe,1); call setcyclic3D(diff_fe)
      do k=1,km-1
       do j=js,je
        do i=1,imt
         diff_fn(i,j,k)=diffx*(var(i,j+1,k)-var(i,j,k))/dx
     &                  *maskV(i,j,k)
        enddo
       enddo
      enddo
      call border_exchg3D(diff_fn,1); call setcyclic3D(diff_fn)
      del2=0.0
      do k=2,km-1
       do j=js,je
        do i=2,imt-1
	  del2(i,j,k) = (DIFF_Tx(i,j,k) + DIFF_Ty(i,j,k))*maskT(i,j,k)
        enddo
       enddo
      enddo
      call border_exchg3D(del2,1); call setcyclic3D(del2)
      end subroutine biha_diss_b




      subroutine biha_diss_u (var,del2)
      use cpflame_module
      implicit none
      integer :: i,j,k,js,je
      real :: var(imt,jmt,km) 
      real :: diff_fe(imt,jmt,km), diff_fn(imt,jmt,km) 
      real :: del2(imt,jmt,km),diff_tx,diff_ty,diffx
      DIFF_Tx(i,j,k) = (diff_fe(i,j,k)-diff_fe(i-1,j,k))/dx
      DIFF_Ty(i,j,k) = (diff_fn(i,j,k)-diff_fn(i,j-1,k))/dx

      js=max(2,js_pe); je = min(je_pe,jmt-1)
      diffx = sqrt(abs(ahbi)); diff_fe=0.0;diff_fn=0.0
      do k=1,km
       do j=js,je
        do i=1,imt-1
         diff_fe(i,j,k)=diffx*(var(i+1,j,k)-var(i,j,k))/dx
        enddo
       enddo
      enddo
      call border_exchg3D(diff_fe,1); call setcyclic3D(diff_fe)
      do k=1,km-1
       do j=js,je
        do i=1,imt
         diff_fn(i,j,k)=diffx*(var(i,j+1,k)-var(i,j,k))/dx
     &                  *maskU(i,j+1,k)*maskU(i,j,k)
        enddo
       enddo
      enddo
      if (enable_noslip) then
       do j=js-1,je
         diff_fn(:,j,:)=diff_fn(:,j,:)-2*diffx*var(:,j,:)/dx
     &                   *(1-maskU(:,j+1,:))*maskU(:,j,:)
     &                                +2*diffx*var(:,j+1,:)/dx
     &                   *(1-maskU(:,j,:))*maskU(:,j+1,:)
       enddo
      endif
      call border_exchg3D(diff_fn,1); call setcyclic3D(diff_fn)
      del2=0.0
      do k=2,km-1
       do j=js,je
        do i=2,imt-1
	  del2(i,j,k) = (DIFF_Tx(i,j,k) + DIFF_Ty(i,j,k))*maskU(i,j,k)
        enddo
       enddo
      enddo
      call border_exchg3D(del2,1); call setcyclic3D(del2)
      end subroutine biha_diss_u




      subroutine biha_diss_v (var,del2)
      use cpflame_module
      implicit none
      integer :: i,j,k,js,je
      real :: var(imt,jmt,km) 
      real :: diff_fe(imt,jmt,km), diff_fn(imt,jmt,km) 
      real :: del2(imt,jmt,km),diff_tx,diff_ty,diffx
      DIFF_Tx(i,j,k)=(diff_fe(i,j,k)-diff_fe(i-1,j,k))/dx
      DIFF_Ty(i,j,k)=(diff_fn(i,j,k)-diff_fn(i,j-1,k))/dx

      js=max(2,js_pe); je = min(je_pe,jmt-1)
      diffx = sqrt(abs(ahbi)); diff_fe=0.0;diff_fn=0.0
      do k=1,km
       do j=js,je
        do i=1,imt-1
         diff_fe(i,j,k)=diffx*(var(i+1,j,k)-var(i,j,k))/dx
     &                  *maskV(i,j,k)*maskV(i+1,j,k)
        enddo
       enddo
      enddo
      if (enable_noslip) then
       do i=1,imt-1
         diff_fe(i,:,:)=diff_fe(i,:,:)-2*diffx*var(i,:,:)/dx
     &                   *(1-maskV(i+1,:,:))*maskV(i,:,:)
     &                                +2*diffx*var(i+1,:,:)/dx
     &                   *(1-maskV(i,:,:))*maskV(i+1,:,:)
       enddo
      endif
      call border_exchg3D(diff_fe,1); call setcyclic3D(diff_fe)
      do k=1,km-1
       do j=js,je
        do i=1,imt
         diff_fn(i,j,k)=diffx*(var(i,j+1,k)-var(i,j,k))/dx
        enddo
       enddo
      enddo
      call border_exchg3D(diff_fn,1); call setcyclic3D(diff_fn)
      del2=0.0
      do k=2,km-1
       do j=js,je
        do i=2,imt-1
	  del2(i,j,k) = (DIFF_Tx(i,j,k) + DIFF_Ty(i,j,k))*maskV(i,j,k)
        enddo
       enddo
      enddo
      call border_exchg3D(del2,1); call setcyclic3D(del2)
      end subroutine biha_diss_v
#endif



#else
      subroutine zonal_averages_dummy
      end
#endif
