!--------------------------------------- LICENCE BEGIN -----------------------------------
!Environment Canada - Atmospheric Science and Technology License/Disclaimer,
!                     version 3; Last Modified: May 7, 2008.
!This is free but copyrighted software; you can use/redistribute/modify it under the terms
!of the Environment Canada - Atmospheric Science and Technology License/Disclaimer
!version 3 or (at your option) any later version that should be found at:
!http://collaboration.cmc.ec.gc.ca/science/rpn.comm/license.html
!
!This software is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY;
!without even the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
!See the above mentioned License/Disclaimer for more details.
!You should have received a copy of the License/Disclaimer along with this software;
!if not, you can write to: EC-RPN COMM Group, 2121 TransCanada, suite 500, Dorval (Quebec),
!CANADA, H9P 1J3; or send e-mail to service.rpn@ec.gc.ca
!-------------------------------------- LICENCE END --------------------------------------


module mpivar_mod 31,1
  use mpi_mod
  implicit none
  save
  private

  ! public procedures
  public :: mpivar_setup_latbands, mpivar_setup_lonbands
  public :: mpivar_setup_m, mpivar_setup_n
  public :: mpivar_setup_levels_npex, mpivar_setup_levels_npey

  ! public variables through inheritance
  public :: mpi_myid, mpi_nprocs, mpi_npex, mpi_npey, mpi_myidx, mpi_myidy
  ! public procedures through inheritance
  public :: mpi_initialize, mpi_getptopo, mpi_allreduce_sumreal8scalar

  contains



  subroutine mpivar_setup_latbands(nj,latPerPE,myLatBeg,myLatEnd,myLatHalfBeg,myLatHalfEnd) 9,2
    ! Purpose: compute parameters that define the mpi distribution of
    !          latitudes over tasks in Y direction (npey)
    implicit none
    integer          :: nj,latPerPE,myLatBeg,myLatEnd,njlath
    integer,optional :: myLatHalfBeg,myLatHalfEnd

    latPerPE=ceiling(dble(nj)/dble(mpi_npey))
    myLatBeg=1+mpi_myidy*latPerPE
    if(myLatBeg.gt.nj) then
      call abort3d('mpivar_setup_latbands: latitudes not separable on mpi processes!')
    endif
    myLatEnd=(1+mpi_myidy)*latPerPE
    if(myLatEnd>nj) then
      call abort3d('mpivar_setup_latbands: latitudes not divisible by MPI numprocs!')
    endif

    if(present(myLatHalfBeg).and.present(myLatHalfEnd)) then
      njlath = (nj + 1)/2
      if(myLatBeg<=njlath .and. myLatEnd<=njlath) then
        myLatHalfBeg=myLatBeg
        myLatHalfEnd=myLatEnd
      elseif(myLatBeg>=njlath .and. myLatEnd>=njlath) then
        myLatHalfBeg=1+nj-myLatEnd
        myLatHalfEnd=1+nj-myLatBeg
      else
        myLatHalfBeg=min(myLatBeg,1+nj-myLatEnd)
        myLatHalfEnd=njlath
      endif
    endif

  end subroutine mpivar_setup_latbands



  subroutine mpivar_setup_lonbands(ni,lonPerPE,myLonBeg,myLonEnd) 9,2
    ! Purpose: compute parameters that define the mpi distribution of
    !          longitudes over tasks in X direction (npex)
    implicit none
    integer          :: ni,lonPerPE,myLonBeg,myLonEnd

    lonPerPE=ceiling(dble(ni)/dble(mpi_npex))
    myLonBeg=1+mpi_myidx*lonPerPE
    if(myLonBeg.gt.ni) then
      call abort3d('mpivar_setup_lonbands: longitudes not separable on mpi processes!')
    endif
    myLonEnd=(1+mpi_myidx)*lonPerPE
    if(myLonEnd>ni) then
      call abort3d('mpivar_setup_lonbands: longitudes not divisible by MPI npex!')
    endif

  end subroutine mpivar_setup_lonbands



  subroutine mpivar_setup_m(ntrunc,mymBeg,mymEnd,mymSkip,mymCount) 7,1
    ! Purpose: compute parameters that define the mpi distribution of
    !          wavenumber m over tasks in Y direction (npey)
    implicit none
    integer :: ntrunc,mymBeg,mymEnd,mymSkip,mymCount,jm

    if((ntrunc+1) .lt. mpi_npey) then
      write(*,*) 'mpivar_setup_m: NPEY (=',mpi_npey,') ',  &
                 'must be less than or equal to ntrunc+1 (=',ntrunc+1,')!'
      call abort3d('mpivar_setup_m')
    endif

    mymBeg=mpi_myidy
    mymEnd=ntrunc
    mymSkip=mpi_npey
    mymCount=0
    do jm=mymBeg,mymEnd,mymSkip
      mymCount=mymCount+1
    enddo

  end subroutine mpivar_setup_m

 

  subroutine mpivar_setup_n(ntrunc,mynBeg,mynEnd,mynSkip,mynCount) 3,1
    ! Purpose: compute parameters that define the mpi distribution of
    !          wavenumber n over tasks in X direction (npex)
    implicit none
    integer :: ntrunc,mynBeg,mynEnd,mynSkip,mynCount,jn

    if((ntrunc+1) .lt. mpi_npex) then
      write(*,*) 'mpivar_setup_n: NPEX (=',mpi_npex,') ',  &
                 'must be less than or equal to ntrunc+1 (=',ntrunc+1,')!'
      call abort3d('mpivar_setup_n')
    endif

    mynBeg=mpi_myidx
    mynEnd=ntrunc
    mynSkip=mpi_npex
    mynCount=0
    do jn=mynBeg,mynEnd,mynSkip
      mynCount=mynCount+1
    enddo

  end subroutine mpivar_setup_n



  subroutine mpivar_setup_levels_npey(numlevels,myLevBeg,myLevEnd,myLevCount) 1,1
    ! Purpose: compute parameters that define the mpi distribution of
    !          levels over tasks in Y direction (npey)
    implicit none
    integer :: numlevels,myLevBeg,myLevEnd,myLevCount
    integer :: jlev,jproc
    integer :: myLevCounts(mpi_npey)

    if(numlevels .lt. mpi_npey) then
      write(*,*) 'mpivar_setup_levels_npey: NPEY (=',mpi_npey,') ',  &
                 'must be less than or equal to number of levels (=',numlevels,')!'
      call abort3d('mpivar_setup_levels_npey')
    endif

    myLevCounts(:)=0
    do jproc=1,mpi_npey
      do jlev=jproc,numlevels,mpi_npey
        myLevCounts(jproc)=myLevCounts(jproc)+1
      enddo
    enddo
    myLevCount=myLevCounts(mpi_myidy+1)

    myLevBeg=1
    do jproc=1,mpi_myidy
      myLevBeg=myLevBeg+myLevCounts(jproc)
    enddo
    myLevEnd=myLevBeg+myLevCount-1

  end subroutine mpivar_setup_levels_npey



  subroutine mpivar_setup_levels_npex(numlevels,myLevBeg,myLevEnd,myLevCount) 4,2
    ! Purpose: compute parameters that define the mpi distribution of
    !          levels over tasks in X direction (npex)
    implicit none
    integer :: numlevels,myLevBeg,myLevEnd,myLevCount
    integer :: jlev,jproc,factor
    integer :: myLevCounts(mpi_npex)
    logical :: makeEven = .true. ! for simplicity (for now) always divide into even number of levels per MPI task

    if(numlevels .lt. mpi_npex) then
      write(*,*) 'mpivar_setup_levels_npex: NPEX (=',mpi_npex,') ',  &
                 'must be less than or equal to number of levels (=',numlevels,')!'
      call abort3d('mpivar_setup_levels_npex')
    endif

    if(makeEven) then
      if(mod(numlevels,2).ne.0) then
        write(*,*) 'mpivar_setup_levels_npex: total number of levels must be even, now=',numlevels
        call abort3d('mpivar_setup_levels_npex')
      endif
      factor = 2
    else 
      factor = 1
    endif

    myLevCounts(:)=0
    do jproc=1,mpi_npex
      do jlev=jproc,(numlevels/factor),mpi_npex
        myLevCounts(jproc)=myLevCounts(jproc)+1
      enddo
    enddo
    do jproc=1,mpi_npex
      myLevCounts(jproc)=myLevCounts(jproc)*factor
    enddo

    myLevCount=myLevCounts(mpi_myidx+1)

    myLevBeg=1
    do jproc=1,mpi_myidx
      myLevBeg=myLevBeg+myLevCounts(jproc)
    enddo
    myLevEnd=myLevBeg+myLevCount-1

  end subroutine mpivar_setup_levels_npex

 
end module mpivar_mod