! @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@ !
! @@                                                                @@ !
! @@       PROGRAM  ASCOT 2014.440 (ver.4.53)                       @@ !
! @@     "Abinitio Simulation Code for Quantum Transport"           @@ !
! @@                                                                @@ !
! @@                                                                @@ !
! @@  AUTHOR(S): Naoki WATANABE, Nobutaka NISHIKAWA (Mizuho I.R.)   @@ !
! @@             Hisashi KONDO (Univ. Tokyo)                        @@ !
! @@                                                09/May/2014     @@ !
! @@                                                                @@ !
! @@  Contact address: Phase System Consortium                      @@ !
! @@                                                                @@ !
! @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@ !

module ac_mpi_module
  use ac_parameter
  implicit none
  include "mpif.h"

  type MPI_type
     integer :: size 
     integer :: rank 
     integer :: sizeE, sizeM 
     integer :: rankE, rankM 
     integer :: info 
     logical :: root 

     integer :: isatom 
     integer :: ieatom 

     integer,pointer :: vnmat(:), vsmat(:), vemat(:)
     integer         ::  nmat,     smat,    emat
     integer,pointer :: ispao(:) 
     integer,pointer :: iepao(:) 

     integer :: localsize1
     integer :: localsize2

     real(8) :: num_sc_work2space = -1000.d0
     integer :: num_check_sc = 1

  end type MPI_type

  type(MPI_type), public, save :: MPI 

contains

  subroutine MPI__Initialize
    implicit none

    MPI%rank = 0
    MPI%size = 1
    MPI%root = .true.

    return
  end subroutine MPI__Initialize

  subroutine MPI__Finalize
    implicit none

    if( associated(MPI%vnmat) ) then
       deallocate( MPI%vnmat, MPI%vsmat, MPI%vemat )
    end if

    if( associated(MPI%ispao) ) then
       deallocate(MPI%ispao,MPI%iepao)
    end if

    return
  end subroutine MPI__Finalize

  subroutine MPI__setup
    implicit none
    integer :: a

    MPI%sizeE = 1
    MPI%sizeM = 1

    MPI%isatom = 1
    MPI%ieatom = Param%Data%natom

    if( associated(MPI%ispao) ) then
       deallocate(MPI%ispao, MPI%iepao)
    end if
    allocate( MPI%ispao(Param%Data%natom), MPI%iepao(Param%Data%natom) )

    do a=1, Param%Data%natom
       MPI%ispao(a) = 0+1
       MPI%iepao(a) = Param%Data%vnpao(a)-1+1
    end do

    if( associated(MPI%vnmat) ) then
       deallocate( MPI%vnmat, MPI%vsmat, MPI%vemat )
    end if
    allocate( MPI%vnmat(0:0), MPI%vsmat(0:0), MPI%vemat(0:0) )

    MPI%vsmat(0) = 0+1
    MPI%vemat(0) = MPI%vsmat(0) + Param%Data%npao - 1
    MPI%vnmat(0) = Param%Data%npao

    MPI%smat = MPI%vsmat(0)
    MPI%emat = MPI%vemat(0)
    MPI%nmat = MPI%vnmat(0)

    MPI%localsize1 = Param%Data%npao
    MPI%localsize2 = 2*Param%Data%npao

    return
  end subroutine MPI__setup

  subroutine MPI__show
    implicit none

    write(*,'(a)')    '# disabled MPI parallelization'

    return
  end subroutine MPI__show

  subroutine MPI__Bcast_Inputfile( filename )
    implicit none
    character(64), intent(inout) :: filename

    return
  end subroutine MPI__Bcast_Inputfile

  subroutine MPI__Allreduce_Density( density )
    implicit none
    real(8), intent(inout) :: density (Param%Cell%Na,Param%Cell%Nb,Param%Cell%Nc)

    return
  end subroutine MPI__Allreduce_Density

  subroutine MPI__Allreduce_DensityLS( densityLS )
    implicit none
    complex(8), intent(inout) :: densityLS(Param%Option%nspin,Param%Cell%Na,Param%Cell%Nb,Param%Cell%Nc)

    return
  end subroutine MPI__Allreduce_DensityLS

  subroutine MPI__Allgather_Matrix( matrix )
    implicit none
    complex(8), intent(inout) :: matrix(0:Param%Data%npao-1,0:Param%Data%npao-1)

    return
  end subroutine MPI__Allgather_Matrix

  subroutine MPI__AlltoAll_Matrix( matrix )
    implicit none
    complex(8), intent(inout) :: matrix(Param%Data%npao,Param%Data%npao)
    integer :: i, j

    do i=1, Param%Data%npao
       do j=i, Param%Data%npao
          matrix(i,j) = dconjg(matrix(i,j)) + matrix(j,i)
          matrix(j,i) = dconjg(matrix(i,j))
       end do
    end do

    return
  end subroutine MPI__AlltoAll_Matrix

  subroutine MPI__AlltoAll_Matrix2( matrix2, matrix3 )
    implicit none
    complex(8), intent(inout) :: matrix2(Param%Data%npao,Param%Data%npao)
    complex(8), intent(inout) :: matrix3(Param%Data%npao,Param%Data%npao)
    integer :: i, j

    do i=1, Param%Data%npao
       do j=1, Param%Data%npao
          matrix2(i,j) = dconjg(matrix2(i,j)) + matrix3(j,i)
          matrix3(j,i) = dconjg(matrix2(i,j))
       end do
    end do

    return
  end subroutine MPI__AlltoAll_Matrix2

  subroutine MPI__Allreduce_Force( vatom )
    implicit none
    type(Atom_type), intent(inout) :: vatom(Param%Data%natom)

    return
  end subroutine MPI__Allreduce_Force

  subroutine MPI__Allreduce_HistoryPulay( A )
    implicit none
    real(8), intent(inout) :: A(Param%SCF%mix_history,Param%SCF%mix_history)
    return
  end subroutine MPI__Allreduce_HistoryPulay

  subroutine MPI__Allreduce_HistoryAnderson( A )
    implicit none
    real(8), intent(inout) :: A(Param%SCF%mix_history,Param%SCF%mix_history)
    return
  end subroutine MPI__Allreduce_HistoryAnderson

  subroutine MPI__Allreduce_Error( dEden )
    implicit none
    real(8), intent(inout) :: dEden

    return
  end subroutine MPI__Allreduce_Error

  subroutine MPI__Barrier
    implicit none

    return
  end subroutine MPI__Barrier

  subroutine MPI__ZHEGV( size, matrixH, matrixS, vectorE, iter )
    implicit none

    integer, intent(in)       :: size 
    complex(8), intent(inout) :: matrixH(size,size) 
    complex(8), intent(in)    :: matrixS(size,size) 
    real(8),    intent(out)   :: vectorE(size) 

    complex(8), allocatable   :: localS(:,:)
    complex(8), allocatable   :: work1(:)
    real(8),    allocatable   :: work2(:)

    integer, intent(in) :: iter

    allocate( work1(4*size), work2(7*size) )
    allocate( localS(size,size) )
    localS = matrixS 

    call ZHEGV( 1, 'V', 'L', size, matrixH, size, localS, size, &
         vectorE, work1, 4*size, work2, MPI%info )

    deallocate( localS )
    deallocate( work1, work2 )

    return
  end subroutine MPI__ZHEGV

  subroutine MPI__ZHEGVLS( size, matrixH, matrixS, vectorE, iter )
    implicit none

    integer, intent(in)       :: size 
    complex(8), intent(inout) :: matrixH(2*size,2*size) 
    complex(8), intent(in)    :: matrixS(2*size,2*size) 
    real(8),    intent(out)   :: vectorE(2*size) 

    integer, intent(in) :: iter

    complex(8), allocatable   :: localS(:,:)
    complex(8), allocatable :: work1(:)
    real(8),    allocatable :: work2(:)

    allocate( work1(4*2*size), work2(7*2*size) )
    allocate( localS(2*size,2*size) )
    localS = matrixS 

    call ZHEGV( 1, 'V', 'L', 2*size, matrixH, 2*size, localS, 2*size, &
         vectorE, work1, 4*2*size, work2, MPI%info )

    deallocate( localS )
    deallocate( work1, work2 )

    return
  end subroutine MPI__ZHEGVLS

  subroutine MPI__ZGEMM( transa, transb, size, &
       alpha, matrixA, matrixB, beta, matrixC )
    implicit none

    character, intent(in)     :: transa, transb
    integer, intent(in)       :: size 
    complex(8), intent(in)    :: matrixA(size,size) 
    complex(8), intent(in)    :: matrixB(size,size) 
    complex(8), intent(out)   :: matrixC(size,size) 
    complex(8), intent(in)    :: alpha, beta 

    call ZGEMM( transa, transb, size, size, size, &
         alpha, matrixA, size, matrixB, size, &
         beta,  matrixC, size )

    return
  end subroutine MPI__ZGEMM

  subroutine MPI__ZHEMM( size, &
       alpha, matrixA, matrixB, beta, matrixC )
    implicit none

    integer, intent(in)       :: size 
    complex(8), intent(in)    :: matrixA(size,size) 
    complex(8), intent(in)    :: matrixB(size,size) 
    complex(8), intent(out)   :: matrixC(size,size) 
    complex(8), intent(in)    :: alpha, beta 

    call ZHEMM( 'L', 'L', size, size, size, &
         alpha, matrixA, size, matrixB, size, &
         beta,  matrixC, size )

    return
  end subroutine MPI__ZHEMM

  subroutine MPI__ZPOTRI( size, matrixA )
    implicit none

    integer, intent(in)       :: size 
    complex(8), intent(inout) :: matrixA(size,size) 

    call ZPOTRF( 'L', size, matrixA, size, MPI%info )
    call ZPOTRI( 'L', size, matrixA, size, MPI%info )

    return
  end subroutine MPI__ZPOTRI

  subroutine MPI__ZGETRI( size, matrixA )
    implicit none

    integer, intent(in)       :: size 
    complex(8), intent(inout) :: matrixA(size,size) 

    integer, allocatable :: ipiv(:)
    complex(8), allocatable :: work1(:)
    integer, allocatable :: work2(:)
    integer :: lwork1, lwork2

    allocate( work1(1), work2(1) )
    lwork1=-1; lwork2=-1

    call ZGETRI( size, matrixA, size, &
         ipiv, work1, lwork1, work2, lwork2, MPI%info )

    lwork1 = int(work1(1))+100
    lwork2 = int(work2(1))+100
    deallocate( work1, work2 )

    allocate( work1(lwork1), work2(lwork2) )
    allocate( ipiv(size) )

    call ZGETRF( size, size, matrixA, size, &
         ipiv, MPI%info )

    call ZGETRI( size, matrixA, size, &
         ipiv, work1, lwork1, work2, lwork2, MPI%info )

    deallocate( ipiv )
    deallocate( work1, work2 )

    return
  end subroutine MPI__ZGETRI

end module ac_mpi_module
