MODULE openacc_mod
  
  LOGICAL,SAVE :: openacc_initialized=.FALSE.
  !$OMP THREADPRIVATE(openacc_initialized)
  
  INTEGER,SAVE :: openacc_device_id=-1
  !$OMP THREADPRIVATE(openacc_device_id)

  INTERFACE set_openacc_device
    MODULE PROCEDURE set_device_manual, set_device_mpi
  END INTERFACE
  
CONTAINS

#ifdef _OPENACC
  SUBROUTINE set_device_manual()
    USE openacc
    IMPLICIT NONE
    
    CHARACTER(len=32) :: device_id_value, slurm_local_id_value
    INTEGER :: local_id, numdev, device_id_ierr, slurm_local_id_ierr, mydev
    
    IF (openacc_initialized) RETURN
    
    ! Use user set DEVICE_ID, then slurm set SLURM_LOCAL_ID, then fallback to default value
    CALL get_environment_variable("DEVICE_ID", value=device_id_value, status=device_id_ierr) 
    CALL get_environment_variable("SLURM_LOCALID", value=slurm_local_id_value, status=slurm_local_id_ierr) 
    IF (device_id_ierr == 0) then
       READ(device_id_value,*) local_id
       PRINT *, "setDevice : found env variable DEVICE_ID =", local_id
    ELSE IF (slurm_local_id_ierr == 0) then
       READ(slurm_local_id_value,*) local_id
       PRINT *, "setDevice : found env variable SLURM_LOCAL_ID =", local_id 
    ELSE
     RETURN
    END IF

    ! get the number of device on this node
    numdev = acc_get_num_devices(ACC_DEVICE_NVIDIA)
    WRITE( *, '("local_id=",i3," numdev=",i3)') local_id, numdev

    IF (numdev < 1) then
       PRINT *, "Error: there are no devices available on this host. ABORTING"
       STOP
    END IF

    ! print a warning if the number of devices is less than the number of processes on this node. Having multiple processes share a devices is not recommended
    IF (numdev == local_id) then
       ! print warning message only once per node
       PRINT *, "WARNING: the number of process is greater than the number of GPUs."
       mydev = mod(local_id, numdev)
    ELSE
       mydev = local_id
    END IF
  
    CALL acc_init(ACC_DEVICE_NVIDIA)
    CALL acc_set_device_num(mydev,ACC_DEVICE_NVIDIA)
    openacc_device_id = acc_get_device_num(ACC_DEVICE_NVIDIA)
    openacc_initialized=.TRUE.
    
  END SUBROUTINE set_device_manual

#else
  SUBROUTINE set_device_manual()
  IMPLICIT NONE
     
  END SUBROUTINE set_device_manual
#endif

#ifdef _OPENACC
  SUBROUTINE set_device_mpi(nprocs, myrank)
    USE iso_c_binding
    USE openacc
    USE mpi_mod
    IMPLICIT NONE
    
    INTERFACE 
       FUNCTION gethostid() bind(C)
         USE iso_c_binding
         INTEGER(C_INT) :: gethostid
       END FUNCTION gethostid
    END INTERFACE
  
    INTEGER, INTENT(in) :: nprocs, myrank
    INTEGER :: hostids(nprocs), localprocs(nprocs)
    INTEGER :: hostid, ierr, numdev, mydev, i, numlocal

    IF (openacc_initialized) RETURN

    ! get the hostids so we can determine what other processes are on this node
    hostid = gethostid()
    CALL mpi_allgather(hostid,1,MPI_INTEGER,hostids,1,MPI_INTEGER, MPI_COMM_WORLD, ierr)

    ! determine which processors are on this node 
    numlocal = 0
    localprocs(:) = 0
    DO i = 1, nprocs
       IF (hostid == hostids(i)) THEN
          localprocs(i) = numlocal
          numlocal = numlocal + 1
       END IF
    END DO

    ! get the number of device on this node
    numdev = acc_get_num_devices(ACC_DEVICE_NVIDIA)

    IF (numdev < 1) THEN
      PRINT *, "Error: there are no devices available on this host. ABORTING", myrank
      STOP
    END IF

    ! print a warning if the number of devices is less than the number of processes on this node. Having multiple processes share a devices is not recommended
    IF (numdev < numlocal) THEN
       IF (localprocs(myrank+1) == 1) THEN
          ! print warning message only once per node
          PRINT *, "WARNING: the number of process is greater than the number of GPUs.", myrank
       END IF
       mydev = mod(localprocs(myrank+1), numdev)
    ELSE
       mydev = localprocs(myrank+1)
    END IF

    CALL acc_set_device_num(mydev,ACC_DEVICE_NVIDIA)
    CALL acc_init(ACC_DEVICE_NVIDIA)
    openacc_device_id = acc_get_device_num(ACC_DEVICE_NVIDIA)
    openacc_initialized=.TRUE.

  END SUBROUTINE set_device_mpi

#else 

  SUBROUTINE set_device_mpi(nprocs, myrank)
  IMPLICIT NONE
    INTEGER, INTENT(in) :: nprocs, myrank
  END SUBROUTINE set_device_mpi

#endif  

END MODULE openacc_mod