SUBROUTINE BALVERT(PSPZZP,YDSPEC,LDJBBAL,PSPDIV,PSPT,PSPSP)

!      BALVERT - Recombine balanced and unbalanced fields

!     Purpose. transform from unbal div, unbal (T,ps) and unbal O3 to full fields
!     --------

!     Arguments. 
!     ----------
!        Input : PSPZZP(NFLEVG,NSPEC2) - balanced P values from BALSTAT

!     Externals. TRMTOS MXMAOP TRSTOM
!     ----------

!     Method.
!     -------
!       Solve statistical balance equation in spectral space
!       to convert vorticity into geopotential(P) and divergence.
!            (total = balanced + unbalanced )
!         The balance leads from correlation space :
!            (total Vor, unbal Div, unbal T/ps, unbal O3)
!         To model space :
!            (total Vor, total Div, total T/ps, total O3)
!         In 2 steps :

!       (1)     [P]tot = [A].[Vor]tot  (A=tridiag matrix for each m)

!       (2)  [T,ps]tot = [T,ps]unbal + [B].[P]tot + [C].[Div]unbal
!            [Div]tot  = [Div]unbal  + [D].[P]tot
!            [O3]tot   = [O3]unbal   + [E].[P]tot

!               (B,C,D,E=full vertical matrices for each n)

!        This routine solves step 2. Step 1 is solved in BALSTAT.

!     Author. John Derber *ECMWF*
!     -------

!     Modifications.
!     --------------
!        J. Derber     96-10-07 write original
!        F. Bouttier   96-12-13 cleaning
!        E. Holm       97-09-29 include ozone balance
!        M.Hamrud      03-08-01 GFL introduction
!        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 : NFLEVG   ,NFLSUR   ,NSPEC2
USE YOMLAP   , ONLY : NVALUE
USE YOMJG    , ONLY : STPS     ,SDIV     ,SO3, SPJB_VARS_INFO, &
                    & N_SPJB_VARS
USE YOMMP    , ONLY : MYSETN   ,NPTRSV   ,NSPEC2V, NPSP
USE YOMGRB   , ONLY : NGRBO3
USE YOMCVA   , ONLY : NVA3D
USE SPECTRAL_FIELDS_MOD

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

IMPLICIT NONE

REAL(KIND=JPRB)   ,INTENT(IN)    :: PSPZZP(NFLSUR,NSPEC2) 
TYPE(SPECTRAL_FIELD),OPTIONAL, INTENT(INOUT) :: YDSPEC
LOGICAL             ,OPTIONAL, INTENT(IN)    :: LDJBBAL 
REAL(KIND=JPRB)     ,OPTIONAL, INTENT(OUT)   :: PSPDIV(NFLSUR,NSPEC2)
REAL(KIND=JPRB)     ,OPTIONAL, INTENT(OUT)   :: PSPT(NFLSUR,NSPEC2)
REAL(KIND=JPRB)     ,OPTIONAL, INTENT(OUT)   :: PSPSP(NSPEC2)

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

INTEGER(KIND=JPIM) :: IN, JSP

REAL(KIND=JPRB) :: ZIN(2*NFLEVG)
REAL(KIND=JPRB) :: ZTPS(NFLEVG+1),ZDIV(NFLEVG),ZO3(NFLEVG)
REAL(KIND=JPRB) :: ZSZPG(NFLEVG,NSPEC2V), ZSDIVG(NFLEVG,NSPEC2V)
REAL(KIND=JPRB) :: ZSPTG(NFLEVG,NSPEC2V)
REAL(KIND=JPRB), ALLOCATABLE :: ZSO3G(:,:,:)
REAL(KIND=JPRB) :: ZSP(NSPEC2)
REAL(KIND=JPRB) :: ZSPSPG(NSPEC2V)
LOGICAL         :: LLSELECT3D(NVA3D), LLJBBAL
INTEGER(KIND=JPIM) :: JFIELD

REAL(KIND=JPRB) :: ZHOOK_HANDLE

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

INTERFACE
#include "mxmaop.h"
END INTERFACE

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

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

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

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

IF(PRESENT(LDJBBAL)) THEN
  LLJBBAL=LDJBBAL
ELSE
  LLJBBAL=.TRUE.
ENDIF

!     1.    Transpose data to vertical columns
!           ----------------------------------

IF (.NOT.LLJBBAL .AND. .NOT.(PRESENT(PSPDIV) .AND. PRESENT(PSPT) .AND. PRESENT(PSPSP)) ) THEN
  CALL ABOR1('BALVERT: OUTPUT ARRAYS PSPDIV, PSPT AND PSPSP REQUIRED FOR SPBS')
ENDIF

LLSELECT3D(:)=.FALSE.
IF (LLJBBAL) THEN
  DO JFIELD=1,N_SPJB_VARS
    IF (SPJB_VARS_INFO(JFIELD)%IGRIBCODE==NGRBO3) THEN
      LLSELECT3D(SPJB_VARS_INFO(JFIELD)%IPTJB)=.TRUE.
    ENDIF
  ENDDO

  ALLOCATE (ZSO3G(NFLEVG,NSPEC2V,COUNT(LLSELECT3D)))
  IF (NPSP==1) ZSP=YDSPEC%SPR

  CALL TRMTOS (PSPVOR=PSPZZP,PSPDIV=YDSPEC%DIV, PSPT=YDSPEC%TEM, &
             & PSPSEL3D=YDSPEC%SP3D, PSPSP=ZSP, &
             & PSPVORG=ZSZPG, PSPDIVG=ZSDIVG,PSPTG=ZSPTG,PSPSEL3DG=ZSO3G, &
             & PSPSPG=ZSPSPG, LDSELECT3D=LLSELECT3D)
ELSE 
  CALL TRMTOS (PSPVOR=PSPZZP, PSPVORG=ZSZPG)
  ZSDIVG(:,:)=0.0_JPRB
  ZSPTG(:,:)=0.0_JPRB
  ZSPSPG(:)=0.0_JPRB
ENDIF

!     2.    Convert from P and unbalanced fields to full fields
!           ---------------------------------------------------
CALL GSTATS(1090,0)
!$OMP PARALLEL DO SCHEDULE(STATIC) PRIVATE(JSP,IN,ZIN,ZTPS,ZDIV)
DO JSP=1,NSPEC2V
  IN=NVALUE(JSP+NPTRSV(MYSETN)-1)

!           2.1 (T,ps)

  ZIN(1:NFLEVG)          = ZSZPG(1:NFLEVG,JSP)
  ZIN(NFLEVG+1:2*NFLEVG) = ZSDIVG(1:NFLEVG,JSP)
  CALL MXMAOP(STPS(1,1,IN),1,NFLEVG+1,&
   & ZIN,1,2*NFLEVG,ZTPS,1,NFLEVG+1,&
   & NFLEVG+1,2*NFLEVG,1)  
  ZSPTG(1:NFLEVG,JSP) = ZSPTG(1:NFLEVG,JSP)+ZTPS(1:NFLEVG)
  ZSPSPG(JSP)         = ZSPSPG(JSP)+ZTPS(NFLEVG+1)

!           2.2 Divergence

  CALL MXMAOP(SDIV(1,1,IN),1,NFLEVG,&
   & ZSZPG(1,JSP),1,NFLEVG,ZDIV,1,NFLEVG,&
   & NFLEVG,NFLEVG,1)  
  ZSDIVG(1:NFLEVG,JSP) = ZSDIVG(1:NFLEVG,JSP)+ZDIV(1:NFLEVG)
ENDDO
!$OMP END PARALLEL DO

!           2.3 Ozone

IF (ANY(SPJB_VARS_INFO(:)%IGRIBCODE==NGRBO3) .AND. LLJBBAL) THEN
!$OMP PARALLEL DO SCHEDULE(STATIC) PRIVATE(JSP,IN,ZO3)
  DO JSP=1,NSPEC2V
    IN=NVALUE(JSP+NPTRSV(MYSETN)-1)
    CALL MXMAOP(SO3(1,1,IN),1,NFLEVG,&
     & ZSZPG(1,JSP),1,NFLEVG,ZO3,1,NFLEVG,&
     & NFLEVG,NFLEVG,1)  
    ZSO3G(1:NFLEVG,JSP,1) = ZSO3G(1:NFLEVG,JSP,1)+ZO3(1:NFLEVG)
  ENDDO
!$OMP END PARALLEL DO
ENDIF
CALL GSTATS(1090,1)

!     3.    Transpose data from vertical columns back to level structure
!           ------------------------------------------------------------

IF (LLJBBAL) THEN
  CALL TRSTOM (PSPDIVG=ZSDIVG,PSPTG=ZSPTG,PSPSEL3DG=ZSO3G,PSPSPG=ZSPSPG,&
             & PSPDIV=YDSPEC%DIV,PSPT=YDSPEC%TEM, PSPSEL3D=YDSPEC%SP3D, &
             & PSPSP=ZSP, LDSELECT3D=LLSELECT3D)
  IF (NPSP==1) YDSPEC%SPR=ZSP
  DEALLOCATE (ZSO3G)
ELSE
  CALL TRSTOM (PSPDIVG=ZSDIVG,PSPTG=ZSPTG,PSPSPG=ZSPSPG,&
             & PSPDIV=PSPDIV,PSPT=PSPT, PSPSP=PSPSP, &
             & LDSELECT3D=LLSELECT3D)
ENDIF

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

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