MODULE omp_para

  INTEGER,SAVE :: omp_size
  INTEGER,SAVE :: omp_rank
!$OMP THREADPRIVATE(omp_rank)

  LOGICAL,SAVE :: is_omp_first_level
  LOGICAL,SAVE :: is_omp_last_level
  LOGICAL,SAVE :: is_omp_master
!$OMP THREADPRIVATE(is_omp_first_level, is_omp_last_level,is_omp_master)

  INTEGER,SAVE :: ll_begin
  INTEGER,SAVE :: ll_beginp1
  INTEGER,SAVE :: ll_end
  INTEGER,SAVE :: ll_endm1
  INTEGER,SAVE :: ll_endp1
!$OMP THREADPRIVATE(ll_begin,ll_beginp1,ll_end,ll_endm1,ll_endp1)
  LOGICAL,SAVE :: using_openmp
  
  INTEGER,SAVE :: omp_domain_size
  INTEGER,SAVE :: omp_domain_rank
  INTEGER,SAVE :: omp_level_size
  INTEGER,SAVE :: omp_level_rank
!$OMP THREADPRIVATE( omp_domain_size, omp_level_size,omp_domain_rank,omp_level_rank)
  LOGICAL,SAVE :: is_omp_domain_master
  LOGICAL,SAVE :: is_omp_level_master
!$OMP THREADPRIVATE(is_omp_domain_master,is_omp_level_master )
  
  LOGICAL,PARAMETER :: omp_by_domain=.TRUE. 
  LOGICAL,SAVE :: is_master
!$OMP THREADPRIVATE(is_master)


   LOGICAL,SAVE :: is_omp_first_level_full
   LOGICAL,SAVE :: is_omp_last_level_full
   INTEGER,SAVE :: ll_begin_full
   INTEGER,SAVE :: ll_beginp1_full
   INTEGER,SAVE :: ll_end_full
   INTEGER,SAVE :: ll_endm1_full
   INTEGER,SAVE :: ll_endp1_full
!$OMP THREADPRIVATE(is_omp_first_level_full,is_omp_last_level_full)   
!$OMP THREADPRIVATE( ll_begin_full, ll_beginp1_full, ll_end_full, ll_endm1_full, ll_endp1_full)
  PRIVATE :: is_omp_first_level_full,is_omp_last_level_full
  PRIVATE :: ll_begin_full, ll_beginp1_full, ll_end_full, ll_endm1_full, ll_endp1_full 
  

   LOGICAL,SAVE :: is_omp_first_level_distrib
   LOGICAL,SAVE :: is_omp_last_level_distrib
   INTEGER,SAVE :: ll_begin_distrib
   INTEGER,SAVE :: ll_beginp1_distrib
   INTEGER,SAVE :: ll_end_distrib
   INTEGER,SAVE :: ll_endm1_distrib
   INTEGER,SAVE :: ll_endp1_distrib
!$OMP THREADPRIVATE(is_omp_first_level_distrib,is_omp_last_level_distrib)   
!$OMP THREADPRIVATE( ll_begin_distrib, ll_beginp1_distrib, ll_end_distrib, ll_endm1_distrib, ll_endp1_distrib)

  PRIVATE :: is_omp_first_level_distrib,is_omp_last_level_distrib
  PRIVATE :: ll_begin_distrib, ll_beginp1_distrib, ll_end_distrib, ll_endm1_distrib, ll_endp1_distrib 

CONTAINS

  FUNCTION get_omp_size() result(omp_size)
#ifdef CPP_USING_OMP
    use omp_lib, only : omp_get_max_threads
    integer :: omp_size
    omp_size = omp_get_max_threads()
#else
    integer :: omp_size
    omp_size = 1
#endif
  END FUNCTION

  SUBROUTINE init_omp_para(is_mpi_master)
  USE grid_param
  USE ioipsl, ONLY : getin
#ifdef CPP_USING_OMP
  USE omp_lib
#endif
  IMPLICIT NONE
  LOGICAL, INTENT(IN) :: is_mpi_master
  INTEGER :: ll_nb,i,llb,lle

#ifdef CPP_USING_OMP
  using_openmp=.TRUE.
#else
  using_openmp=.FALSE.
#endif

  IF (using_openmp) THEN    
!$OMP PARALLEL PRIVATE(ll_nb,i,llb,lle)
  
!$OMP MASTER
#ifdef CPP_USING_OMP
    omp_size=OMP_GET_NUM_THREADS()
#endif
!$OMP END MASTER
!$OMP BARRIER
#ifdef CPP_USING_OMP
    omp_rank=OMP_GET_THREAD_NUM()
#endif
    
    is_omp_master=.FALSE.
    is_master=.FALSE.
      
    IF (omp_rank==0) THEN
      is_omp_master=.TRUE.
      IF (is_mpi_master) is_master=.TRUE.
    ENDIF

    !$OMP CRITICAL
    omp_level_size=1
    CALL getin("omp_level_size",omp_level_size)
    !$OMP END CRITICAL

    IF(is_mpi_master) PRINT *,'GETIN omp_level_size', ' = ', omp_level_size

    IF (MOD(omp_size,omp_level_size)/=0) THEN
      IF (is_mpi_master) PRINT*,"omp_size /= omp_level_size x omp_domain_size  => disable omp threads on vertical layers"
      omp_level_size=1
    ENDIF
    omp_domain_size=omp_size/omp_level_size
    omp_domain_rank = omp_rank / omp_level_size    
    omp_level_rank = MOD(omp_rank, omp_level_size)    
    
    IF (is_mpi_master) PRINT*,"omp_domain_size",omp_domain_size,"omp_domain_rank", omp_domain_rank
    IF (is_mpi_master) PRINT*,"omp_level_size",omp_level_size,"omp_level_rank", omp_level_rank
    
    is_omp_first_level=.FALSE.
    is_omp_last_level= .FALSE.
    is_omp_domain_master=.FALSE.
    is_omp_level_master=.FALSE.

    IF (omp_domain_rank==0) is_omp_domain_master = .TRUE.
    IF (omp_level_rank==0) is_omp_level_master = .TRUE.
    IF (omp_level_rank==0) is_omp_first_level=.TRUE.
    
    IF (omp_level_rank==omp_level_size-1) is_omp_last_level=.TRUE.
    
    lle=0
    
    DO i=0,omp_level_rank
      llb=lle+1
      ll_nb=llm/omp_level_size
      IF (MOD(llm,omp_level_size)>i) ll_nb=ll_nb+1
      lle=llb+ll_nb-1
    ENDDO
    ll_begin=llb
    ll_end=lle
    
    ll_beginp1=ll_begin
    ll_endp1=ll_end
    ll_endm1=ll_end

    IF (is_omp_first_level) ll_beginp1=ll_begin+1
    IF (is_omp_last_level) ll_endp1=ll_endp1+1
    IF (is_omp_last_level) ll_endm1=ll_endm1-1
    
    
    
    is_omp_first_level_distrib = is_omp_first_level
    is_omp_last_level_distrib  = is_omp_last_level
    ll_begin_distrib           = ll_begin
    ll_beginp1_distrib         = ll_beginp1
    ll_end_distrib             = ll_end
    ll_endm1_distrib           = ll_endm1
    ll_endp1_distrib           = ll_endp1
    
    is_omp_first_level_full = .TRUE.
    is_omp_last_level_full  = .TRUE.
    ll_begin_full           = 1
    ll_beginp1_full         = 2
    ll_end_full             = llm
    ll_endm1_full           = llm-1
    ll_endp1_full           = llm+1    

!$OMP END PARALLEL

   ELSE
     omp_size=1
     omp_level_size=1
     omp_domain_size=1
     omp_rank=0
     omp_level_rank=0
     omp_domain_rank=0
     is_master=is_mpi_master
     is_omp_first_level=.TRUE.
     is_omp_last_level=.TRUE.
     is_omp_master=.TRUE.
     is_omp_domain_master=.TRUE.
     is_omp_level_master=.TRUE.
     ll_begin=1
     ll_beginp1=2
     ll_end=llm
     ll_endm1=llm-1
     ll_endp1=llm+1
     
     is_omp_first_level_distrib = is_omp_first_level
     is_omp_last_level_distrib  = is_omp_last_level
     ll_begin_distrib           = ll_begin
     ll_beginp1_distrib         = ll_beginp1
     ll_end_distrib             = ll_end
     ll_endm1_distrib           = ll_endm1
     ll_endp1_distrib           = ll_endp1
    
     is_omp_first_level_full = .TRUE.
     is_omp_last_level_full  = .TRUE.
     ll_begin_full           = 1
     ll_beginp1_full         = 2
     ll_end_full             = llm
     ll_endm1_full           = llm-1
     ll_endp1_full           = llm+1    

   ENDIF
   
  END SUBROUTINE init_omp_para

  SUBROUTINE distrib_level(ibegin,iend, lbegin,lend)
  IMPLICIT NONE
    INTEGER,INTENT(IN)  :: ibegin,iend  
    INTEGER,INTENT(OUT) :: lbegin  
    INTEGER,INTENT(OUT) :: lend  
    INTEGER :: size,div,rest
    size=iend-ibegin+1
    div=size/omp_level_size
    rest=MOD(size,omp_level_size)
    IF (omp_level_rank<rest) THEN
      lbegin=(div+1)*omp_level_rank + ibegin
      lend=lbegin+div
    ELSE
      lbegin=(div+1)*rest + (omp_level_rank-rest)*div + ibegin
      lend=lbegin+div-1
    ENDIF
  END SUBROUTINE distrib_level


  SUBROUTINE switch_omp_distrib_level
  IMPLICIT NONE
     is_omp_first_level = is_omp_first_level_distrib 
     is_omp_last_level  = is_omp_last_level_distrib 
     ll_begin           = ll_begin_distrib 
     ll_beginp1         = ll_beginp1_distrib 
     ll_end             = ll_end_distrib 
     ll_endm1           = ll_endm1_distrib 
     ll_endp1           = ll_endp1_distrib 
  
  END SUBROUTINE switch_omp_distrib_level
  

  SUBROUTINE switch_omp_no_distrib_level
  IMPLICIT NONE

     is_omp_first_level = is_omp_first_level_full 
     is_omp_last_level  = is_omp_last_level_full 
     ll_begin           = ll_begin_full 
     ll_beginp1         = ll_beginp1_full 
     ll_end             = ll_end_full 
     ll_endm1           = ll_endm1_full 
     ll_endp1           = ll_endp1_full 
  
  END SUBROUTINE switch_omp_no_distrib_level
  

  FUNCTION omp_in_parallel()
#ifdef CPP_USING_OMP
  USE omp_lib, ONLY : omp_in_parallel_=>omp_in_parallel
#endif
  IMPLICIT NONE
    LOGICAL :: omp_in_parallel
    
#ifdef CPP_USING_OMP
    omp_in_parallel=omp_in_parallel_()
#else
    omp_in_parallel=.FALSE.
#endif

  END FUNCTION  omp_in_parallel  
      
END MODULE omp_para