!
      module epflux
!
! Eliasson-Palm flux calculations (fields EPVY, EPVZ, and EPVDIV).
! B. Foster and Hanli Liu, 2-4/98.
!
! Reference (book):
! Volland, Hans, "Atmospheric Tidal and Planetary Waves", 
!   Kluwer Academic Publishers (Norwell, MA), Section 2.7
!
! Sy is EP vector in latitude (field EPVY), 
! Sz is EP vector in vertical (field EPVZ), and 
! S  is (Sy,Sz) divergence (forcing) (field EPVDIV).
!
!               v'T'    duzm
! Sy = -u'v' - ----- * ------	! err in sign, see epfluxy below
!              gamma     dz
!
!       [        (    1      d(uzm*cos(lat))    )    v'T' ]
! Sz = -[ u'w' + (-------- * --------------- - f) * ----- ]
!       [        (cos(lat)         dy           )   gamma ]
!
!            1        d(cos(lat)**2 * Sy)      1     d(rhozm*Sz)
! S  = ------------ * -------------------- + ----- * -----------
!      cos(lat)**2             dy            rhozm       dz
!
      use proc,only: spval,dlon,p0,gcmlat,nlat
      implicit none
!
      logical mkepv
      real,parameter,private ::
     +  g=9.8,			! gravity
     +  cp=1004.,		! specific heat
     +  ts=300.,		! surface temperature
     +  re=6371.,		! earth radius (km)
     +  pi=3.1415,		! pi
     +  om=2.*pi/86400.,	! omega (for coriolis)
     +  dy=2.*pi*re*1000.*dlon/360., ! delta lon 5 deg (m)
     +  r=287.,			! scale height reference constant
     +  d2r=pi/180.             ! degree to radian
!
! 3d fields necessary for calculation of epvz and epvdiv:
! (will be dimensioned (imx,kmx,jmx), and saved by save_epv)
!
      real,allocatable ::
     +  t3d(:,:,:),u3d(:,:,:),v3d(:,:,:),w3d(:,:,:),epvy3d(:,:,:)
!
! Zonal mean fields for ep fluxes (dimensioned (kmx,jmx) and saved
! by save_epv):
!
      real,allocatable ::
     +  tzm(:,:),uzm(:,:),vzm(:,:),wzm(:,:),rhozm(:,:),zpht(:),
     +  gamma(:,:)
      contains
!-------------------------------------------------------------------
      subroutine epfluxy(t,u,v,fout,nlon,nlev,lat)
!
! Calculate Eliassen-Palm flux in Y (lat) direction.
!
! S(y) = -u'v' - (v't'/gamma * du(i)/dz) 
! Where primed values are (f(i)-f(zm)), and df is f(k+1)-f(k-1),
!   and gamma = dt(zm)/dz(zm)+(g*t)/(ts*cp)
!
! Args:
      integer,intent(in) :: nlon,nlev,lat
      real,dimension(nlon,nlev),intent(in)  :: t,u,v
      real,dimension(nlon,nlev),intent(out) :: fout
!
! Locals:
      integer :: i,k      
      real,dimension(nlev) :: dudz
!
! Calculate dudz (extrapolate boundaries): 
      do k=2,nlev-1
        dudz(k) = (uzm(k+1,lat)-uzm(k-1,lat)) / (zpht(k+1)-zpht(k-1))
      enddo
      dudz(1) = 2.*dudz(2)-dudz(3)
      dudz(nlev) = 2.*dudz(nlev-1)-dudz(nlev-2)
!
! Longitude and vertical loops:
! 3/12/98: changed sign due to error in Hans ref cited above.
!   Old fout before sign fix:
!         fout(i,k) = -((u(i,k)-uzm(k,lat))*(v(i,k)-vzm(k,lat))) -
!    +                 ((v(i,k)-vzm(k,lat))*(t(i,k)-tzm(k,lat))) / 
!    +                 gamma(k,lat) * dudz(k)
!
      do k=1,nlev
        do i=1,nlon
          fout(i,k) = -((u(i,k)-uzm(k,lat))*(v(i,k)-vzm(k,lat))) +
     +                 ((v(i,k)-vzm(k,lat))*(t(i,k)-tzm(k,lat))) / 
     +                 gamma(k,lat) * dudz(k)
        enddo
      enddo
      end subroutine epfluxy
!-------------------------------------------------------------------
      subroutine epfluxz(t,u,v,w,fout,nlon,nlev,glat,lat)
!
! Calculate EP flux in Z (vertical) direction:
! S(z) = u'w' + (1/cos(lat)*d(uzm*cos(lat))/dy-f) * v't'/gamma
!
! Args:
      integer,intent(in) :: nlon,nlev,lat
      real,dimension(nlon,nlev),intent(in)  :: t,u,v,w
      real,dimension(nlon,nlev),intent(out) :: fout
      real,intent(in) :: glat	! current geog latitude
!
! Locals:
      integer :: i,k      
      real :: f,dudy
!
! This routine should be called w/ lat=2,nlat-1:
      if (lat<2.or.lat>nlat) write(6,"('>>> WARNING epfluxz: ',
     +  ' lat=',i3,'  This routine does not calculate latitude ',
     +  ' boundaries.',/,' It must not be called with lat==1 or ',
     +  'lat==nlat (',i3,').')") lat,nlat
      f = 2.*om*sin(glat*d2r)	! coriolis
!
! Calculate epvz:
      do k=1,nlev
        dudy = (uzm(k,lat+1)*cos(gcmlat(lat+1)*d2r)-
     +          uzm(k,lat-1)*cos(gcmlat(lat-1)*d2r)) / 
     +          (cos(glat*d2r)*2.*dy)
        do i=1,nlon
          fout(i,k) = 
     +      -(((u(i,k)-uzm(k,lat))*(w(i,k)-wzm(k,lat))) + (dudy-f)*
     +      (v(i,k)-vzm(k,lat))*(t(i,k)-tzm(k,lat)) / gamma(k,lat))
        enddo
      enddo
      end subroutine epfluxz
!-------------------------------------------------------------------
      subroutine epfluxdiv(t,u,v,w,epvz,fout,nlon,nlev,glat,lat)
!
! Calculate Eliassen-Palm flux divergence (forcing) at current lat:
! S(div) = 1/cos2(lat) * d(cos2(lat)*Sy)/dy + 1/rho*d(rhozm*Sz)/dz
!
! Args:
      integer,intent(in) :: nlon,nlev,lat
      real,dimension(nlon,nlev),intent(in)  :: t,u,v,w,epvz
      real,dimension(nlon,nlev),intent(out) :: fout
      real,intent(in) :: glat	! current geog latitude
!
! Locals:
      real :: depydy,depzdz(nlev),f
      integer :: i,k,ier
!
! This routine should be called w/ lat=2,nlat-1:
      if (lat<2.or.lat>nlat) write(6,"('>>> WARNING epfluxdiv: ',
     +  ' lat=',i3,'  This routine does not calculate latitude ',
     +  ' boundaries.',/,' It must not be called with lat==1 or ',
     +  'lat==nlat (',i3,').')") lat,nlat
!
! Calculate epvdiv:
      f = 2.*om*sin(glat*d2r)	! coriolis
      do i=1,nlon
        do k=2,nlev-1
          depzdz(k) = 
     +      (rhozm(k+1,lat)*epvz(i,k+1)-rhozm(k-1,lat)*epvz(i,k-1))/
     +      (rhozm(k,lat) * (zpht(k+1)-zpht(k-1)))
        enddo
        depzdz(1) = 2.*depzdz(2)-depzdz(3)
        depzdz(nlev) = 2.*depzdz(nlev-1)-depzdz(nlev-2)
        do k=1,nlev
          depydy = (cos(gcmlat(lat+1)*d2r)**2 * epvy3d(i,k,lat+1)  -
     +              cos(gcmlat(lat-1)*d2r)**2 * epvy3d(i,k,lat-1)) /
     +             (cos(glat*d2r)**2 * 2.*dy)
          fout(i,k) = (depydy+depzdz(k))*86400.
        enddo
      enddo
      end subroutine epfluxdiv
!-------------------------------------------------------------------
      subroutine save_epv(flat,rho,h,imx,kmx,jmx,glat,lat)
      use hist,only: history
!
! Save 3d temperature and winds, 2d zonal means, and 3d edpvy for
!   later calculation of epvz and epvdiv (this routine called at 
!   each latitude by mkderived, which is called from 1st latitude
!   loop in getflds). Epvz and epvdiv are calculated from 2nd 
!   latitude loop in getflds.
!
! Args:
      integer,intent(in) :: imx,kmx,jmx,lat
      type(history),intent(in) :: h
      real,intent(in) :: flat(imx,kmx,h%nflds),rho(imx,kmx),glat
!
! Locals:
      integer,save :: ixt,ixu,ixv,ixw,ixz
      integer :: ier,k
      real :: dlev
!
! Externals:
      real,external :: fmean
      integer,external :: ixfindc
!
! Assume 1st entry if t3d not allocated:
      if (.not.allocated(t3d)) then
!
! Save indices to needed fields in flat:
! (w is needed only for epvz and epvdiv)
!
        ixt  = ixfindc(h%fnames,h%nflds,'TN      ')
        ixu  = ixfindc(h%fnames,h%nflds,'UN      ')
        ixv  = ixfindc(h%fnames,h%nflds,'VN      ')
        ixw  = ixfindc(h%fnames,h%nflds,'W       ')
        ixz  = ixfindc(h%fnames,h%nflds,'Z       ')
        if (ixt==0.or.ixu==0.or.ixv==0.or.ixz==0) then
          write(6,"('>>> WARNING save_epv: need t,u,v,z for ',
     +      'ep fluxes: ixt,u,v,z=',4i4)") ixt,ixu,ixv,ixz 
          stop 'save_epv'
        endif
!
! Allocate 3d fields:
! (assume all are not allocated if t3d is not)
!
        allocate(t3d(imx,kmx,jmx),stat=ier)
        if (ier /= 0) call allocerr(ier,"save_epv allocating t3d")
        allocate(u3d(imx,kmx,jmx),stat=ier)
        if (ier /= 0) call allocerr(ier,"save_epv allocating u3d")
        allocate(v3d(imx,kmx,jmx),stat=ier)
        if (ier /= 0) call allocerr(ier,"save_epv allocating v3d")
        if (ixw > 0) then
          allocate(w3d(imx,kmx,jmx),stat=ier)
          if (ier /= 0) call allocerr(ier,"save_epv allocating w3d")
        endif
        allocate(epvy3d(imx,kmx,jmx),stat=ier)
        if (ier /= 0) call allocerr(ier,"save_epv allocating epvy3d")
!
! Allocate 2d zonal means:
        allocate(tzm(kmx,jmx),stat=ier)
        if (ier /= 0) call allocerr(ier,"save_epv allocating tzm")
        allocate(uzm(kmx,jmx),stat=ier)
        if (ier /= 0) call allocerr(ier,"save_epv allocating uzm")
        allocate(vzm(kmx,jmx),stat=ier)
        if (ier /= 0) call allocerr(ier,"save_epv allocating vzm")
        if (ixw > 0) then
          allocate(wzm(kmx,jmx),stat=ier)
          if (ier /= 0) call allocerr(ier,"save_epv allocating wzm")
        endif
        allocate(rhozm(kmx,jmx),stat=ier)
        if (ier /= 0) call allocerr(ier,"save_epv allocating rhozm")
        allocate(gamma(kmx,jmx),stat=ier)
        if (ier /= 0) call allocerr(ier,"save_epv allocating gamma")
!
! Allocate height, made from zp and reference scale height:
        allocate(zpht(kmx),stat=ier)
        if (ier /= 0) call allocerr(ier,"save_epv allocating zpht")
      endif
!
! Save 3d (imx,kmx,jmx) temperature and winds for current lat:
      t3d(:,:,lat) = flat(:,:,ixt)
      u3d(:,:,lat) = flat(:,:,ixu)
      v3d(:,:,lat) = flat(:,:,ixv)
      if (ixw > 0) w3d(:,:,lat) = flat(:,:,ixw)
!
! Save 2d zonal means and gamma (kmx,jmx) for current lat:
      dlev = (h%zpt-h%zpb)/float(h%nzp-1)
      do k=1,kmx
        tzm(k,lat) = fmean(flat(:,k,ixt),imx,spval,0)
        uzm(k,lat) = fmean(flat(:,k,ixu),imx,spval,0)
        vzm(k,lat) = fmean(flat(:,k,ixv),imx,spval,0)
        if (ixw > 0) wzm(k,lat) = fmean(flat(:,k,ixw),imx,spval,0)
        rhozm(k,lat) = fmean(rho(:,k),imx,spval,0)
!
! Heights are -Hln(p/p0), where H = r*ts/g is reference scale height (m)
! (zpb is bottom zp -ln(p/p0), dlev is calculated above from zpt, zpb, 
!  and nzp):
        zpht(k) = (h%zpb+(k-1)*dlev)*r*ts/g
      enddo
!
! 3/12/98: remove tzm and ts terms from gamma. Old gamma:
!       gamma(k,lat) = ((tzm(k+1,lat)-tzm(k-1,lat)) / 
!    +                  (zzm(k+1,lat)-zzm(k-1,lat))) + 
!    +                 ((g*tzm(k,lat)) / (ts*cp))
! 3/31/98: change zzm to zpht:
!       gamma(k,lat) = ((tzm(k+1,lat)-tzm(k-1,lat)) / 
!    +                  (zpht(k+1)-zpht(k-1))) + (g / cp)
! 5/11/98: use original 3/12 equation, but with zzm changed to zpht: 
!
      do k=2,kmx-1
        gamma(k,lat) = ((tzm(k+1,lat)-tzm(k-1,lat)) / 
     +                  (zpht(k+1)-zpht(k-1))) + 
     +                 ((g*tzm(k,lat)) / (ts*cp))
      enddo
      gamma(1,lat) = 2.*gamma(2,lat)-gamma(3,lat)
      gamma(kmx,lat) = 2.*gamma(kmx-1,lat)-gamma(kmx-2,lat)
!
! Save 3d epvy:
!     subroutine epfluxy(t,u,v,fout,nlon,nlev,glat,lat)
!
      call epfluxy(flat(:,:,ixt),flat(:,:,ixu),flat(:,:,ixv),
     +  epvy3d(:,:,lat),imx,kmx,lat)
!
      end subroutine save_epv
!-------------------------------------------------------------------
      subroutine cleanepv
!
! Remove space allocated by save_epv (called from getflds):
      if (allocated(t3d)) deallocate(t3d)
      if (allocated(u3d)) deallocate(u3d)
      if (allocated(v3d)) deallocate(v3d)
! w3d was allocated only if epvz and/or epvdiv
      if (allocated(w3d)) deallocate(w3d)
      if (allocated(epvy3d)) deallocate(epvy3d)
      if (allocated(tzm)) deallocate(tzm)
      if (allocated(uzm)) deallocate(uzm)
      if (allocated(vzm)) deallocate(vzm)
! wzm was allocated only if epvz and/or epvdiv
      if (allocated(wzm)) deallocate(wzm)
      if (allocated(zpht)) deallocate(zpht)
      if (allocated(rhozm)) deallocate(rhozm)
      if (allocated(gamma)) deallocate(gamma)
      end subroutine cleanepv
      end module epflux
