! (C) Copyright 2025- ECMWF.
! (C) Copyright 2025- Meteo-France.
!
! This software is licensed under the terms of the Apache Licence Version 2.0
! which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
! In applying this licence, ECMWF does not waive the privileges and immunities
! granted to it by virtue of its status as an intergovernmental organisation
! nor does it submit to any jurisdiction.
!

MODULE TRGL_MOD

USE PARKIND1, ONLY : JPIM

IMPLICIT NONE

TYPE TRGL_BUFFERS
  INTEGER(KIND=JPIM) :: ISENDCOUNT = -9999
  INTEGER(KIND=JPIM) :: IRECVCOUNT = -9999
  INTEGER(KIND=JPIM) :: INSEND = -9999
  INTEGER(KIND=JPIM) :: INRECV = -9999
  INTEGER(KIND=JPIM) :: IFLDS = 0
  LOGICAL :: LLTRGTOL = .FALSE.
  LOGICAL :: LLPGPONLY = .FALSE.
  LOGICAL :: LLINDER = .FALSE.

  INTEGER(KIND=JPIM), ALLOCATABLE :: ISENDTOT (:)
  INTEGER(KIND=JPIM), ALLOCATABLE :: IRECVTOT (:)
  INTEGER(KIND=JPIM), ALLOCATABLE :: ISEND(:)
  INTEGER(KIND=JPIM), ALLOCATABLE :: IRECV(:)
  INTEGER(KIND=JPIM), ALLOCATABLE :: IINDEX(:)
  INTEGER(KIND=JPIM), ALLOCATABLE :: INDOFF(:)
  INTEGER(KIND=JPIM), ALLOCATABLE :: IGPTRSEND(:,:,:)
  INTEGER(KIND=JPIM), ALLOCATABLE :: ISETWL(:)
  INTEGER(KIND=JPIM), ALLOCATABLE :: ISETVL(:)
  INTEGER(KIND=JPIM), ALLOCATABLE :: ISETW(:)
  INTEGER(KIND=JPIM), ALLOCATABLE :: IJPOS(:,:)
  INTEGER(KIND=JPIM), ALLOCATABLE :: IPOSPLUS(:)
  INTEGER(KIND=JPIM), ALLOCATABLE :: IFLDA(:,:)
END TYPE TRGL_BUFFERS

TYPE TRGL_VARS
  INTEGER(KIND=JPIM), ALLOCATABLE :: IUVLEVS(:)
  INTEGER(KIND=JPIM), ALLOCATABLE :: IUVPARS(:)
  INTEGER(KIND=JPIM), ALLOCATABLE :: IGP2PARS(:)
  INTEGER(KIND=JPIM), ALLOCATABLE :: IFLDOFF(:)
  INTEGER(KIND=JPIM), ALLOCATABLE :: IGPTROFF(:)
  INTEGER(KIND=JPIM), ALLOCATABLE :: IGP3APARS(:)
  INTEGER(KIND=JPIM), ALLOCATABLE :: IGP3ALEVS(:)
  INTEGER(KIND=JPIM), ALLOCATABLE :: IGP3BPARS(:)
  INTEGER(KIND=JPIM), ALLOCATABLE :: IGP3BLEVS(:)
  LOGICAL, ALLOCATABLE   :: LLUV(:)
  LOGICAL, ALLOCATABLE   :: LLGP2(:)
  LOGICAL, ALLOCATABLE   :: LLGP3A(:)
  LOGICAL, ALLOCATABLE   :: LLGP3B(:)
END TYPE TRGL_VARS

CONTAINS

SUBROUTINE ALLOCATE_BUFFERS_CST(SELF)
  USE TPM_DISTR       ,ONLY : D, NPRTRNS, NPROC
  USE TPM_TRANS       ,ONLY : NGPBLKS

  CLASS(TRGL_BUFFERS), INTENT(INOUT) :: SELF
  ALLOCATE (SELF%ISENDTOT (NPROC))
  ALLOCATE (SELF%IRECVTOT (NPROC))
  ALLOCATE (SELF%ISEND    (NPROC))
  ALLOCATE (SELF%IRECV    (NPROC))
  ALLOCATE (SELF%IINDEX(D%NLENGTF))
  ALLOCATE (SELF%INDOFF(NPROC))
  ALLOCATE (SELF%IGPTRSEND(2,NGPBLKS,NPRTRNS))
  ALLOCATE (SELF%ISETWL(NPROC))
  ALLOCATE (SELF%ISETVL(NPROC))

END SUBROUTINE ALLOCATE_BUFFERS_CST


SUBROUTINE ALLOCATE_BUFFERS_SR(SELF, KF_GP)
  USE TPM_TRANS       ,ONLY : NGPBLKS

  CLASS(TRGL_BUFFERS), INTENT(INOUT) :: SELF
  INTEGER(KIND=JPIM),INTENT(IN) :: KF_GP

  IF (SELF%LLTRGTOL) THEN
    ALLOCATE (SELF%ISETW(SELF%INSEND))
    ALLOCATE (SELF%IJPOS(NGPBLKS,SELF%INSEND))
    ALLOCATE (SELF%IPOSPLUS(SELF%INSEND))
    ALLOCATE (SELF%IFLDA(KF_GP,SELF%INSEND))
  ELSE
    ALLOCATE (SELF%ISETW(SELF%INRECV))
    ALLOCATE (SELF%IJPOS(NGPBLKS,SELF%INRECV))
    ALLOCATE (SELF%IPOSPLUS(SELF%INRECV))
    ALLOCATE (SELF%IFLDA(KF_GP,SELF%INRECV))
  ENDIF
END SUBROUTINE ALLOCATE_BUFFERS_SR

SUBROUTINE TRGL_ALLOCATE_VARS(SELF, KF_GP, KF_FS)
  USE TPM_TRANS       ,ONLY : NGPBLKS

  CLASS(TRGL_VARS), INTENT(INOUT) :: SELF
  INTEGER(KIND=JPIM), INTENT(IN) :: KF_GP, KF_FS

  ALLOCATE(SELF%IUVLEVS(KF_GP))
  ALLOCATE(SELF%IUVPARS(KF_GP))
  ALLOCATE(SELF%IGP2PARS(KF_GP))
  ALLOCATE(SELF%IFLDOFF(KF_FS))
  ALLOCATE(SELF%IGPTROFF(NGPBLKS))
  ALLOCATE(SELF%LLUV(KF_GP))
  ALLOCATE(SELF%LLGP2(KF_GP))
  ALLOCATE(SELF%LLGP3A(KF_GP))
  ALLOCATE(SELF%LLGP3B(KF_GP))
  ALLOCATE(SELF%IGP3APARS(KF_GP))
  ALLOCATE(SELF%IGP3ALEVS(KF_GP))
  ALLOCATE(SELF%IGP3BPARS(KF_GP))
  ALLOCATE(SELF%IGP3BLEVS(KF_GP))

END SUBROUTINE TRGL_ALLOCATE_VARS

SUBROUTINE TRGL_ALLOCATE_HEAP_BUFFER(Z_HEAP, S1, S2)
  USE PARKIND1  ,ONLY : JPIM, JPRB
  USE ABORT_TRANS_MOD ,ONLY : ABORT_TRANS

  IMPLICIT NONE

  REAL(KIND=JPRB), INTENT(INOUT), ALLOCATABLE :: Z_HEAP(:,:)
  INTEGER(KIND=JPIM), INTENT(IN) :: S1, S2

  IF (ALLOCATED(Z_HEAP) .AND. (S1 /= UBOUND(Z_HEAP,1) .OR. S2 /= SIZE(Z_HEAP,2) )) THEN
    IF (LBOUND(Z_HEAP,1) /= -1) CALL ABORT_TRANS('TRGL_MOD: WRONG Z_HEAP SIZE IN TRGL_ALLOCATE_HEAP_BUFFER ')
    DEALLOCATE(Z_HEAP)
  ENDIF

  IF (.NOT. ALLOCATED(Z_HEAP)) THEN
    ALLOCATE(Z_HEAP(-1:S1,S2))
  ENDIF
END SUBROUTINE TRGL_ALLOCATE_HEAP_BUFFER


SUBROUTINE TRGL_INIT_VARS(YLVARS, KF_SCALARS_G, PGP, PGPUV, PGP3A, PGP3B, PGP2)
  USE PARKIND1  ,ONLY : JPIM, JPRB
  USE TPM_TRANS       ,ONLY : LDIVGP, LSCDERS, LUVDER, LVORGP
  IMPLICIT NONE

  TYPE(TRGL_VARS), INTENT(INOUT) :: YLVARS
  INTEGER(KIND=JPIM), INTENT(IN) :: KF_SCALARS_G
  REAL(KIND=JPRB),OPTIONAL  :: PGP(:,:,:)
  REAL(KIND=JPRB),OPTIONAL  :: PGPUV(:,:,:,:)
  REAL(KIND=JPRB),OPTIONAL  :: PGP3A(:,:,:,:)
  REAL(KIND=JPRB),OPTIONAL  :: PGP3B(:,:,:,:)
  REAL(KIND=JPRB),OPTIONAL  :: PGP2(:,:,:)

  ! Local variables
  INTEGER(KIND=JPIM) :: IUVPAR,IUVLEV,IGP2PAR,IGP3ALEV,IGP3APAR,IGP3BLEV,IGP3BPAR,IPAROFF
  INTEGER(KIND=JPIM) :: IOFF,IOFF1,IOFFNS,IOFFEW,J1,J2

  INTEGER(KIND=JPIM) ::  J

  ASSOCIATE(IUVLEVS=>YLVARS%IUVLEVS, IUVPARS=>YLVARS%IUVPARS, IGP2PARS=>YLVARS%IGP2PARS, &
    &       LLUV=>YLVARS%LLUV, LLGP2=>YLVARS%LLGP2, LLGP3A=>YLVARS%LLGP3A, LLGP3B=>YLVARS%LLGP3B, &
    &       IGP3APARS=>YLVARS%IGP3APARS, IGP3ALEVS=>YLVARS%IGP3ALEVS, IGP3BPARS=>YLVARS%IGP3BPARS, &
    &       IGP3BLEVS=>YLVARS%IGP3BLEVS)

  IUVPAR=0
  IUVLEV=0
  IOFF1=0
  IOFFNS=KF_SCALARS_G
  IOFFEW=2*KF_SCALARS_G

  LLUV(:) = .FALSE.
  IUVPARS(:) = -99
  IUVLEVS(:) = -99

  IF (PRESENT(PGPUV)) THEN
    IOFF=0
    IUVLEV=UBOUND(PGPUV,2)
    IF(LVORGP) THEN
      IUVPAR=IUVPAR+1
      DO J=1,IUVLEV
        IUVLEVS(IOFF+J)=J
        IUVPARS(IOFF+J)=IUVPAR
        LLUV(IOFF+J)=.TRUE.
      ENDDO
      IOFF=IOFF+IUVLEV
    ENDIF
    IF(LDIVGP) THEN
      IUVPAR=IUVPAR+1
      DO J=1,IUVLEV
        IUVLEVS(IOFF+J)=J
        IUVPARS(IOFF+J)=IUVPAR
        LLUV(IOFF+J)=.TRUE.
      ENDDO
      IOFF=IOFF+IUVLEV
    ENDIF
    DO J=1,IUVLEV
      IUVLEVS(IOFF+J)=J
      IUVPARS(IOFF+J)=IUVPAR+1
      IUVLEVS(IOFF+J+IUVLEV)=J
      IUVPARS(IOFF+J+IUVLEV)=IUVPAR+2
    ENDDO
    IUVPAR=IUVPAR+2
    LLUV(IOFF+1:IOFF+2*IUVLEV)=.TRUE.
    IOFF=IOFF+2*IUVLEV
    IOFF1=IOFF
    IOFFNS=IOFFNS+IOFF
    IOFFEW=IOFFEW+IOFF

    IOFF=IUVPAR*IUVLEV+KF_SCALARS_G
    IF(LUVDER) THEN
      IF(LSCDERS) IOFF=IOFF+KF_SCALARS_G
      DO J=1,IUVLEV
        IUVLEVS(IOFF+J)=J
        IUVPARS(IOFF+J)=IUVPAR+1
        LLUV(IOFF+J)=.TRUE.
        IUVLEVS(IOFF+J+IUVLEV)=J
        IUVPARS(IOFF+J+IUVLEV)=IUVPAR+2
        LLUV(IOFF+J+IUVLEV)=.TRUE.
      ENDDO
      IUVPAR=IUVPAR+2
      IOFF=IOFF+2*IUVLEV
      IOFFEW=IOFFEW+2*IUVLEV
    ENDIF
  ENDIF

  LLGP2(:)=.FALSE.
  IF (PRESENT(PGP2)) THEN
    IOFF=IOFF1
    IGP2PAR=UBOUND(PGP2,2)
    IF(LSCDERS) IGP2PAR=IGP2PAR/3
    DO J=1,IGP2PAR
      LLGP2(J+IOFF) = .TRUE.
      IGP2PARS(J+IOFF)=J
    ENDDO
    IOFF1=IOFF1+IGP2PAR
    IF(LSCDERS) THEN
      IOFF=IOFFNS
      DO J=1,IGP2PAR
        LLGP2(J+IOFF) = .TRUE.
        IGP2PARS(J+IOFF)=J+IGP2PAR
      ENDDO
      IOFFNS=IOFF+IGP2PAR
      IOFF=IOFFEW
      DO J=1,IGP2PAR
        LLGP2(J+IOFF) = .TRUE.
        IGP2PARS(J+IOFF)=J+2*IGP2PAR
      ENDDO
      IOFFEW=IOFF+IGP2PAR
    ENDIF
  ENDIF

  LLGP3A(:) = .FALSE.
  IF (PRESENT(PGP3A)) THEN
    IGP3ALEV=UBOUND(PGP3A,2)
    IGP3APAR=UBOUND(PGP3A,3)
    IF(LSCDERS) IGP3APAR=IGP3APAR/3
    IOFF=IOFF1
    DO J1=1,IGP3APAR
      DO J2=1,IGP3ALEV
        LLGP3A(J2+(J1-1)*IGP3ALEV+IOFF) = .TRUE.
        IGP3APARS(J2+(J1-1)*IGP3ALEV+IOFF)=J1
        IGP3ALEVS(J2+(J1-1)*IGP3ALEV+IOFF)=J2
      ENDDO
    ENDDO
    IPAROFF=IGP3APAR
    IOFF1=IOFF1+IGP3APAR*IGP3ALEV
    IF(LSCDERS) THEN
      IOFF=IOFFNS
      DO J1=1,IGP3APAR
        DO J2=1,IGP3ALEV
          LLGP3A(J2+(J1-1)*IGP3ALEV+IOFF) = .TRUE.
          IGP3APARS(J2+(J1-1)*IGP3ALEV+IOFF)=J1+IPAROFF
          IGP3ALEVS(J2+(J1-1)*IGP3ALEV+IOFF)=J2
        ENDDO
      ENDDO
      IPAROFF=IPAROFF+IGP3APAR
      IOFFNS=IOFFNS+IGP3APAR*IGP3ALEV
      IOFF=IOFFEW
      DO J1=1,IGP3APAR
        DO J2=1,IGP3ALEV
          LLGP3A(J2+(J1-1)*IGP3ALEV+IOFF) = .TRUE.
          IGP3APARS(J2+(J1-1)*IGP3ALEV+IOFF)=J1+IPAROFF
          IGP3ALEVS(J2+(J1-1)*IGP3ALEV+IOFF)=J2
        ENDDO
      ENDDO
      IOFFEW=IOFFEW+IGP3APAR*IGP3ALEV
    ENDIF
  ENDIF

  LLGP3B(:) = .FALSE.
  IF (PRESENT(PGP3B)) THEN
    IGP3BLEV=UBOUND(PGP3B,2)
    IGP3BPAR=UBOUND(PGP3B,3)
    IF(LSCDERS) IGP3BPAR=IGP3BPAR/3
    IOFF=IOFF1
    DO J1=1,IGP3BPAR
      DO J2=1,IGP3BLEV
        LLGP3B(J2+(J1-1)*IGP3BLEV+IOFF) = .TRUE.
        IGP3BPARS(J2+(J1-1)*IGP3BLEV+IOFF)=J1
        IGP3BLEVS(J2+(J1-1)*IGP3BLEV+IOFF)=J2
      ENDDO
    ENDDO
    IPAROFF=IGP3BPAR
    IOFF1=IOFF1+IGP3BPAR*IGP3BLEV
    IF(LSCDERS) THEN
      IOFF=IOFFNS
      DO J1=1,IGP3BPAR
        DO J2=1,IGP3BLEV
          LLGP3B(J2+(J1-1)*IGP3BLEV+IOFF) = .TRUE.
          IGP3BPARS(J2+(J1-1)*IGP3BLEV+IOFF)=J1+IPAROFF
          IGP3BLEVS(J2+(J1-1)*IGP3BLEV+IOFF)=J2
        ENDDO
      ENDDO
      IPAROFF=IPAROFF+IGP3BPAR
      IOFFNS=IOFFNS+IGP3BPAR*IGP3BLEV
      IOFF=IOFFEW
      DO J1=1,IGP3BPAR
        DO J2=1,IGP3BLEV
          LLGP3B(J2+(J1-1)*IGP3BLEV+IOFF) = .TRUE.
          IGP3BPARS(J2+(J1-1)*IGP3BLEV+IOFF)=J1+IPAROFF
          IGP3BLEVS(J2+(J1-1)*IGP3BLEV+IOFF)=J2
        ENDDO
      ENDDO
      IOFFEW=IOFFEW+IGP3BPAR*IGP3BLEV
    ENDIF
  ENDIF

  END ASSOCIATE
END SUBROUTINE TRGL_INIT_VARS

SUBROUTINE TRGL_INIT_OFF_VARS(YDBUFS,YLVARS,KVSET,KPTRGP,KF_GP)
  USE TPM_DISTR       ,ONLY : MYSETV, MYSETW
  USE TPM_TRANS       ,ONLY : NGPBLKS

  TYPE(TRGL_BUFFERS),  INTENT(INOUT) :: YDBUFS
  TYPE(TRGL_VARS), INTENT(INOUT) :: YLVARS
  INTEGER(KIND=JPIM), INTENT(IN) :: KVSET(:)
  INTEGER(KIND=JPIM) ,OPTIONAL, INTENT(IN) :: KPTRGP(:)
  INTEGER(KIND=JPIM), INTENT(IN) :: KF_GP

  !local variables
  INTEGER(KIND=JPIM) :: JFLD, IFIRST, ILAST,IPOS, JBLK

  ASSOCIATE(KGPTRSEND=>YDBUFS%IGPTRSEND, IFLDS=>YDBUFS%IFLDS, IFLDOFF=>YLVARS%IFLDOFF, &
    &       IGPTROFF=>YLVARS%IGPTROFF, LLINDER=>YDBUFS%LLINDER)

  IFLDS = 0
  DO JFLD=1,KF_GP
    IF(KVSET(JFLD) == MYSETV .OR. KVSET(JFLD) == -1) THEN
      IFLDS = IFLDS+1
      IF(LLINDER) THEN
        IFLDOFF(IFLDS) = KPTRGP(JFLD)
      ELSE
        IFLDOFF(IFLDS) = JFLD
      ENDIF
    ENDIF
  ENDDO

  IPOS=0
  DO JBLK=1,NGPBLKS
    IGPTROFF(JBLK)=IPOS
    IFIRST = KGPTRSEND(1,JBLK,MYSETW)
    IF(IFIRST > 0) THEN
      ILAST = KGPTRSEND(2,JBLK,MYSETW)
      IPOS=IPOS+ILAST-IFIRST+1
    ENDIF
  ENDDO

  END ASSOCIATE
END SUBROUTINE TRGL_INIT_OFF_VARS

SUBROUTINE TGRL_INIT_PACKING_VARS(YDBUFS,YLVARS, KVSET, KF_GP, PCOMBUFS)
  USE PARKIND1  ,ONLY : JPIM, JPRB
  USE TPM_TRANS       ,ONLY : NGPBLKS
  TYPE(TRGL_BUFFERS),  INTENT(INOUT) :: YDBUFS
  TYPE(TRGL_VARS), INTENT(INOUT) :: YLVARS

  INTEGER(KIND=JPIM), INTENT(IN) :: KVSET(:)
  INTEGER(KIND=JPIM), INTENT(IN) :: KF_GP
  REAL(KIND=JPRB), POINTER,INTENT(IN), OPTIONAL :: PCOMBUFS(:,:)
  !local variables
  INTEGER(KIND=JPIM) :: IFLD, IPOS, JFLD, IFIRST, ILAST, JBLK
  INTEGER(KIND=JPIM) :: KINRS, IV, ISETV, INRS

  ASSOCIATE(KGPTRSEND=>YDBUFS%IGPTRSEND, IPOSPLUS=>YDBUFS%IPOSPLUS, IJPOS=>YDBUFS%IJPOS, &
    &       IFLDA=>YDBUFS%IFLDA, ISETW=>YDBUFS%ISETW)

  IF (YDBUFS%LLTRGTOL) THEN
      KINRS = YDBUFS%INSEND
  ELSE
      KINRS = YDBUFS%INRECV
  ENDIF

  !$OMP PARALLEL DO SCHEDULE(STATIC) PRIVATE(INRS, IV, ISETV, JBLK, IFIRST, ILAST, IFLD, IPOS, JFLD)
  DO INRS=1,KINRS
    IF (YDBUFS%LLTRGTOL) THEN
      IV=YDBUFS%ISEND(INRS)
    ELSE
      IV=YDBUFS%IRECV(INRS)
    ENDIF
    YDBUFS%ISETW(INRS)=YDBUFS%ISETWL(IV)
    ISETV=YDBUFS%ISETVL(IV)

    IFLD = 0
    IPOS = 0
    IPOSPLUS(INRS)=0
    DO JFLD=1,KF_GP
      IF(KVSET(JFLD) == ISETV .OR. KVSET(JFLD) == -1 ) THEN
        IFLD = IFLD+1
        IFLDA(IFLD,INRS)=JFLD
      ENDIF
    ENDDO

    DO JBLK=1,NGPBLKS
      IFIRST = KGPTRSEND(1,JBLK,ISETW(INRS))
      IF(IFIRST > 0) THEN
        ILAST = KGPTRSEND(2,JBLK,ISETW(INRS))
        IJPOS(JBLK,INRS)=IPOS
        IPOSPLUS(INRS)=IPOSPLUS(INRS)+(ILAST-IFIRST+1)
        IPOS=IPOS+(ILAST-IFIRST+1)
      ENDIF
    ENDDO
    IF (PRESENT(PCOMBUFS)) THEN
      PCOMBUFS(-1,INRS) = 1
      PCOMBUFS(0,INRS) = IFLD
    ENDIF
  ENDDO
  !$OMP END PARALLEL DO

  END ASSOCIATE
END SUBROUTINE TGRL_INIT_PACKING_VARS

SUBROUTINE TGRL_COPY_ZCOMBUF(YDBUFS, YLVARS, INRS, ZCOMBUF, KPTRGP, PGP, PGPUV, PGP3A, PGP3B, PGP2)

  USE PARKIND1  ,ONLY : JPIM, JPRB
  USE TPM_TRANS       ,ONLY : NGPBLKS

  TYPE(TRGL_BUFFERS),  INTENT(INOUT) :: YDBUFS
  TYPE(TRGL_VARS), INTENT(INOUT) :: YLVARS
  INTEGER(KIND=JPIM), INTENT(IN) :: INRS
  REAL(KIND=JPRB), POINTER, CONTIGUOUS, INTENT(INOUT) :: ZCOMBUF(:,:)
  INTEGER(KIND=JPIM) ,OPTIONAL, INTENT(IN) :: KPTRGP(:)
  REAL(KIND=JPRB),OPTIONAL     :: PGP(:,:,:)
  REAL(KIND=JPRB),OPTIONAL     :: PGPUV(:,:,:,:)
  REAL(KIND=JPRB),OPTIONAL     :: PGP3A(:,:,:,:)
  REAL(KIND=JPRB),OPTIONAL     :: PGP3B(:,:,:,:)
  REAL(KIND=JPRB),OPTIONAL     :: PGP2(:,:,:)

  !Local variables
  INTEGER(KIND=JPIM) :: I_FLD_START,I_FLD_END
  INTEGER(KIND=JPIM) :: IFIRST, ILAST
  INTEGER(KIND=JPIM) :: JJ,JI,JK,IFLDT, JBLK, IPOS

  ASSOCIATE(IUVLEVS=>YLVARS%IUVLEVS, IUVPARS=>YLVARS%IUVPARS, IGP2PARS=>YLVARS%IGP2PARS, &
    &       LLUV=>YLVARS%LLUV, LLGP2=>YLVARS%LLGP2, LLGP3A=>YLVARS%LLGP3A, LLGP3B=>YLVARS%LLGP3B, &
    &       IGP3APARS=>YLVARS%IGP3APARS, IGP3ALEVS=>YLVARS%IGP3ALEVS, IGP3BPARS=>YLVARS%IGP3BPARS, &
    &       IGP3BLEVS=>YLVARS%IGP3BLEVS, KGPTRSEND =>YDBUFS%IGPTRSEND, IFLDA=>YDBUFS%IFLDA, &
    &       IPOSPLUS=>YDBUFS%IPOSPLUS, JPOS=>YDBUFS%IJPOS, ISETW=>YDBUFS%ISETW, &
    &       LLPGPONLY=>YDBUFS%LLPGPONLY, LLINDER=>YDBUFS%LLINDER)

  IPOS=IPOSPLUS(INRS)
  I_FLD_START = ZCOMBUF(-1,INRS)
  I_FLD_END   = ZCOMBUF(0,INRS)

  !$OMP PARALLEL DO SCHEDULE(STATIC) PRIVATE(IFLDT,IFIRST,ILAST,JK,JJ,JI,JBLK)
  DO JJ=I_FLD_START,I_FLD_END
    IFLDT=IFLDA(JJ,INRS)
    DO JBLK=1,NGPBLKS
      IFIRST = KGPTRSEND(1,JBLK,ISETW(INRS))
      IF(IFIRST > 0) THEN
        ILAST = KGPTRSEND(2,JBLK,ISETW(INRS))
        IF(LLINDER) THEN
          DO JK=IFIRST,ILAST
            JI=(JJ-I_FLD_START)*IPOS+JPOS(JBLK,INRS)+JK-IFIRST+1
            IF (YDBUFS%LLTRGTOL) THEN
              ZCOMBUF(JI,INRS) = PGP(JK,KPTRGP(IFLDT),JBLK)
            ELSE
              PGP(JK,KPTRGP(IFLDT),JBLK) = ZCOMBUF(JI,INRS)
            ENDIF
          ENDDO
        ELSEIF(LLPGPONLY) THEN
          DO JK=IFIRST,ILAST
            JI=(JJ-I_FLD_START)*IPOS+JPOS(JBLK,INRS)+JK-IFIRST+1
            IF (YDBUFS%LLTRGTOL) THEN
              ZCOMBUF(JI,INRS) = PGP(JK,IFLDT,JBLK)
            ELSE
              PGP(JK,IFLDT,JBLK) = ZCOMBUF(JI,INRS)
            ENDIF
          ENDDO
        ELSEIF(LLUV(IFLDT)) THEN
          DO JK=IFIRST,ILAST
            JI=(JJ-I_FLD_START)*IPOS+JPOS(JBLK,INRS)+JK-IFIRST+1
            IF (YDBUFS%LLTRGTOL) THEN
              ZCOMBUF(JI,INRS) = PGPUV(JK,IUVLEVS(IFLDT),IUVPARS(IFLDT),JBLK)
            ELSE
              PGPUV(JK,IUVLEVS(IFLDT),IUVPARS(IFLDT),JBLK) = ZCOMBUF(JI,INRS)
            ENDIF
          ENDDO
        ELSEIF(LLGP2(IFLDT)) THEN
          DO JK=IFIRST,ILAST
            JI=(JJ-I_FLD_START)*IPOS+JPOS(JBLK,INRS)+JK-IFIRST+1
            IF (YDBUFS%LLTRGTOL) THEN
              ZCOMBUF(JI,INRS) = PGP2(JK,IGP2PARS(IFLDT),JBLK)
            ELSE
              PGP2(JK,IGP2PARS(IFLDT),JBLK) = ZCOMBUF(JI,INRS)
            ENDIF
          ENDDO
        ELSEIF(LLGP3A(IFLDT)) THEN
          DO JK=IFIRST,ILAST
            JI=(JJ-I_FLD_START)*IPOS+JPOS(JBLK,INRS)+JK-IFIRST+1
            IF (YDBUFS%LLTRGTOL) THEN
              ZCOMBUF(JI,INRS) = PGP3A(JK,IGP3ALEVS(IFLDT),IGP3APARS(IFLDT),JBLK)
            ELSE
              PGP3A(JK,IGP3ALEVS(IFLDT),IGP3APARS(IFLDT),JBLK) = ZCOMBUF(JI,INRS)
            ENDIF
          ENDDO
        ELSEIF(LLGP3B(IFLDT)) THEN
          DO JK=IFIRST,ILAST
            JI=(JJ-I_FLD_START)*IPOS+JPOS(JBLK,INRS)+JK-IFIRST+1
            IF (YDBUFS%LLTRGTOL) THEN
              ZCOMBUF(JI,INRS) = PGP3B(JK,IGP3BLEVS(IFLDT),IGP3BPARS(IFLDT),JBLK)
            ELSE
              PGP3B(JK,IGP3BLEVS(IFLDT),IGP3BPARS(IFLDT),JBLK) = ZCOMBUF(JI,INRS)
            ENDIF
          ENDDO
        ENDIF
      ENDIF
    ENDDO
  ENDDO
  !$OMP END PARALLEL DO

  END ASSOCIATE
END SUBROUTINE TGRL_COPY_ZCOMBUF


SUBROUTINE TGRL_COPY_PGLAT(PGLAT, YDBUFS, YLVARS, PGP, PGPUV, PGP3A, PGP3B, PGP2)

  USE PARKIND1  ,ONLY : JPIM, JPRB, JPIB
  USE TPM_DISTR       ,ONLY : MYSETW, MYPROC
  USE TPM_GEN         ,ONLY : NOUT
  USE ABORT_TRANS_MOD ,ONLY : ABORT_TRANS
  USE TPM_TRANS       ,ONLY : NGPBLKS

  REAL(KIND=JPRB),OPTIONAL   :: PGLAT(:,:)
  TYPE(TRGL_BUFFERS),  INTENT(INOUT) :: YDBUFS
  TYPE(TRGL_VARS), INTENT(INOUT) :: YLVARS
  REAL(KIND=JPRB),OPTIONAL    :: PGP(:,:,:)
  REAL(KIND=JPRB),OPTIONAL    :: PGPUV(:,:,:,:)
  REAL(KIND=JPRB),OPTIONAL    :: PGP3A(:,:,:,:)
  REAL(KIND=JPRB),OPTIONAL    :: PGP3B(:,:,:,:)
  REAL(KIND=JPRB),OPTIONAL    :: PGP2(:,:,:)

    !Local variables
  INTEGER(KIND=JPIM) :: IFIRST, ILAST, IFLD, IPOS, JBLK, JK
  INTEGER(KIND=JPIB) :: JFLD64

  ASSOCIATE(IUVLEVS=>YLVARS%IUVLEVS,IFLDOFF=>YLVARS%IFLDOFF, IGPTROFF=>YLVARS%IGPTROFF, &
    &       IUVPARS=>YLVARS%IUVPARS, IGP2PARS=>YLVARS%IGP2PARS, LLUV=>YLVARS%LLUV, &
    &       LLGP2=>YLVARS%LLGP2, LLGP3A=>YLVARS%LLGP3A, LLGP3B=>YLVARS%LLGP3B, &
    &       IGP3APARS=>YLVARS%IGP3APARS, IGP3ALEVS=>YLVARS%IGP3ALEVS, &
    &       IGP3BPARS=>YLVARS%IGP3BPARS, IGP3BLEVS=>YLVARS%IGP3BLEVS, KINDEX=>YDBUFS%IINDEX, &
    &       KNDOFF=>YDBUFS%INDOFF, KGPTRSEND =>YDBUFS%IGPTRSEND, IFLDS=>YDBUFS%IFLDS, &
    &       LLPGPONLY=>YDBUFS%LLPGPONLY)

#ifdef __NEC__
  ! Loops inversion is still better on Aurora machines, according to CHMI. REK.
  !$OMP PARALLEL DO SCHEDULE(DYNAMIC) PRIVATE(JFLD64,JBLK,JK,IFLD,IPOS,IFIRST,ILAST)
#else
  !$OMP PARALLEL DO SCHEDULE(STATIC) PRIVATE(JFLD64,JBLK,JK,IFLD,IPOS,IFIRST,ILAST)
#endif
  DO JBLK=1,NGPBLKS
    IFIRST = KGPTRSEND(1,JBLK,MYSETW)
    IF(IFIRST > 0) THEN
      ILAST = KGPTRSEND(2,JBLK,MYSETW)
      ! Address PGLAT over 64 bits because its size may exceed 2 GB for big data and
      ! small number of tasks.
      IF(LLPGPONLY) THEN
        DO JFLD64=1,IFLDS
          IFLD = IFLDOFF(JFLD64)
          !DIR$ VECTOR ALWAYS
          DO JK=IFIRST,ILAST
            IPOS = KNDOFF(MYPROC)+IGPTROFF(JBLK)+JK-IFIRST+1
            IF (YDBUFS%LLTRGTOL) THEN
              PGLAT(JFLD64,KINDEX(IPOS)) = PGP(JK,IFLD,JBLK)
            ELSE
              PGP(JK,IFLD,JBLK) = PGLAT(JFLD64,KINDEX(IPOS))
            ENDIF
          ENDDO
        ENDDO
      ELSE
        DO JFLD64=1,IFLDS
          IFLD = IFLDOFF(JFLD64)
          IF(LLUV(IFLD)) THEN
            !DIR$ VECTOR ALWAYS
            DO JK=IFIRST,ILAST
              IPOS = KNDOFF(MYPROC)+IGPTROFF(JBLK)+JK-IFIRST+1
              IF (YDBUFS%LLTRGTOL) THEN
                PGLAT(JFLD64,KINDEX(IPOS)) = PGPUV(JK,IUVLEVS(IFLD),IUVPARS(IFLD),JBLK)
              ELSE
                PGPUV(JK,IUVLEVS(IFLD),IUVPARS(IFLD),JBLK) = PGLAT(JFLD64,KINDEX(IPOS))
              ENDIF
            ENDDO
          ELSEIF(LLGP2(IFLD)) THEN
            !DIR$ VECTOR ALWAYS
            DO JK=IFIRST,ILAST
              IPOS = KNDOFF(MYPROC)+IGPTROFF(JBLK)+JK-IFIRST+1
              IF (YDBUFS%LLTRGTOL) THEN
                PGLAT(JFLD64,KINDEX(IPOS)) = PGP2(JK,IGP2PARS(IFLD),JBLK)
              ELSE
                PGP2(JK,IGP2PARS(IFLD),JBLK)=PGLAT(JFLD64,KINDEX(IPOS))
              ENDIF
            ENDDO
          ELSEIF(LLGP3A(IFLD)) THEN
            !DIR$ VECTOR ALWAYS
            DO JK=IFIRST,ILAST
              IPOS = KNDOFF(MYPROC)+IGPTROFF(JBLK)+JK-IFIRST+1
              IF (YDBUFS%LLTRGTOL) THEN
                PGLAT(JFLD64,KINDEX(IPOS)) = PGP3A(JK,IGP3ALEVS(IFLD),IGP3APARS(IFLD),JBLK)
              ELSE
                PGP3A(JK,IGP3ALEVS(IFLD),IGP3APARS(IFLD),JBLK)=PGLAT(JFLD64,KINDEX(IPOS))
              ENDIF
            ENDDO
          ELSEIF(LLGP3B(IFLD)) THEN
            !DIR$ VECTOR ALWAYS
            DO JK=IFIRST,ILAST
              IPOS = KNDOFF(MYPROC)+IGPTROFF(JBLK)+JK-IFIRST+1
              IF (YDBUFS%LLTRGTOL) THEN
                PGLAT(JFLD64,KINDEX(IPOS)) = PGP3B(JK,IGP3BLEVS(IFLD),IGP3BPARS(IFLD),JBLK)
              ELSE
                PGP3B(JK,IGP3BLEVS(IFLD),IGP3BPARS(IFLD),JBLK)=PGLAT(JFLD64,KINDEX(IPOS))
              ENDIF
            ENDDO
          ELSE
            WRITE(NOUT,*)'TRGTOL_MOD: ERROR',JFLD64,IFLD
            CALL ABORT_TRANS('TRGTOL_MOD: ERROR')
          ENDIF
        ENDDO
      ENDIF
    ENDIF
  ENDDO
  !$OMP END PARALLEL DO

  END ASSOCIATE
END SUBROUTINE TGRL_COPY_PGLAT


SUBROUTINE TRGL_PROLOG(KF_FS,KF_GP,KVSET,YDBUFS)

  USE PARKIND1  ,ONLY : JPIM

  USE TPM_DISTR       ,ONLY : D, MYSETW, NPRTRNS, MYPROC, NPROC

  USE INIGPTR_MOD     ,ONLY : INIGPTR
  USE PE2SET_MOD      ,ONLY : PE2SET

  IMPLICIT NONE

  INTEGER(KIND=JPIM),INTENT(IN) :: KF_FS,KF_GP
  TYPE (TRGL_BUFFERS), INTENT(INOUT) :: YDBUFS
  INTEGER(KIND=JPIM),INTENT(IN) :: KVSET(KF_GP)
  INTEGER(KIND=JPIM) :: IGPTRRECV(NPRTRNS)
  INTEGER(KIND=JPIM) :: IFIRSTLAT, IGL, IGLL, ILASTLAT, IPOS, ISETA, ISETB, ISETV
  INTEGER(KIND=JPIM) :: JFLD, JGL, JL, ISETW, JROC, J
  INTEGER(KIND=JPIM) :: INDOFFX

  !     ------------------------------------------------------------------
  !*       0.    Some initializations
  !              --------------------

  CALL INIGPTR(YDBUFS%IGPTRSEND,IGPTRRECV)

  INDOFFX  = 0
  YDBUFS%INRECV = 0
  YDBUFS%INSEND = 0

  DO JROC=1,NPROC

    CALL PE2SET(JROC,ISETA,ISETB,YDBUFS%ISETWL(JROC),YDBUFS%ISETVL(JROC))

    ISETW=YDBUFS%ISETWL(JROC)
    ISETV=YDBUFS%ISETVL(JROC)

    ! Count up expected number of fields
    IPOS = COUNT(KVSET == ISETV .OR. KVSET == -1)

    IF (YDBUFS%LLTRGTOL) THEN
      YDBUFS%ISENDTOT(JROC) = IGPTRRECV(ISETW)*IPOS
      IF( JROC /= MYPROC) THEN
        IF(YDBUFS%ISENDTOT(JROC) > 0) THEN
          YDBUFS%INSEND = YDBUFS%INSEND+1
          YDBUFS%ISEND(YDBUFS%INSEND)=JROC
        ENDIF
      ENDIF
    ELSE
      YDBUFS%IRECVTOT(JROC) = IGPTRRECV(ISETW)*IPOS
      IF(YDBUFS%IRECVTOT(JROC) > 0 .AND. MYPROC /= JROC) THEN
        YDBUFS%INRECV = YDBUFS%INRECV + 1
        YDBUFS%IRECV(YDBUFS%INRECV)=JROC
      ENDIF
    ENDIF
    IFIRSTLAT = MAX(D%NPTRLS(MYSETW),D%NFRSTLAT(ISETA))
    ILASTLAT  = MIN(D%NPTRLS(MYSETW)+D%NULTPP(MYSETW)-1,D%NLSTLAT(ISETA))

    IPOS = 0
    DO JGL=IFIRSTLAT,ILASTLAT
      IGL  = D%NPTRFRSTLAT(ISETA)+JGL-D%NFRSTLAT(ISETA)
      IPOS = IPOS+D%NONL(IGL,ISETB)
    ENDDO

    IF (YDBUFS%LLTRGTOL) THEN
      YDBUFS%IRECVTOT(JROC) = IPOS*KF_FS
      IF(YDBUFS%IRECVTOT(JROC) > 0 .AND. MYPROC /= JROC) THEN
        YDBUFS%INRECV = YDBUFS%INRECV + 1
        YDBUFS%IRECV(YDBUFS%INRECV)=JROC
      ENDIF
    ELSE
      YDBUFS%ISENDTOT(JROC) = IPOS*KF_FS
      IF( JROC /= MYPROC) THEN
        IF(YDBUFS%ISENDTOT(JROC) > 0) THEN
          YDBUFS%INSEND = YDBUFS%INSEND+1
          YDBUFS%ISEND(YDBUFS%INSEND)=JROC
        ENDIF
      ENDIF
    ENDIF

    IF(IPOS > 0) THEN
      YDBUFS%INDOFF(JROC) = INDOFFX
      INDOFFX = INDOFFX+IPOS
      IPOS = 0
      DO JGL=IFIRSTLAT,ILASTLAT
        IGL  = D%NPTRFRSTLAT(ISETA)+JGL-D%NFRSTLAT(ISETA)
        IGLL = JGL-D%NPTRLS(MYSETW)+1
        DO JL=D%NSTA(IGL,ISETB)+D%NSTAGTF(IGLL),&
        &D%NSTA(IGL,ISETB)+D%NSTAGTF(IGLL)+D%NONL(IGL,ISETB)-1
          IPOS = IPOS+1
          YDBUFS%IINDEX(IPOS+YDBUFS%INDOFF(JROC)) = JL
        ENDDO
      ENDDO
    ENDIF
  ENDDO

  YDBUFS%ISENDCOUNT = MAXVAL(YDBUFS%ISENDTOT)
  YDBUFS%IRECVCOUNT = MAXVAL(YDBUFS%IRECVTOT)

END SUBROUTINE TRGL_PROLOG

END MODULE TRGL_MOD
