

 module quasi_stokes_module
 implicit none
 real*8,allocatable :: b_scale(:),harea(:,:,:),htrans(:,:,:),hforc(:,:,:)
 real*8,allocatable :: varea(:,:,:),utrans(:,:,:),vtrans(:,:,:)
 real*8,allocatable :: umean(:,:,:),vmean(:,:,:),wmean(:,:,:),bmean(:,:,:)

 integer :: nb
 ! User input
 real*8 :: b_start=-0.01, b_end=0.08
 logical :: enable_quasi_stokes_ver = .true.
 logical :: enable_quasi_stokes_hor = .false.
 logical :: enable_quasi_stokes_mean = .false.
 ! end user input
 real*8,parameter :: epsln     = 1.0d-20
 integer :: counter=0, quasi_stokes_snapint 
 end module quasi_stokes_module



 subroutine init_quasi_stokes
 use pyOM_module 
 use quasi_stokes_module
 implicit none
 integer :: m

 nb=ny*1.3
 quasi_stokes_snapint = 86400.0 
 if (my_pe==0) then
   print*,' Initializing module for mean quasi stokes streamfunction'
   print*,' using ',nb,' buoyancy levels from ',b_start,' to ',b_end
   print*,' averaging each ',quasi_stokes_snapint,' sec'
   if (enable_quasi_stokes_hor) print*,' calculating horizontal quasi stokes streamfct.'
   if (enable_quasi_stokes_ver) print*,' calculating vertical quasi stokes streamfct.'
   if (enable_quasi_stokes_mean) print*,' calculating also Eulerian means'
 endif

 allocate( b_scale(nb) ); b_scale=0.0
 if (enable_quasi_stokes_hor) then
  allocate( harea(nx,nb,nz),htrans(nx,nb,nz),hforc(nx,nb,nz) )
  harea=0.0;htrans=0.0;hforc=0.0
 endif
 if (enable_quasi_stokes_ver) then
  allocate( varea(nx,ny,nb),utrans(nx,ny,nb),vtrans(nx,ny,nb) )
  varea=0.0;utrans=0.0;vtrans=0.0
 endif

 do m=1,nb
  b_scale(m) = b_start+(m-1)*(b_end-b_start)/(nb-1)
 enddo
 if (enable_quasi_stokes_mean) then
  allocate( umean(nx,ny,nz),vmean(nx,ny,nz),wmean(nx,ny,nz),bmean(nx,ny,nz) )
  umean=0.0;vmean=0.0;wmean=0.0;bmean=0.0
 endif
 counter = 0.0
 call quasi_stokes_read_restart
 if (my_pe==0) then
  print*,' done'
 endif
 end subroutine init_quasi_stokes


 subroutine quasi_stokes
 use pyOM_module 
 use fcontrol_module
 use quasi_stokes_module
 implicit none
 integer :: i,j,k,m,js,je
 real*8 :: bb,fxa,ww
 real*8 :: harea_(nx),htrans_(nx),hforc_(nx)

 if ( mod(itt,int(quasi_stokes_snapint/dt))  == 0)  then
   
  counter = counter + 1
  js=max(2,js_pe); je = min(je_pe,ny-1)

  if (enable_quasi_stokes_mean) then
   umean(:,js:je,:) = umean(:,js:je,:) + u(:,js:je,:,tau) 
   vmean(:,js:je,:) = vmean(:,js:je,:) + v(:,js:je,:,tau) 
   wmean(:,js:je,:) = wmean(:,js:je,:) + w(:,js:je,:,tau) 
   bmean(:,js:je,:) = bmean(:,js:je,:) + b(:,js:je,:,tau) 
  endif

  if (enable_quasi_stokes_hor) then
   do m=1,nb
    do k=2,nz-1
     htrans_=0.0;harea_=0.0;hforc_=0.0
     do j=js,je
      do i=2,nx-1
       bb=b(i,j,k,tau)
       ww=dx*maskT(i,j,k)
       if (bb < b_scale(m) ) ww=0
       harea_(i)  = harea_(i)+ww
       ww=dx*maskU(i,j,k)*u(i,j,k,tau)
       if (bb < b_scale(m) ) ww=0
       htrans_(i) = htrans_(i)+ww
       ww=(w(i,j,k,tau)+w(i,j,k-1,tau))/2.0*dx*maskT(i,j,k)
       if (bb < b_scale(m) ) ww=0
       hforc_(i)  = hforc_(i)+ww
      enddo
     enddo
     call global_sum_vec(harea_,nx )
     call global_sum_vec(htrans_,nx )
     call global_sum_vec(hforc_,nx ) 
     harea(:,m,k) =harea(:,m,k) + harea_
     htrans(:,m,k)=htrans(:,m,k)+ htrans_
     hforc(:,m,k) =hforc(:,m,k) + hforc_
    enddo
   enddo
  endif

  if (enable_quasi_stokes_ver) then
   do m=1,nb
    do k=2,nz-1
       do j=js,je
        do i=2,nx-1
         bb=b(i,j,k,tau)
         if (enable_back_state) bb=bb+back(i,j,k,tau)
         ww=(u(i-1,j,k,tau)*maskU(i-1,j,k)+u(i,j,k,tau)*maskU(i,j,k))/(maskU(i-1,j,k)+maskU(i,j,k)+epsln)
         if( bb<b_scale(m) ) ww=0
         utrans(i,j,m)=utrans(i,j,m)+ww*dz*maskT(i,j,k)
         ww=(v(i,j-1,k,tau)*maskV(i,j-1,k)+v(i,j,k,tau)*maskV(i,j,k))/(maskV(i,j-1,k)+maskV(i,j,k)+epsln)
         if( bb<b_scale(m) ) ww=0
         vtrans(i,j,m)=vtrans(i,j,m)+ww*dz*maskT(i,j,k)
         ww=dz*maskT(i,j,k)
         if( bb<b_scale(m) ) ww=0
         varea(i,j,m)=varea(i,j,m)+ww
        enddo
       enddo
    enddo
   enddo
  endif
 endif
 end subroutine quasi_stokes

 subroutine quasi_stokes_write_restart
!=======================================================================
!    write restart
!=======================================================================
 use pyOM_module 
 use quasi_stokes_module
 implicit none
 integer :: m,io,ierr,j

 if (enable_quasi_stokes_ver) then
  do m=1,nb
   call pe0_recv_2D(nx,ny,utrans(1,1,m))
   call pe0_recv_2D(nx,ny,vtrans(1,1,m))
   call pe0_recv_2D(nx,ny,varea(1,1,m))
  enddo
 endif

 if (enable_quasi_stokes_mean) then
  call pe0_recv_3D(nx,ny,nz,umean)
  call pe0_recv_3D(nx,ny,nz,vmean)
  call pe0_recv_3D(nx,ny,nz,wmean)
  call pe0_recv_3D(nx,ny,nz,bmean)
 endif

 if (my_pe==0) then
     print*,' writing restart file for mean quasi stokes streamfunction'
     call get_free_iounit(io,ierr)
     open(io,file='restart_over.dta', form='unformatted',status='unknown')
     write(io) nx,ny,nz,nb,b_start,b_end
     write(io) counter
     if (enable_quasi_stokes_mean) then
      do j=1,ny
       write(io) umean(:,j,:),vmean(:,j,:),wmean(:,j,:),bmean(:,j,:)
      enddo
     endif
     do m=1,nb
      if (enable_quasi_stokes_ver) write(io) varea(:,:,m),utrans(:,:,m),vtrans(:,:,m)
      if (enable_quasi_stokes_hor) write(io) harea(:,m,:),htrans(:,m,:),hforc(:,m,:)
     enddo
     close(io)
     print*,' done '
 endif
 end subroutine quasi_stokes_write_restart


 subroutine quasi_stokes_read_restart
!=======================================================================
!       read restart
!=======================================================================
 use pyOM_module 
 use quasi_stokes_module
 implicit none
 integer :: m,io,ierr,j
 integer :: nx_,ny_,nz_,nb_
 real*8 :: b_start_,b_end_
 if (my_pe==0) then
     print*,' reading restart file for mean quasi stokes streamfunction'
     call get_free_iounit(io,ierr)
     open(io,file='restart_over.dta', form='unformatted',status='old',err=10)
     read(io) nx_,ny_,nz_,nb_,b_start_,b_end_
     if (nx/=nx_ .or. ny/=ny_ .or. nz/= nz_ .or. nb/= nb_) then 
          print*,' read from restart dimensions: ',nx_,ny_,nz_,nb_
          print*,' does not match dimensions   : ',nx,ny,nz,nb
          goto 10
     endif
     if (b_start_ /= b_start .or. b_end_ /= b_end ) then
        print*,' read buoyancy scale from ',b_start_,' to ',b_end_
        print*,' does not match scale from ',b_start,' to ',b_end 
        goto 10
     endif
     read(io) counter
     print*,' read counter = ',counter
     if (enable_quasi_stokes_mean) then
      do j=1,ny
       read(io) umean(:,j,:),vmean(:,j,:),wmean(:,j,:),bmean(:,j,:)
      enddo
     endif
     do m=1,nb
      if (enable_quasi_stokes_ver) read(io) varea(:,:,m),utrans(:,:,m),vtrans(:,:,m)
      if (enable_quasi_stokes_hor) read(io) harea(:,m,:),htrans(:,m,:),hforc(:,m,:)
     enddo
     close(io)
     goto 20
 10  print*,' error while reading restart file '
 20  continue
 endif
 call bcast_integer(counter,1,0)
 if (enable_quasi_stokes_mean) then
  call pe0_send_3d(nx,ny,nz,umean)
  call pe0_send_3d(nx,ny,nz,vmean)
  call pe0_send_3d(nx,ny,nz,wmean)
  call pe0_send_3d(nx,ny,nz,bmean)
 endif
 if (enable_quasi_stokes_ver) then
  do m=1,nb
   call pe0_send_2d(nx,ny,varea(1,1,m))
   call pe0_send_2d(nx,ny,utrans(1,1,m))
   call pe0_send_2d(nx,ny,vtrans(1,1,m))
  enddo
 endif
 if (enable_quasi_stokes_hor) then
  m=nx*nb*nz
  call bcast_real(harea,m,0)
  call bcast_real(htrans,m,0)
  call bcast_real(hforc,m,0)
 endif
 end subroutine quasi_stokes_read_restart




 subroutine quasi_stokes_write_cdf
!=======================================================================
!      write results to netcdf file
!=======================================================================
 use pyOM_module 
 use fcontrol_module
 use quasi_stokes_module
 implicit none
 integer :: i,j,k,m,js,je
 real*8 :: time,fxa
 include "netcdf.inc"
 integer :: iret,ncid,tdim,tid,ilen,varid,nbdim,nbid
 integer :: corner(4), edges(4),dims(4)
 integer :: xtdim,xudim,ytdim,yudim,ztdim,zudim
 character :: name*24, unit*16
 real*8,parameter :: spval = -1.0d33

 if (enable_quasi_stokes_ver) then
  do m=1,nb
   call pe0_recv_2D(nx,ny,utrans(1,1,m))
   call pe0_recv_2D(nx,ny,vtrans(1,1,m))
   call pe0_recv_2D(nx,ny,varea(1,1,m))
  enddo
 endif
 if (enable_quasi_stokes_mean) then
  call pe0_recv_3D(nx,ny,nz,umean)
  call pe0_recv_3D(nx,ny,nz,vmean)
  call pe0_recv_3D(nx,ny,nz,wmean)
  call pe0_recv_3D(nx,ny,nz,bmean)
 endif

 if (my_pe==0) then 

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

  iret=nf_inq_dimid(ncid,'xt',xtdim)
  iret=nf_inq_dimid(ncid,'xu',xudim)
  iret=nf_inq_dimid(ncid,'yt',ytdim)
  iret=nf_inq_dimid(ncid,'yu',yudim)
  iret=nf_inq_dimid(ncid,'zt',ztdim)
  iret=nf_inq_dimid(ncid,'zu',zudim)
  iret=nf_inq_dimid(ncid,'Time',tdim)

  if (enable_quasi_stokes_hor) then
   dims = (/xtdim, nbdim,ztdim, tdim/)
   varid = ncvdef (ncid,'harea', NCFLOAT,4,dims,iret)
   name = 'Area'; unit = 'm'
   call dvcdf(ncid,varid,name,len_trim(name),unit,len_trim(unit),spval) 

   dims = (/xtdim, nbdim,ztdim, tdim/)
   varid = ncvdef (ncid,'hforc', NCFLOAT,4,dims,iret)
   name = 'Forcing'; unit = 'm^2/s'
   call dvcdf(ncid,varid,name,len_trim(name),unit,len_trim(unit),spval) 

   dims = (/xtdim, nbdim,ztdim, tdim/)
   varid = ncvdef (ncid,'htrans', NCFLOAT,4,dims,iret)
   name = 'Transport'; unit = 'm^2/s'
   call dvcdf(ncid,varid,name,len_trim(name),unit,len_trim(unit),spval) 
  endif

  if (enable_quasi_stokes_ver) then
   dims = (/xtdim, ytdim, nbdim, tdim/)
   varid = ncvdef (ncid,'utrans', NCFLOAT,4,dims,iret)
   name = 'Transport'; unit = 'm^2/s'
   call dvcdf(ncid,varid,name,len_trim(name),unit,len_trim(unit),spval) 

   dims = (/xtdim, ytdim, nbdim, tdim/)
   varid = ncvdef (ncid,'vtrans', NCFLOAT,4,dims,iret)
   name = 'Transport'; unit = 'm^2/s'
   call dvcdf(ncid,varid,name,len_trim(name),unit,len_trim(unit),spval) 

   dims = (/xtdim, ytdim, nbdim, tdim/)
   varid = ncvdef (ncid,'varea', NCFLOAT,4,dims,iret)
   name = 'Area'; unit = 'm'
   call dvcdf(ncid,varid,name,len_trim(name),unit,len_trim(unit),spval) 
  endif

  if (enable_quasi_stokes_mean) then
   dims = (/xudim, ytdim,ztdim, tdim/)
   varid = ncvdef (ncid,'um', NCFLOAT,4,dims,iret)
   name = 'Mean zonal velocity'; unit = 'm/s'
   call dvcdf(ncid,varid,name,len_trim(name),unit,len_trim(unit),spval) 

   dims = (/xtdim, yudim,ztdim, tdim/)
   varid = ncvdef (ncid,'vm', NCFLOAT,4,dims,iret)
   name = 'Mean meridional velocity'; unit = 'm/s'
   call dvcdf(ncid,varid,name,len_trim(name),unit,len_trim(unit),spval) 

   dims = (/xtdim, ytdim,zudim, tdim/)
   varid = ncvdef (ncid,'wm', NCFLOAT,4,dims,iret)
   name = 'Mean vertical velocity'; unit = 'm/s'
   call dvcdf(ncid,varid,name,len_trim(name),unit,len_trim(unit),spval) 

   dims = (/xtdim, ytdim,ztdim, tdim/)
   varid = ncvdef (ncid,'bm', NCFLOAT,4,dims,iret)
   name = 'Mean buoyancy'; unit = 'm/s^2'
   call dvcdf(ncid,varid,name,len_trim(name),unit,len_trim(unit),spval) 
  endif

  call ncendf(ncid, iret)
  iret= nf_put_vara_double(ncid,nbid,1,nb,b_scale)

  iret=nf_inq_dimid(ncid,'Time',tdim)
  iret=nf_inq_dimlen(ncid, tdim,ilen)
  iret=nf_inq_varid(ncid,'Time',tid)
  ilen=ilen+1
  time = itt*dt ! current_time-initial_time
  fxa = time/86400.0 ! time%days + time%seconds/86400.
  iret= nf_put_vara_double(ncid,tid,ilen,1,fxa)

  if (counter>0) then
   if (enable_quasi_stokes_hor) then
     harea=harea/counter
     htrans=htrans/counter
     hforc=hforc/counter
   endif
   if (enable_quasi_stokes_ver) then
     varea=varea/counter
     utrans=utrans/counter
     vtrans=vtrans/counter
   endif
   if (enable_quasi_stokes_mean) then
     umean=umean/counter
     vmean=vmean/counter
     wmean=wmean/counter
     bmean=bmean/counter
   endif
  endif

  Corner = (/1,1,1,ilen/); 
  if (enable_quasi_stokes_hor) then
   edges  = (/nx,nb,nz,1/)
   iret=nf_inq_varid(ncid,'harea',varid)
   iret= nf_put_vara_double(ncid,varid,corner,edges,harea)
   iret=nf_inq_varid(ncid,'htrans',varid)
   iret= nf_put_vara_double(ncid,varid,corner,edges,htrans)
   iret=nf_inq_varid(ncid,'hforc',varid)
   iret= nf_put_vara_double(ncid,varid,corner,edges,hforc)
  endif

  if (enable_quasi_stokes_ver) then
   edges  = (/nx,ny,nb,1/)
   iret=nf_inq_varid(ncid,'varea',varid)
   iret= nf_put_vara_double(ncid,varid,corner,edges,varea)
   iret=nf_inq_varid(ncid,'utrans',varid)
   iret= nf_put_vara_double(ncid,varid,corner,edges,utrans)
   iret=nf_inq_varid(ncid,'vtrans',varid)
   iret= nf_put_vara_double(ncid,varid,corner,edges,vtrans)
  endif

  if (enable_quasi_stokes_mean) then
   edges  = (/nx,ny,nz,1/)
   iret=nf_inq_varid(ncid,'um',varid)
   iret= nf_put_vara_double(ncid,varid,corner,edges,umean)
   iret=nf_inq_varid(ncid,'vm',varid)
   iret= nf_put_vara_double(ncid,varid,corner,edges,vmean)
   iret=nf_inq_varid(ncid,'wm',varid)
   iret= nf_put_vara_double(ncid,varid,corner,edges,wmean)
   iret=nf_inq_varid(ncid,'bm',varid)
   iret= nf_put_vara_double(ncid,varid,corner,edges,bmean)
  endif

  call ncclos (ncid, iret)

  if (enable_quasi_stokes_hor) then
   harea=harea*counter
   htrans=htrans*counter
   hforc=hforc*counter
  endif
  if (enable_quasi_stokes_ver) then
   varea=varea*counter
   utrans=utrans*counter
   vtrans=vtrans*counter
  endif
  if (enable_quasi_stokes_mean) then
   umean=umean*counter
   vmean=vmean*counter
   wmean=wmean*counter
   bmean=bmean*counter
  endif
 endif
 end subroutine quasi_stokes_write_cdf
