program tnf_xyz_multilevel

! Wave-activity flux based on Eq (38) of Takaya and Nakamura (2001)

! Takaya, K., Nakamura, H., 2001: 
!   A formulation of a phase-independent wave-activity flux 
!   for stationary and migratory quasi-geostrophic eddies 
!   on zonally-varying basic flow. J. Atmos. Sci., 58, 608--627. 
 
! Input file format
!  4-byte, real
!  unformatted direct access 
!  lon; 144 (0 - 357.5E)
!  lat;  73 (90S - 90N)
!  23-level (1000,925,850,700,600,500,400,300,250,200,150,
!            100,70,50,30,20,
!            10, 7, 5,3,2,1,0.4)

!  Z(m),U(m/s),V(m/s),T(Kelvin) for 
!   fields to be subtracted by basic state
!   and basic state.

! Output data
!   x,y,z component 
!   Vertical component is log-p coordinate
!   The number of vertical levels are reducecd by 2 (uppermost and lowermost levels)
!   (144x73x21)

!  Scale height (H=RT/g) is caluculates with T=250K,R=2.9E2 J/kg/K
!  and g=9.8. so, H=7400m

!   p = p/1000.
!   z = -H*log(p)
!   f=2*7.29E-5 *sin(psi)
!   N^2 = g/theta * D(theta)/Dz

! Examples of comilation 
! ifort -o  tnf_xyz_multilevel.run tnf_xyz_multilevel.f90 -assume byterecl
! ifc -o tnf_xyz_multilevel.run tnf_xyz_multilevel.f90 -lIEPCF90
! g95 -o tnf_xyz_multilevel.run tnf_xyz_multilevel.f90

! This program may work like this; 
!
! ./tnf_xyz_multilevel.run << EOF
! z.dat
! z.clim.dat
! u.dat
! u.clim.dat
! v.dat
! v.clim.dat
! T.dat
! T.clim.dat
! tnfx.dat
! tnfy.dat
! tnfz.dat
! EOF

  implicit none

  real,parameter :: unavl=9.999e+20 !data not available
  integer,parameter :: mlon=144,mlat=73,mlev=23,llev=mlev-2
  integer,parameter :: mmax=mlon*mlat*mlev,lmax=mlon*mlat*llev
! grid interval in degree
!  integer,parameter :: dlati=2.5,dloni=2.5
! (corrected 2012/11/5)
  real,parameter :: dlati=2.5,dloni=2.5
  real :: dzm

! The limit latitudes of QG approximation
!  ( This may be changed)
! 43-> 15N, 31-> 15S
  integer,parameter :: neqlats=31,neqlatn=43

! Unit number for input file
  integer,parameter :: un_z=10,un_zb=11
  integer,parameter :: un_T=12,un_Tb=13
  integer,parameter :: un_u=14,un_ub=15
  integer,parameter :: un_v=16,un_vb=17
! Unit number for output file
  integer,parameter :: un_tnfx=30,un_tnfy=31,un_tnfz=32

  real :: z(mlon,mlat,mlev) ,zb(mlon,mlat,mlev)
  real :: T(mlon,mlat,mlev) ,Tb(mlon,mlat,mlev)
  real :: u(mlon,mlat,mlev) ,ub(mlon,mlat,mlev)
  real :: v(mlon,mlat,mlev) ,vb(mlon,mlat,mlev)

  real :: tnfx(mlon,mlat,2:mlev-1),tnfy(mlon,mlat,2:mlev-1)
  real :: tnfz(mlon,mlat,2:mlev-1)

  real :: Tbldd(mlon,mlat),Tbmdd(mlon,mlat),Tbudd(mlon,mlat)
  real :: Tmdd(mlon,mlat)
  real :: umdd(mlon,mlat),ubmdd(mlon,mlat)
  real :: vmdd(mlon,mlat),vbmdd(mlon,mlat)
  real :: zmdd(mlon,mlat),zbmdd(mlon,mlat)
  real :: stabmdd(mlon,mlat)
  real :: tnfxdd(mlon,mlat),tnfydd(mlon,mlat),tnfzdd(mlon,mlat)

  integer irec,irecb,ilev,ilon,ilat

  real :: preu,prem,prel
  real,parameter:: kappa=0.286
  real,parameter :: Ts=250.,gscnst=2.9E2,gra=9.8
  real,parameter :: sclhght=gscnst*Ts/gra
  integer nendfile

  real,dimension(mlev) :: pre=(/1000.,925.,850.,700. &
       &        ,600.,500.,400.,300.,250.,200.,150.  &
       &        ,100.,70.,50.,30.,20.,10. &
       &        ,7.,5.,3.,2.,1.,0.4/)

  real,parameter :: ubmin=1.,tday=60.*60.*24.,erd=6.371E+06
  real :: pai,dgrd,rloni,rloni2, rlati,rlati2,oum2

! constant
  pai   = acos(-1.)
  dgrd  = pai/180.
  rloni = dloni*dgrd
  rloni2= 2.*rloni
  rlati = dlati*dgrd
  rlati2= 2.*rlati
  oum2 = 2.*2.*pai/tday
  
! Note that block size is multiplied by 4. 
! Some compilers do not need this.

  write(6,*) "input file name of z"
  call of_udo (un_z,mmax*4)
  
  write(6,*) "input file name of z of basic state"
  call of_udo (un_zb,mmax*4)

  write(6,*) "input file name of u"
  call of_udo (un_u,mmax*4)

  write(6,*) "input file name of u of basic state"
  call of_udo (un_ub,mmax*4)

  write(6,*) "input file name of v "
  call of_udo (un_v,mmax*4)

  write(6,*) "input file name of v of basic state"
  call of_udo (un_vb,mmax*4)

  write(6,*) "input file name of T"
  call of_udo (un_T,mmax*4)

  write(6,*) "input file name of T of basic state"
  call of_udo (un_Tb,mmax*4)

! for output file
  write(6,*) "The output file name of TN flux x"
  call of_udr (un_tnfx,lmax*4)
  write(6,*) "The output file name of TN flux y"
  call of_udr (un_tnfy,lmax*4)
  write(6,*) "The output file name of TN flux z"
  call of_udr (un_tnfz,lmax*4)

  irec=0
  irecb=0

! for irec
  do 

! If records of anomaly and basic state fields are different, 
!  modify the increment.

     irec=irec+1

! This is an example that the basic state is climatology for each month
     irecb=irecb+1
     if(irecb .eq. 13)  irecb = 1

! test
!     if(irec .eq. 2) exit

     read(un_z,rec=irec,err=300)  z
     read(un_zb,rec=irecb,err=300) zb
     read(un_T,rec=irec,err=300)  T
     read(un_Tb,rec=irecb,err=300) Tb
     read(un_u,rec=irec,err=300)  u
     read(un_ub,rec=irecb,err=300) ub
     read(un_v,rec=irec,err=300 ) v
     read(un_vb,rec=irecb,err=300) vb


     do ilev=2,mlev-1

        preu=pre(ilev+1)/1000.
        prem=pre(ilev)/1000.
        prel=pre(ilev-1)/1000.
        dzm=-sclhght*alog(preu/prel)

! for calculating flux for each level

        do ilat=1,mlat
           do ilon=1,mlon
              zmdd(ilon,ilat)=z(ilon,ilat,ilev)
              Tmdd(ilon,ilat)=T(ilon,ilat,ilev)
              umdd(ilon,ilat)=u(ilon,ilat,ilev)
              vmdd(ilon,ilat)=v(ilon,ilat,ilev)

              zbmdd(ilon,ilat)=zb(ilon,ilat,ilev)
              Tbmdd(ilon,ilat)=Tb(ilon,ilat,ilev)
              Tbldd(ilon,ilat)=Tb(ilon,ilat,ilev-1)
              Tbudd(ilon,ilat)=Tb(ilon,ilat,ilev+1)
              ubmdd(ilon,ilat)=ub(ilon,ilat,ilev)
              vbmdd(ilon,ilat)=vb(ilon,ilat,ilev)
              
! ! For resorting order in latitude for basic state if necessary.
!               zbmdd(ilon,ilat)=zb(ilon,mlat-ilat+1,ilev)
!               Tbmdd(ilon,ilat)=Tb(ilon,mlat-ilat+1,ilev)
!               Tbldd(ilon,ilat)=Tb(ilon,mlat-ilat+1,ilev-1)
!               Tbudd(ilon,ilat)=Tb(ilon,mlat-ilat+1,ilev+1)
!               ubmdd(ilon,ilat)=ub(ilon,mlat-ilat+1,ilev)
!               vbmdd(ilon,ilat)=vb(ilon,mlat-ilat+1,ilev)

           enddo !ilon
        enddo !ilat

! log-p coordinate

        call stability 

        call tnflux_comp
         
        do ilat=1,mlat
           do ilon=1,mlon
              tnfx(ilon,ilat,ilev)=tnfxdd(ilon,ilat)
              tnfy(ilon,ilat,ilev)=tnfydd(ilon,ilat)
              tnfz(ilon,ilat,ilev)=tnfzdd(ilon,ilat)
           enddo
        enddo


     enddo !ilev
     write(un_tnfx,rec=irec)tnfx
     write(un_tnfy,rec=irec)tnfy
     write(un_tnfz,rec=irec)tnfz


  enddo

300 write(*,*) "file ended"
  stop

contains
!     caluculate Brunt Vaisala frequency
  subroutine stability 
    implicit none
      real :: dTdz

      do ilat=1,mlat
         do ilon=1,mlon
            stabmdd(ilon,ilat)=unavl

            if(abs(Tbudd(ilon,ilat)-unavl).le.0.) cycle
            if(abs(Tbmdd(ilon,ilat)-unavl).le.0.) cycle
            if(abs(Tbldd(ilon,ilat)-unavl).le.0.) cycle

            dTdz=(Tbudd(ilon,ilat)-Tbldd(ilon,ilat))/dzm
            stabmdd(ilon,ilat)= &
                 &        gscnst/sclhght* &
                 &(dTdz+kappa/sclhght*Tbmdd(ilon,ilat))
         
            if(stabmdd(ilon,ilat).le.0.)then
               stabmdd(ilon,ilat)=unavl
            endif
         enddo
      enddo
      return
    end subroutine stability
!
    subroutine tnflux_comp
      implicit none
      integer :: ilatn,ilats
      real :: termxU,termxV,termyU,termyV,termzU,termzV
      real :: Tam,dTdx,dTdy
      real :: wspeed
      real :: Tan,Tas,Taw,Tae
      real :: stab,ub,vb,psiam,zam
      real :: uam,vam,dudx,dvdx,dudy,dvdy

      real :: rlat,coriol,coslat
      integer :: ilon,ilone,ilonw
      real :: uamn,uams,uame,uamw,vame,vamw
      real :: fx,fy,fz
! You can input zero into flux where the flux is undefined.      
      do ilat=1,mlat
         do ilon=1,mlon
            tnfxdd(ilon,ilat)=unavl
            tnfydd(ilon,ilat)=unavl
            tnfzdd(ilon,ilat)=unavl
         enddo
      enddo

      do ilat=2,mlat-1

! If the latitude of input data is north to south, change the following lines 
!  and definition of neqlats and neqlatn

         if(ilat .gt. neqlats .and. ilat .lt.neqlatn) cycle

         ilatn=ilat+1
         ilats=ilat-1

!         rlat=dgrd*(90.-real((ilat-1))*dlati)
! (corrected 2012/11/14)
         rlat=dgrd*(real(ilat-1)*dlati-90.)

         coriol=oum2*sin(rlat)
         coslat=cos(rlat)

         do ilon=1,mlon

! If westerly wind of the basic state is smaller than ubmin or easterly, 
! flux is unavailable.

            ub=ubmdd(ilon,ilat)
            if(ub.lt.ubmin) cycle

            vb=vbmdd(ilon,ilat)
            wspeed = sqrt(ub*ub+vb*vb)

            stab=stabmdd(ilon,ilat)    
            if(abs(stab-unavl).le.0.) cycle
      
            ilone=ilon+1
            if(ilone.gt.mlon) ilone=1

            ilonw=ilon-1
            if(ilonw.lt.1) ilonw=mlon

! anomaly
            zam=zmdd(ilon,ilat)-zbmdd(ilon,ilat)
            
            Tan = Tmdd(ilon,ilatn) - Tbmdd(ilon,ilatn)
            Tam = Tmdd(ilon,ilat)  - Tbmdd(ilon,ilat)
            Tas = Tmdd(ilon,ilats) - Tbmdd(ilon,ilats)
            Taw = Tmdd(ilonw,ilat) - Tbmdd(ilonw,ilat)
            Tae = Tmdd(ilone,ilat) - Tbmdd(ilone,ilat)     
         
! geopotential height -> QG streamfunction

            psiam=gra/coriol*zam

            uam=umdd(ilon,ilat)-ub
            vam=vmdd(ilon,ilat)-vb

            uamn=umdd(ilon,ilatn)-ubmdd(ilon,ilatn)
            uams=umdd(ilon,ilats)-ubmdd(ilon,ilats)

            uame=umdd(ilone,ilat)-ubmdd(ilone,ilat)
            uamw=umdd(ilonw,ilat)-ubmdd(ilonw,ilat)

            vame=vmdd(ilone,ilat)-vbmdd(ilone,ilat)
            vamw=vmdd(ilonw,ilat)-vbmdd(ilonw,ilat)

            dudx = (uame-uamw)/(erd*rloni2*coslat)
            dvdx = (vame-vamw)/(erd*rloni2*coslat)

            dudy = (uamn-uams)/(erd*rlati2)

            dTdx=(Tae-Taw)/(erd*rloni2*coslat)
            dTdy=(Tan-Tas)/(erd*rlati2)

            termxU=vam*vam-psiam*dvdx
            termxV=psiam*dudx-uam*vam
            termyU=termxV
            termyV=uam*uam+psiam*dudy

! Here, derivative QG stream function by z is replaced 
!  by temperature by using the following relation
!       d (psim)/d z = gra/(coriol*Ts)* T
            termzU=vam*Tam-psiam*dTdx
!            termzV=-uam*Tam+psiam*dTdy
! (corrected 2012/11/5)
            termzV=-uam*Tam-psiam*dTdy
! tnf
            fx=ub*termxU+vb*termxV
            fy=ub*termyU+vb*termyV
            fz=ub*termzU+vb*termzV

            tnfxdd(ilon,ilat)=prem*coslat/(2.*wspeed)*fx
            tnfydd(ilon,ilat)=prem*coslat/(2.*wspeed)*fy

!  Note that the coefficient of z component of Eq (38) is
!   multipled by gra/(corili*Ts), 
!   which is not multipled at the calculations of termzU and termzV
            tnfzdd(ilon,ilat)= &
                 & prem*coslat*coriol*gra/(2.*wspeed*stab*Ts)*fz
         enddo
      enddo

      return
    end subroutine tnflux_comp

  end program tnf_xyz_multilevel

! external subroutine

  subroutine of_udo (un,bs)
! unformatted, direct access, old file
  implicit none
  character(160) :: file
  integer :: un,bs ! unit number and block size
  integer :: ier

  read(5,'(A160)') file
  write(6,'(A160)') file
  open(unit=un,status='old',file=file,FORM='UNFORMATTED', &
  & recl=bs,access='direct',iostat=ier)

  select case (ier)
  case(1:)
     write(6,*) " has error in opening."
     stop
  case(0)
     write(6,*) " opened successfully."
     return
  case default
     write(6,*) " file end"
     stop
  end select


end subroutine of_udo

subroutine of_udr (un,bs)
! unformatted, direct access, replace file
  implicit none
  character(160) :: file
  integer :: un,bs ! unit number and block size
  integer :: ier

  read(5,'(A160)') file
  write(6,'(A160)') file
  open(unit=un,status='replace',file=file,FORM='UNFORMATTED', &
  & recl=bs,access='direct',iostat=ier)

  select case (ier)
  case(1:)
     write(6,*) " has error in opening."
     stop
  case(0)
     write(6,*) " opened successfully."
     return
  case default
     write(6,*) " file end"
     stop
  end select

end subroutine of_udr
