SUBROUTINE BALOMEGA (PVOR,PT,PDIV)

!**** *BALOMEGA*  - approximate quasi-geotrophic omega equation

!     Purpose.
!     --------
!           calculate balanced divergence given temperature and vorticity

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

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

!      Output:
!        PDIV(KFLEV,KSPEC2) omega-eqn divergence

!     Method.
!     -------
!       Solve omega equation. NB: I've made some pretty crude
!       approximations (e.g. pretending I'm in pressure coordinates).
!       See Hoskins, Draghici and Davies, 1978, QJRMS pp33-38 for
!       the derivation of the Q-vector form of the omega equation.

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

!     Modifications.
!     --------------
!        Original : 2001-11-10 Mike Fisher 
!        M.Hamrud      01-Oct-2003 CY28 Cleaning
!        M.Hamrud      10-Jan-2004 CY28R1 Cleaning
!     ------------------------------------------------------------------

USE PARKIND1  ,ONLY : JPIM     ,JPRB
USE YOMHOOK   ,ONLY : LHOOK,   DR_HOOK
USE YOMDIM   , ONLY : NFLEVL, NSPEC2, NPROMA, NFLEVG
USE YOMMP    , ONLY : NBSETLEV, NPTRSV, NSPEC2V, MYSETN
USE YOMGEM   , ONLY : NGPTOT
USE YOMGC    , ONLY : GEMU, GSQM2
USE YOMSTA   , ONLY : STPRE, STTEM
USE YOMCST   , ONLY : RD, RKAPPA, ROMEGA, RA
USE YOMLAP   , ONLY : RLAPDI, NVALUE
USE YOMMSC   , ONLY : N_DEFAULT_REAL_KIND, N_DOUBLE_KIND
USE YOMLUN   , ONLY : NULOUT

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

IMPLICIT NONE

REAL(KIND=JPRB)   ,INTENT(INOUT) :: PVOR(NFLEVL,NSPEC2) 
REAL(KIND=JPRB)   ,INTENT(INOUT) :: PT(NFLEVL,NSPEC2) 
REAL(KIND=JPRB)   ,INTENT(OUT)   :: PDIV(NFLEVL,NSPEC2) 

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

REAL(KIND=JPRB) :: ZGP (NPROMA,9*NFLEVG,(NGPTOT+NPROMA-1)/NPROMA), &
 & ZGPNSU (NPROMA,NFLEVG,(NGPTOT+NPROMA-1)/NPROMA), &
 & ZGPNSV (NPROMA,NFLEVG,(NGPTOT+NPROMA-1)/NPROMA), &
 & ZGP2(NPROMA,3*NFLEVG,(NGPTOT+NPROMA-1)/NPROMA), &
 & ZVORL(NFLEVL,NSPEC2), ZBTXL(NFLEVL,NSPEC2), &
 & ZDIVL(NFLEVL,NSPEC2), ZDIVG(NFLEVG,NSPEC2V), &
 & ZLOCALG(NFLEVG,NSPEC2V), &
 & ZSTAB(NFLEVG), ZDIAG(NFLEVG), ZBELOW(NFLEVG-1),  &
 & ZABOVE(NFLEVG-1),ZWORK(NFLEVG), ZDWORK(NFLEVG), &
 & ZBWORK(NFLEVG-1), ZAWORK(NFLEVG-1)  

REAL(KIND=JPRB) :: ZCORI, ZDP

INTEGER(KIND=JPIM) :: JKGLO, ICEND, IBL, IUP, IDOWN, INFO, &
 & JLEV, JSP, IN, IOFF_VOR, IOFF_DIV, IOFF_U, IOFF_V, IOFF_NS_T, &
 & IOFF_EW_U, IOFF_EW_V, IOFF_EW_T, ILEV  
REAL(KIND=JPRB) :: ZHOOK_HANDLE

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

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

#include "abor1.intfb.h"
#include "trmtos.intfb.h"
#include "trstom.intfb.h"

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

IF (LHOOK) CALL DR_HOOK('BALOMEGA',0,ZHOOK_HANDLE)

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

!      1. Transform to grid space

ZDIVL(:,:) = 0.0_JPRB

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

IOFF_VOR  = 0
IOFF_DIV  = NFLEVG
IOFF_U    = 2*NFLEVG
IOFF_V    = 3*NFLEVG

IOFF_NS_T = 5*NFLEVG
IOFF_EW_U = 6*NFLEVG
IOFF_EW_V = 7*NFLEVG
IOFF_EW_T = 8*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 Q-vector (and additional beta-effect term)

    ZGP2(1:ICEND,JLEV,IBL) = &
     & -( ZGP (1:ICEND,IOFF_EW_U+JLEV,IBL)*ZGP (1:ICEND,IOFF_EW_T+JLEV,IBL)   &
     & +ZGP (1:ICEND,IOFF_EW_V+JLEV,IBL)*ZGP (1:ICEND,IOFF_NS_T+JLEV,IBL)   &
     & )*RD/STPRE(JLEV)  

    ZGP2(1:ICEND,NFLEVG+JLEV,IBL) = &
     & -( ZGPNSU (1:ICEND,JLEV,IBL)*ZGP (1:ICEND,IOFF_EW_T+JLEV,IBL)   &
     & +ZGPNSV (1:ICEND,JLEV,IBL)*ZGP (1:ICEND,IOFF_NS_T+JLEV,IBL)   &
     & )*RD/STPRE(JLEV)  

    ZGP2(1:ICEND,2*NFLEVG+JLEV,IBL) = &
     & - ZGP (1:ICEND,IOFF_EW_T+JLEV,IBL) &
     & *2.0_JPRB*ROMEGA*GSQM2(JKGLO:JKGLO-1+ICEND)/RA &
     & *RD/STPRE(JLEV)  
  ENDDO
ENDDO

!      4. Transform back to spectral space

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

!      5. Calculate RHS of omega equation: -2*div(Q) + beta term

ZDIVL(:,:) = -2.0_JPRB * ZDIVL(:,:) + ZBTXL(:,:)

!      6. Transpose to full vertical columns

CALL TRMTOS (PSPDIV=ZDIVL, PSPDIVG=ZDIVG)

!      7. Calculate the reference stability sigma=-(RT/p) d(ln(theta))/dp

DO JLEV=1,NFLEVG
  IUP   = MAX(1,JLEV-1)
  IDOWN = MIN(NFLEVG,JLEV+1)
  ZSTAB(JLEV)= (-RD*STTEM(JLEV)/STPRE(JLEV)) &
   & *( LOG(STTEM(IUP  )) - RKAPPA*LOG(STPRE(IUP  )) &
   & -LOG(STTEM(IDOWN)) + RKAPPA*LOG(STPRE(IDOWN)) &
   & )/(STPRE(IUP)-STPRE(IDOWN))  
ENDDO

!      8. Invert the linear operator to get the perturbation omega

ZCORI = 2.0_JPRB * ROMEGA / SQRT(2.0_JPRB)

DO JLEV=2,NFLEVG-1
  IUP   = JLEV-1
  IDOWN = JLEV+1

  ZDP = 2.0_JPRB / (STPRE(IUP)-STPRE(IDOWN))
  ZABOVE(JLEV  ) =  ZCORI*ZCORI*ZDP/(STPRE(JLEV)-STPRE(IDOWN))
  ZBELOW(JLEV-1) =  ZCORI*ZCORI*ZDP/(STPRE(IUP )-STPRE(JLEV ))
  ZDIAG(JLEV)    = -(ZABOVE(JLEV)+ZBELOW(JLEV-1))
ENDDO

ILEV=1
IDOWN = ILEV+1
ZDP = -2.0_JPRB / STPRE(IDOWN)
ZABOVE(ILEV) =  ZCORI*ZCORI*ZDP/(STPRE(ILEV)-STPRE(IDOWN))
ZDIAG(ILEV)  = -ZABOVE(ILEV)

ILEV=NFLEVG
IUP =ILEV-1
ZDP = 1.0_JPRB / (STPRE(IUP)-STPRE(ILEV))
ZBELOW(ILEV-1) =  ZCORI*ZCORI*ZDP/(STPRE(IUP)-STPRE(ILEV))
ZDIAG(ILEV)    = -ZBELOW(ILEV-1)

DO JSP=1,NSPEC2V
  IN=NVALUE(JSP+NPTRSV(MYSETN)-1)

  IF (IN==0) THEN
    ZLOCALG(:,JSP) = 0.0_JPRB
  ELSE
    ZDWORK(:) = ZDIAG(:) + RLAPDI(IN)*ZSTAB(:)
    ZAWORK(:) = ZABOVE(:)
    ZBWORK(:) = ZBELOW(:)
    ZWORK(:)  = ZDIVG(:,JSP)

    IF (JPRB ==  N_DEFAULT_REAL_KIND) THEN
      CALL SGTSV (NFLEVG,1,ZBWORK,ZDWORK,ZAWORK,ZWORK,NFLEVG,INFO)
    ELSEIF (JPRB == N_DOUBLE_KIND) THEN
      CALL DGTSV (NFLEVG,1,ZBWORK,ZDWORK,ZAWORK,ZWORK,NFLEVG,INFO)
    ELSE
      call ABOR1('REAL(KIND=JPRB) is neither default real nor double precision')
    ENDIF

    IF (INFO /= 0) THEN
      WRITE (NULOUT,*) 'BALOMEGA: SGTSV/DGTSV RETURNED INFO=',INFO
      CALL ABOR1 ('BALOMEGA: failed to solve tridiagonal system')
    ENDIF

    ZLOCALG(:,JSP) = ZWORK(:)
  ENDIF
ENDDO

!      9. Differentiate omega to get the divergence perturbation

ZDIVG(1,:) = -(ZLOCALG(1,:)-ZLOCALG(2,:))/(STPRE(1)-STPRE(2))

DO JLEV=2,NFLEVG-1
  ZDIVG(JLEV,:) = -(ZLOCALG(JLEV-1,:)-ZLOCALG(JLEV+1,:)) &
   & /(STPRE(JLEV-1)-STPRE(JLEV+1))  
ENDDO

ZDIVG(NFLEVG,:) = -(ZLOCALG(NFLEVG-1,:)-ZLOCALG(NFLEVG,:)) &
 & /(STPRE(NFLEVG-1)-STPRE(NFLEVG))  

!      10. Transpose back to distributed levels

CALL TRSTOM (PSPDIVG=ZDIVG, PSPDIV=PDIV)

!     ------------------------------------------------------------------
IF (LHOOK) CALL DR_HOOK('BALOMEGA',1,ZHOOK_HANDLE)
END SUBROUTINE BALOMEGA
