SUBROUTINE BALNONLIN (PVOR,PZZP)

!**** *BALNONLIN*  - nonlinear balance equation

!     Purpose.
!     --------
!           calculate nonlinearly-balanced geopotential from vorticity

!**   Interface. CALL BALNONLIN (...)
!     ----------

!     Explicit arguments :
!     --------------------
!      Input:
!        PVOR (NFLEVL,KSPEC2) vorticity

!      Output:
!        PZZP (NFLEVL,KSPEC2) nonlinearly-balanced geopotential

!     Method.
!     -------
!       Solve the nonlinear balance equation.

!     Author.        Mike Fisher
!     -------

!     Modifications.
!     --------------
!        Original : 2002-01-17 Mike Fisher 
!        2003-01-10 Mike Fisher : Corrected a bug in the calculation of v.grad(v)
!        M.Hamrud      01-Oct-2003 CY28 Cleaning
!     ------------------------------------------------------------------

USE PARKIND1  ,ONLY : JPIM     ,JPRB
USE YOMHOOK   ,ONLY : LHOOK,   DR_HOOK

USE YOMDIM   , ONLY : NFLEVL, NSPEC2, NPROMA, NFLEVG, NUMP, NSMAX
USE YOMMP    , ONLY : NBSETLEV
USE YOMGEM   , ONLY : NGPTOT
USE YOMGC    , ONLY : GEMU, GSQM2, RCORI
USE YOMCST   , ONLY : RA
USE YOMLAP   , ONLY : NASM0, MYMS, RLAPIN

IMPLICIT NONE

REAL(KIND=JPRB)   ,INTENT(INOUT) :: PVOR(NFLEVL,NSPEC2) 
REAL(KIND=JPRB)   ,INTENT(OUT)   :: PZZP(NFLEVL,NSPEC2) 
REAL(KIND=JPRB) :: ZGP (NPROMA,6*NFLEVG,(NGPTOT+NPROMA-1)/NPROMA), &
 & ZGPNSU (NPROMA,NFLEVG,(NGPTOT+NPROMA-1)/NPROMA), &
 & ZGPNSV (NPROMA,NFLEVG,(NGPTOT+NPROMA-1)/NPROMA), &
 & ZGP2(NPROMA,2*NFLEVG,(NGPTOT+NPROMA-1)/NPROMA), &
 & ZVORL(NFLEVL,NSPEC2), ZDIVL(NFLEVL,NSPEC2)  

INTEGER(KIND=JPIM) :: JKGLO, ICEND, IBL, JLEV, INM, JMLOC, IM, JIR, JN, &
 & IOFF_VOR, IOFF_DIV, IOFF_U, IOFF_V, IOFF_EW_U, IOFF_EW_V  
REAL(KIND=JPRB) :: ZHOOK_HANDLE

INTERFACE
#include "inv_trans.h"
#include "dir_trans.h"
END INTERFACE

!     ------------------------------------------------------------------

!      1. Transform to grid space

IF (LHOOK) CALL DR_HOOK('BALNONLIN',0,ZHOOK_HANDLE)
ZDIVL(:,:) = 0.0_JPRB

CALL INV_TRANS (PSPVOR=PVOR, PSPDIV=ZDIVL, LDUVDER=.TRUE., &
 & LDVORGP=.TRUE., LDDIVGP=.TRUE., &
 & KPROMA=NPROMA, KVSETUV=NBSETLEV, PGP=ZGP)  

IOFF_VOR  = 0
IOFF_DIV  = NFLEVG
IOFF_U    = 2*NFLEVG
IOFF_V    = 3*NFLEVG
IOFF_EW_U = 4*NFLEVG
IOFF_EW_V = 5*NFLEVG

!      2. Calculate meridional derivatives of wind components

DO JKGLO=1,NGPTOT,NPROMA
  ICEND=MIN(NPROMA,NGPTOT-JKGLO+1)
  IBL=(JKGLO-1)/NPROMA+1
  DO JLEV=1,NFLEVG
    ZGPNSU (1:ICEND,JLEV,IBL) =  ZGP (1:ICEND,JLEV+IOFF_EW_V,IBL) &
     & - ZGP (1:ICEND,JLEV+IOFF_VOR,IBL)  &
     & + ZGP (1:ICEND,JLEV+IOFF_U,IBL)    &
     & *GEMU(JKGLO:JKGLO-1+ICEND)   &
     & /GSQM2(JKGLO:JKGLO-1+ICEND)/RA  

    ZGPNSV (1:ICEND,JLEV,IBL) =  ZGP (1:ICEND,JLEV+IOFF_DIV,IBL)  &
     & - ZGP (1:ICEND,JLEV+IOFF_EW_U,IBL) &
     & + ZGP (1:ICEND,JLEV+IOFF_V,IBL)    &
     & *GEMU(JKGLO:JKGLO-1+ICEND)   &
     & /GSQM2(JKGLO:JKGLO-1+ICEND)/RA  

!      3. Calculate : RHS = -(f k x v + v.grad(v)) = -((f+zeta) x v + grad(v.v)/2)

    ZGP2(1:ICEND,JLEV,IBL) = &
     & (RCORI(JKGLO:JKGLO-1+ICEND)+ZGP(1:ICEND,JLEV+IOFF_VOR,IBL))  &
     & *ZGP(1:ICEND,JLEV+IOFF_V,IBL) &
     & -ZGP(1:ICEND,JLEV+IOFF_U,IBL)*ZGP(1:ICEND,JLEV+IOFF_EW_U,IBL) &
     & -ZGP(1:ICEND,JLEV+IOFF_V,IBL)*ZGP(1:ICEND,JLEV+IOFF_EW_V,IBL)   

    ZGP2(1:ICEND,NFLEVG+JLEV,IBL) = &
     & -(RCORI(JKGLO:JKGLO-1+ICEND)+ZGP(1:ICEND,JLEV+IOFF_VOR,IBL))  &
     & *ZGP(1:ICEND,JLEV+IOFF_U,IBL) &
     & -ZGP(1:ICEND,JLEV+IOFF_U,IBL)*ZGPNSU(1:ICEND,JLEV,IBL)        &
     & -ZGP(1:ICEND,JLEV+IOFF_V,IBL)*ZGPNSV(1:ICEND,JLEV,IBL)  
  ENDDO
ENDDO

!      4. Transform back to spectral space

CALL DIR_TRANS (PGP=ZGP2, PSPVOR=ZVORL, PSPDIV=ZDIVL, &
 & KVSETUV=NBSETLEV, KPROMA=NPROMA)  

!      5. Invert Laplacian to solve: del_squared (phi) = div(RHS)

DO JMLOC=1,NUMP
  IM=MYMS(JMLOC)
  DO JIR=0,MIN(1,IM)
    DO JN=IM,NSMAX
      INM=NASM0(IM)+(JN-IM)*2+JIR
      DO JLEV=1,NFLEVL
        PZZP(JLEV,INM)=RLAPIN(JN)*ZDIVL(JLEV,INM)
      ENDDO
    ENDDO
  ENDDO
ENDDO

IF (LHOOK) CALL DR_HOOK('BALNONLIN',1,ZHOOK_HANDLE)
END SUBROUTINE BALNONLIN
