! (C) Copyright 2026- ECMWF.
! 
! 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 UTILS

USE PARKIND1, ONLY: JPRB, JPIM

IMPLICIT NONE

PRIVATE
PUBLIC :: DETECT_MPIRUN, SCALPRODSP, SCALPRODGP

CONTAINS

FUNCTION DETECT_MPIRUN() RESULT(LMPI_REQUIRED)
  USE EC_ENV_MOD, ONLY : EC_PUTENV
  LOGICAL :: LMPI_REQUIRED
  INTEGER :: ILEN
  INTEGER, PARAMETER :: NVARS = 4
  CHARACTER(LEN=32), DIMENSION(NVARS) :: CMPIRUN_DETECT
  CHARACTER(LEN=4) :: CLENV
  INTEGER :: IVAR

  ! Environment variables that are set when mpirun, srun, aprun, ... are used
  CMPIRUN_DETECT(1) = 'OMPI_COMM_WORLD_SIZE'  ! OPENMPI
  CMPIRUN_DETECT(2) = 'ALPS_APP_PE'           ! CRAY PE
  CMPIRUN_DETECT(3) = 'PMI_SIZE'              ! INTEL
  CMPIRUN_DETECT(4) = 'SLURM_NTASKS'          ! SLURM

  LMPI_REQUIRED = .FALSE.
  DO IVAR = 1, NVARS
    CALL GET_ENVIRONMENT_VARIABLE(NAME=TRIM(CMPIRUN_DETECT(IVAR)), LENGTH=ILEN)
    IF (ILEN > 0) THEN
      LMPI_REQUIRED = .TRUE.
      EXIT ! Break
    ENDIF
  ENDDO

  CALL GET_ENVIRONMENT_VARIABLE(NAME="ECTRANS_USE_MPI", VALUE=CLENV, LENGTH=ILEN )
  IF (ILEN > 0) THEN
      LMPI_REQUIRED = .TRUE.
      IF( TRIM(CLENV) == "0" .OR. TRIM(CLENV) == "OFF" .OR. TRIM(CLENV) == "off" .OR. TRIM(CLENV) == "F" ) THEN
        LMPI_REQUIRED = .FALSE.
      ENDIF
      CALL EC_PUTENV("DR_HOOK_ASSERT_MPI_INITIALIZED=0", OVERWRITE=.TRUE.)
  ENDIF
END FUNCTION DETECT_MPIRUN

FUNCTION SCALPRODSP(PSP1, PSP2, KVSET, KLEVL, KLEVG, KSPEC2, KSPEC2G, KSMAX, KMYPROC) RESULT(PSC)
  ! Scalar product in spectral space
  REAL(KIND=JPRB), INTENT(IN) :: PSP1(:,:)
  REAL(KIND=JPRB), INTENT(IN) :: PSP2(:,:)
  INTEGER(KIND=JPIM), INTENT(IN) :: KVSET(:)
  INTEGER(KIND=JPIM), INTENT(IN) :: KLEVL
  INTEGER(KIND=JPIM), INTENT(IN) :: KLEVG
  INTEGER(KIND=JPIM), INTENT(IN) :: KSPEC2
  INTEGER(KIND=JPIM), INTENT(IN) :: KSPEC2G
  INTEGER(KIND=JPIM), INTENT(IN) :: KSMAX
  INTEGER(KIND=JPIM), INTENT(IN) :: KMYPROC

  REAL(KIND=JPRB) :: PSC

  INTEGER(KIND=JPIM), ALLOCATABLE :: MYMS(:), NASM0(:)
  INTEGER(KIND=JPIM) :: JMLOC, IM, JIR, JN, INM, JLEV, NUMP, I
  REAL(KIND=JPRB) :: ZMFACT, ZSP(KLEVL,KSPEC2), ZSPG(KLEVG,KSPEC2G)

#include "trans_inq.h"
#include "gath_spec.h"

  ! Get Ms I'm responsible for (MYMS)
  CALL TRANS_INQ(KNUMP=NUMP)
  ALLOCATE(MYMS(NUMP))
  ALLOCATE(NASM0(0:KSMAX))
  CALL TRANS_INQ(KMYMS=MYMS, KASM0=NASM0)

  PSC = 0.0_JPRB
  ZSP(:,:) = 0.0_JPRB

  !$OMP PARALLEL DO SCHEDULE(STATIC,1) PRIVATE(JLEV,JMLOC,IM,ZMFACT,JIR,JN,INM)
  DO JLEV = 1, KLEVL
    DO JMLOC = 1, NUMP
      IM = MYMS(JMLOC)
      ZMFACT = 1.0_JPRB + REAL(MIN(1, IM), JPRB)
      DO JIR = 0, MIN(1, IM)
        DO JN = IM, KSMAX
          INM = NASM0(IM) + (JN - IM) * 2 + JIR
          ZSP(JLEV,INM) = PSP1(JLEV,INM) * PSP2(JLEV,INM) * ZMFACT
        ENDDO
      ENDDO
    ENDDO
  ENDDO
  !$OMP END PARALLEL DO

  CALL GATH_SPEC(PSPECG=ZSPG, KFGATHG=KLEVG, KTO=(/ (1, I = 1, KLEVG) /), PSPEC=ZSP, KVSET=KVSET)

  IF (KMYPROC == 1) THEN
    PSC = SUM(ZSPG)
  ELSE
    PSC = 0.0_JPRB
  ENDIF
END FUNCTION SCALPRODSP

FUNCTION SCALPRODGP(RGP1, RGP2, KPROMA, KFIELD, KGPBLKS, KGPTOT, KGPTOTG, KMYPROC) RESULT(RSC)

  ! Scalar product in spectral space
  INTEGER(KIND=JPIM), INTENT(IN) :: KPROMA
  INTEGER(KIND=JPIM), INTENT(IN) :: KFIELD
  INTEGER(KIND=JPIM), INTENT(IN) :: KGPBLKS
  REAL(KIND=JPRB), INTENT(IN) :: RGP1(KPROMA,KFIELD,KGPBLKS)
  REAL(KIND=JPRB), INTENT(IN) :: RGP2(KPROMA,KFIELD,KGPBLKS)
  INTEGER(KIND=JPIM), INTENT(IN) :: KGPTOT
  INTEGER(KIND=JPIM), INTENT(IN) :: KGPTOTG
  INTEGER(KIND=JPIM), INTENT(IN) :: KMYPROC

  REAL(KIND=JPRB) :: RSC

  INTEGER(KIND=JPIM) :: JLEV, JKGLO, IEND, IBL, JROF, I
  REAL(KIND=JPRB) :: RGP(KPROMA,KFIELD,KGPBLKS), RGPG(KGPTOTG,KFIELD)

#include "gath_grid.h"

  RSC = 0.0_JPRB

  !$OMP PARALLEL DO SCHEDULE(STATIC,1) PRIVATE(JLEV,JKGLO,IEND,IBL,JROF)
  DO JLEV = 1, KFIELD
    DO JKGLO = 1, KGPTOT, KPROMA
      IEND = MIN(KPROMA, KGPTOT - JKGLO + 1)
      IBL  = (JKGLO - 1) / KPROMA+1
      DO JROF = 1, IEND
        RGP(JROF,JLEV,IBL) = RGP1(JROF,JLEV,IBL) * RGP2(JROF,JLEV,IBL)
      ENDDO
    ENDDO
  ENDDO
  !$OMP END PARALLEL DO

  CALL GATH_GRID(RGPG, KPROMA, KFIELD, KTO=[ (1, I = 1, KFIELD) ], PGP=RGP)

  IF (KMYPROC == 1) THEN
    RSC = SUM(RGPG)
  ELSE
    RSC = 0.0_JPRB
  ENDIF

END FUNCTION SCALPRODGP

END MODULE UTILS
