SUBROUTINE BALOMEGAAD (PVOR,PT,PVOR7,PT7,PDIV)

!**** *BALOMEGAAD* - approximate quasi-geotrophic omega equation *ADJOINT*

!     Purpose.
!     --------
!           calculate a balanced divergence perturbation for Jb

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

!     Explicit arguments :
!     --------------------
!      Input:
!        PVOR (NFLEVL,KSPEC2) vorticity perturbation
!        PT   (NFLEVL,KSPEC2) temperature perturbation
!        PVOR7(NFLEVL,KSPEC2) background vorticity
!        PT7  (NFLEVL,KSPEC2) background temperature
!        PDIV (NFLEVL,KSPEC2) divergence perturbation

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

!     Method.
!     -------
!       Solve omega equation. NB: I've made some pretty crude
!       approximations (e.g. pretending I'm in pressure coordinates).

!       NB: the omega equation is linearized about the background, not
!       the trajectory. This is because, in the change of variable, PVOR
!       and PT are departures from the background, not from the trajectory.

!       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(INOUT) :: PVOR7(NFLEVL,NSPEC2) 
REAL(KIND=JPRB)   ,INTENT(INOUT) :: PT7(NFLEVL,NSPEC2) 
REAL(KIND=JPRB)   ,INTENT(INOUT) :: 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), &
 & ZGP7(NPROMA,9*NFLEVG,(NGPTOT+NPROMA-1)/NPROMA), &
 & ZGPNSU7(NPROMA,NFLEVG,(NGPTOT+NPROMA-1)/NPROMA), &
 & ZGPNSV7(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, IOFF_T  
REAL(KIND=JPRB) :: ZHOOK_HANDLE

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

INTERFACE
#include "inv_trans.h"
#include "inv_transad.h"
#include "dir_transad.h"
END INTERFACE

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

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

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

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

!      1. Transform to grid space (background)

ZDIVL(:,:) = 0.0_JPRB
 
CALL INV_TRANS (PSPVOR=PVOR7, PSPDIV=ZDIVL, PSPSCALAR=PT7, LDUVDER=.TRUE., &
 & LDSCDERS=.TRUE., LDVORGP=.TRUE., LDDIVGP=.TRUE., &
 & KPROMA=NPROMA, KVSETUV=NBSETLEV, KVSETSC=NBSETLEV, PGP=ZGP7)  
IOFF_VOR  = 0
IOFF_DIV  = NFLEVG
IOFF_U    = 2*NFLEVG
IOFF_V    = 3*NFLEVG
IOFF_T    = 4*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 (background)

CALL GSTATS(1842,0)
DO JKGLO=1,NGPTOT,NPROMA
  ICEND=MIN(NPROMA,NGPTOT-JKGLO+1)
  IBL=(JKGLO-1)/NPROMA+1
  DO JLEV=1,NFLEVG
    ZGPNSU7(1:ICEND,JLEV,IBL) =  ZGP7(1:ICEND,JLEV+IOFF_EW_V,IBL) &
     & - ZGP7(1:ICEND,JLEV+IOFF_VOR,IBL)  &
     & + ZGP7(1:ICEND,JLEV+IOFF_U,IBL)    &
     & *GEMU(JKGLO:JKGLO-1+ICEND)   &
     & /GSQM2(JKGLO:JKGLO-1+ICEND)/RA  

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

!      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

!      10. Transpose back to distributed levels

CALL GSTATS(1842,1)

CALL TRMTOS (PSPDIV=PDIV, PSPDIVG=ZDIVG)

!      9. Differentiate omega to get the divergence perturbation

CALL GSTATS(1841,0)
ZLOCALG(:,:) = 0.0_JPRB

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

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

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

!      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.0_JPRB) THEN
    ZDIVG(:,JSP) = 0.0_JPRB
  ELSE
    ZDWORK(:) = ZDIAG(:) + RLAPDI(IN)*ZSTAB(:)
    ZAWORK(:) = ZABOVE(:)
    ZBWORK(:) = ZBELOW(:)
    ZWORK(:)  = ZLOCALG(:,JSP)

    IF (JPRB ==  N_DEFAULT_REAL_KIND) THEN
      CALL SGTSV (NFLEVG,1,ZAWORK,ZDWORK,ZBWORK,ZWORK,NFLEVG,INFO)
    ELSEIF (JPRB == N_DOUBLE_KIND) THEN
      CALL DGTSV (NFLEVG,1,ZAWORK,ZDWORK,ZBWORK,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

    ZDIVG(:,JSP) = ZWORK(:)
  ENDIF
ENDDO
CALL GSTATS(1841,1)

!      6. Transpose to full vertical columns

CALL TRSTOM (PSPDIVG=ZDIVG, PSPDIV=PDIV)

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

CALL GSTATS(1841,0)
ZBTXL(:,:) = PDIV(:,:)

PDIV(:,:) = -2.0_JPRB * PDIV(:,:)

!      4. Transform back to spectral space

DO JKGLO=1,NGPTOT,NPROMA
  ICEND=MIN(NPROMA,NGPTOT-JKGLO+1)
  IBL=(JKGLO-1)/NPROMA+1
  ZGP2(1:ICEND,:,IBL) = 0.0_JPRB
ENDDO

ZVORL(:,:)  = 0.0_JPRB
CALL GSTATS(1841,1)

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

!      3. Calculate Q-vector (and additional beta-effect term)

CALL GSTATS(1899,0)
DO JKGLO=1,NGPTOT,NPROMA
  ICEND=MIN(NPROMA,NGPTOT-JKGLO+1)
  IBL=(JKGLO-1)/NPROMA+1
  DO JLEV=1,NFLEVG

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

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

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

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

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

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

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

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

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

!      2. Calculate meridional derivatives of wind components

    ZGP (1:ICEND,JLEV+IOFF_DIV,IBL)  = ZGPNSV (1:ICEND,JLEV,IBL)

    ZGP (1:ICEND,JLEV+IOFF_EW_U,IBL) = ZGP (1:ICEND,JLEV+IOFF_EW_U,IBL) &
     & -ZGPNSV (1:ICEND,JLEV,IBL)  

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

    ZGP (1:ICEND,JLEV+IOFF_VOR,IBL)  = -ZGPNSU (1:ICEND,JLEV,IBL)

    ZGP (1:ICEND,JLEV+IOFF_U,IBL)    = ZGPNSU (1:ICEND,JLEV,IBL)     &
     & *GEMU(JKGLO:JKGLO-1+ICEND)     &
     & /GSQM2(JKGLO:JKGLO-1+ICEND)/RA  

    ZGP (1:ICEND,JLEV+IOFF_T,IBL)    = 0.0_JPRB
  ENDDO
ENDDO
CALL GSTATS(1899,1)

!      1. Transform to grid space

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

PDIV(:,:) = 0.0_JPRB

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