!
! Copyright (C) 2017- Orbital-free DFT group at University of Florida
!
! This program is free software; you can redistribute it and/or modify 
! it under the terms of the GNU General Public License as published by 
! the Free Software Foundation; either version 2 of the License, or 
! any later version. See the file LICENSE in the root directory of the 
! present distribution, or http://www.gnu.org/copyleft/gpl.txt ,
! or contact developers via e-mail: calderin@lcalderin.net, or paper mail:
!
! Quantum Theory Project
! University of Florida
! P.O. Box 118435
! Gainesville, FL 32611
!
! This program 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
! GNU General Public License for more details.
! 
! Authors: Lazaro Calderin and Valentin Karasiev
!-----------------------------------------------------------------------
! 


module parallel
use mp
use mp_pools
use mp_bands
use mp_global
use io_global
use qevars,   ONLY : DP, nbnd, nks
use kgecvars, ONLY : nproc_kp, nproc_bands, nproc_pw, &
                     nbnds_per_process, nbnds_per_group, nbnds_local, &
                     my_rank, my_rank_old, my_group, &
                     psi_grad_psi
implicit none


contains

!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

   subroutine parallel_bands_init

      nproc_kp    = mp_size( inter_pool_comm )
      nproc_bands = mp_size( inter_bgrp_comm )
      nproc_pw    = mp_size( intra_bgrp_comm )

      write(stdout,*) ' '
      write(stdout,*) '    --------------------------------------------------------------'
      write(stdout,*) '                        PARALLEL DISTRIBUTION'
      write(stdout,*) '    --------------------------------------------------------------'
      write(stdout,'(A,i4,4x,A,i4,4x,A,i4)') &
                      '                     nk=', nproc_kp, 'npw=', nproc_pw, 'nb=', nproc_bands
      write(stdout,*) '    --------------------------------------------------------------'

      my_rank           = mp_rank( inter_bgrp_comm )
      nproc_bands       = mp_size( inter_bgrp_comm )
      nbnds_per_process = nbnd / nproc_bands
      nbnds_per_group   = 0
      my_group          = 0
      my_rank_old       = my_rank
    
      if( my_rank /= (nproc_bands -1) ) then
         nbnds_local = nbnds_per_process
      else
         nbnds_local = nbnds_per_process + mod( nbnd, nproc_bands )
      endif

   end subroutine parallel_bands_init

!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

   subroutine recover_processes

      integer :: npwrec
      integer :: groupid, my_rank_world
      integer :: ib1, ib2

      complex(DP), allocatable :: gmaux(:,:,:,:)

      write(stdout,*) ''
      write(stdout,'(A)',advance='no') &
        '     Recovering pw processes for band parallelization ...'


      nproc_kp    = mp_size( inter_pool_comm ) 
      nproc_bands = mp_size( inter_bgrp_comm )
      nproc_pw    = mp_size( intra_bgrp_comm )
      
!      npwrec = mod( nbnd, nproc_pw*nproc_bands )
!      npwrec = mod( nbnds_per_process, nproc_pw )

      npwrec = nbnds_local / nproc_pw
 
      if ( nproc_pw > 1 .and. npwrec > 0 ) then

#ifdef __KGDEBUGPARA
         call parallel_groups_and_comms
#endif

         ! first free groups and communicators
         call mp_comm_group( inter_bgrp_comm, groupid )
         call mp_group_free( groupid )
         call mp_comm_group( intra_bgrp_comm, groupid )
         call mp_group_free( groupid )
         call mp_comm_free( inter_bgrp_comm )
         call mp_comm_free( intra_bgrp_comm )

         ! recreate inter_bgrp and intra_bgrp from world_comm,
         ! but swaping the names of inter and intra

         my_rank_world = mp_rank( world_comm )
         nproc_bands   = nproc_bands * nproc_pw
         my_group      = my_rank_world / nproc_bands
         my_rank       = MOD( my_rank_world, nproc_bands )

         call mp_comm_split( world_comm, my_group, my_rank_world, inter_bgrp_comm )
         call mp_comm_split( world_comm,  my_rank, my_rank_world, intra_bgrp_comm )  

#ifdef __KGDEBUGPARA
         call parallel_groups_and_comms
#endif

         ! Change the size of slices and redistribute the bands
         my_rank     = mp_rank( inter_bgrp_comm )
         my_group    = int( my_rank / nproc_pw )
         my_rank_old = mod( my_rank, nproc_pw )

         ! Bands per group and new number of bands per processes
         nbnds_per_group   = nbnds_per_process
         nbnds_per_process = nbnds_local / nproc_pw 


         ! Indices inside a band group 
         ib1 = nbnds_per_process * my_rank_old + 1
         ib2 = nbnds_per_process * ( my_rank_old + 1 )


         ! nbnds_local and ib2 index taking into account remainders 
         if( mod( (my_rank+1), nproc_pw ) /= 0 ) then
            nbnds_local = nbnds_per_process
         else
            ib2         = ib2 + mod( nbnds_local, nproc_pw )
            nbnds_local = nbnds_per_process + mod( nbnds_local, nproc_pw )
         endif

         allocate( gmaux( 1:3, 1:nbnd, 1:nbnds_local, 1:nks ) )

                gmaux( 1:3, 1:nbnd,   1:nbnds_local, 1:nks ) = &
         psi_grad_psi( 1:3, 1:nbnd, ib1:ib2        , 1:nks   ) 
   
         deallocate( psi_grad_psi )
         allocate( psi_grad_psi(1:3,1:nbnd,1:nbnds_local, 1:nks ) ) 
   
         psi_grad_psi( 1:3, 1:nbnd, 1:nbnds_local, 1:nks ) = &
                gmaux( 1:3, 1:nbnd, 1:nbnds_local, 1:nks )
   
         deallocate( gmaux )

         write(stdout,'(A)') ' done.'

      else
      
         write(stdout,'(A)') ' done.'

         if( nproc_pw == 1 ) then
            write(stdout,*) " "
            write(stdout,*) &
                  '     WARNING: Parallel: np is 1, nothing to recover'
         endif
         
         if( npwrec < 1 ) then
            write(stdout,*) ''
            write(stdout,*) &
                '     WARNING: Parallel: int(nbnd/nb)/np less than one'
            write(stdout,*) &
                '      ( int(nbnd/nb)=',nbnds_per_process,', np=', nproc_pw,')'
            write(stdout,*) &
                '     WARNING: therefore nothing recovered.'
         endif
      
      endif ! recovering of pw cores for bands parallelization

      write(stdout,*) ' '
      write(stdout,*) '--------------------------------------------------------------'
      write(stdout,*) '          PARALLEL DISTRIBUTION AFTER NPW RECOVERY'
      write(stdout,*) '--------------------------------------------------------------'
      write(stdout,'(A,i4,4x,A,i4,4x)') &
                      '                     nk=', nproc_kp,  'nb=', nproc_bands
      write(stdout,*) '--------------------------------------------------------------'


   end subroutine recover_processes

!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

   subroutine parallel_groups_and_comms
!   use mp
!   use mp_global
!   use mp_pools
!   use mp_bands
!   use io_global
!   implicit none
   
      character(20) :: comm_name(5)
      integer :: groupid, groupsize   
      integer :: maxrank, ip, ic
      integer, allocatable :: ranks(:,:)
 

      comm_name(1)= "world_comm"
      comm_name(2)= "inter_pool_comm"
      comm_name(3)= "intra_pool_comm"
      comm_name(4)= "inter_bgrp_comm"
      comm_name(5)= "intra_bgrp_comm"

       if(ionode) then

         write(*,*) "======================================================"
         write(*,*) "Communicators' group"
         write(*,*) "======================================================"
   
         call mp_comm_group( world_comm, groupid )
         groupsize=mp_size( world_comm )
         write(*,*) "world_comm      id and size", groupid, groupsize 
         call mp_comm_group( inter_pool_comm, groupid )
         groupsize=mp_size( inter_pool_comm )
         write(*,*) "inter_pool_comm id and size", groupid, groupsize 
         call mp_comm_group( intra_pool_comm, groupid )
         groupsize=mp_size( intra_pool_comm )
         write(*,*) "intra_pool_comm id and size", groupid, groupsize 
         call mp_comm_group( inter_bgrp_comm, groupid )
         groupsize=mp_size( inter_bgrp_comm )
         write(*,*) "inter_bgrp_comm id and size", groupid, groupsize 
         call mp_comm_group( intra_bgrp_comm, groupid )
         groupsize=mp_size( intra_bgrp_comm )
         write(*,*) "intra_bgrp_comm id and size", groupid, groupsize 
         write(*,*) " "
         write(*,*) " "
   
      endif ! ionode

      maxrank = mp_size( world_comm ) -1 
      allocate( ranks( 0:maxrank,  5) )
      ranks = 0

      do ip=0, maxrank   
         if ( mp_rank(world_comm)==ip ) then
            ranks(ip,1)=mp_rank(      world_comm )
            ranks(ip,2)=mp_rank( inter_pool_comm )
            ranks(ip,3)=mp_rank( intra_pool_comm )
            ranks(ip,4)=mp_rank( inter_bgrp_comm )
            ranks(ip,5)=mp_rank( intra_bgrp_comm )
          endif
      enddo

      call mp_sum( ranks, world_comm )

      if(ionode) then

         write(*,*) "======================================================"
         write(*,*) "Processes' ranks in each communicator"
         write(*,*) "======================================================"
         do ic=1,5
            write(*,'(1x,a18,64i4)'), comm_name(ic), ((ranks(ip,ic)),ip=0,maxrank)
         enddo

      endif

   end subroutine parallel_groups_and_comms      

end module parallel
