! @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@ !
! @@                                                                @@ !
! @@       PROGRAM  ASCOT 2014.440 (ver.4.40)                       @@ !
! @@     "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                      @@ !
! @@                                                                @@ !
! @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@ !

subroutine BandMatrix__allocate( bmat )
  use ac_misc_module
  use ac_mpi_module
  implicit none
  type(BandMatrix_type) :: bmat
  integer :: i, j

  allocate( bmat%S(Base%npao, MPI%smat:MPI%emat) )
  allocate( bmat%H(Base%npao, MPI%smat:MPI%emat, Param%Option%nspin) )

  if( Param%Option%spin_orbit ) then
     allocate( bmat%E(Base%npao*2, 1) )
     allocate( bmat%F(Base%npao*2, 1) )
  else
     allocate( bmat%E(Base%npao, Param%Option%nspin) )
     allocate( bmat%F(Base%npao, Param%Option%nspin) )
  end if

  return
end subroutine BandMatrix__allocate

subroutine BandMatrix__deallocate( bmat )
  use ac_misc_module
  use ac_mpi_module
  implicit none
  type(BandMatrix_type) :: bmat

  if( associated(bmat%S) ) deallocate(bmat%S)
  if( associated(bmat%H) ) deallocate(bmat%H)
  if( associated(bmat%E) ) deallocate(bmat%E)
  if( associated(bmat%F) ) deallocate(bmat%F)

  return
end subroutine BandMatrix__deallocate

subroutine BandMatrix__calc( bmat, k )
  use ac_misc_module
  use ac_mpi_module
  implicit none
  type(BandMatrix_type), intent(out) :: bmat
  real(8), intent(in)    :: k(3)

  integer    :: spin
  integer    :: l
  integer    :: a, i, i1
  integer    :: b, j, j1
  complex(8) :: bloch

  bmat%S = 0.d0
  bmat%H = 0.d0

  do l=0, Param%Cell%nL-1

     bloch =  polar( dot_product(k,Param%Cell%vL(:,l)) )
     if( l==0 ) then
        bloch = 0.5d0 * bloch
     end if

     do a=MPI%isatom, MPI%ieatom
        i=Base%vipao(a)

        do b=1, Param%Data%natom
           j=Base%vipao(b)

           if( .not. associated(Hamiltonian%vAtomMatrix(b,a,l)%Ptr) ) then
              cycle
           end if

           do i1=MPI%ispao(a), MPI%iepao(a)
              do j1=1, Base%vnpao(b)

                 if( Param%Option%spin_orbit ) then
                    bmat%S(j+j1-1,i+i1-1) = bmat%S(j+j1-1,i+i1-1) &
                         + Hamiltonian%vAtomMatrix(b,a,l)%Ptr%S(j1,i1) * bloch
                    bmat%H(j+j1-1,i+i1-1,:) = bmat%H(j+j1-1,i+i1-1,:) &
                         + Hamiltonian%vAtomMatrix(b,a,l)%Ptr%HLS(j1,i1,:) * bloch
                 else
                    bmat%S(j+j1-1,i+i1-1) = bmat%S(j+j1-1,i+i1-1) &
                         + Hamiltonian%vAtomMatrix(b,a,l)%Ptr%S(j1,i1) * bloch
                    bmat%H(j+j1-1,i+i1-1,:) = bmat%H(j+j1-1,i+i1-1,:) &
                         + Hamiltonian%vAtomMatrix(b,a,l)%Ptr%H(j1,i1,:) * bloch
                 end if
              end do
           end do

        end do
     end do
  end do


  if( Param%Option%spin_orbit ) then
     call MPI__AlltoAll_Matrix( bmat%S )
     call MPI__AlltoAll_Matrix( bmat%H(:,:,1) )
     call MPI__AlltoAll_Matrix( bmat%H(:,:,4) )
     call MPI__AlltoAll_Matrix2( bmat%H(:,:,2), bmat%H(:,:,3) )
  else
     call MPI__AlltoAll_Matrix( bmat%S )
     do spin=1, Param%Option%nspin
        call MPI__AlltoAll_Matrix( bmat%H(:,:,spin) )
     end do
  end if

  return
end subroutine BandMatrix__calc

subroutine BandMatrix__solve( bmat, iter )
  use ac_misc_module
  use ac_mpi_module
  implicit none
  type(BandMatrix_type), intent(inout) :: bmat
  integer, intent(in) :: iter

  integer    :: spin, ispin, jspin
  integer    :: i, j
  complex(8), allocatable :: H_LS(:,:), S_LS(:,:)
  complex(8), allocatable :: work1(:)
  real(8),    allocatable :: work2(:)
  integer  :: info

  if( Param%Option%spin_orbit ) then

     allocate( H_LS(2*Base%npao,2*MPI%smat-1:2*MPI%emat-0) )
     allocate( S_LS(2*Base%npao,2*MPI%smat-1:2*MPI%emat-0) )

     do i=MPI%smat, MPI%emat
        do j=1, Base%npao
           H_LS(2*j-1,2*i-1) = bmat%H(j,i,1)
           H_LS(2*j-1,2*i-0) = bmat%H(j,i,2)
           H_LS(2*j-0,2*i-1) = bmat%H(j,i,3)
           H_LS(2*j-0,2*i-0) = bmat%H(j,i,4)
           S_LS(2*j-1,2*i-1) = bmat%S(j,i)
           S_LS(2*j-1,2*i-0) = 0.0d0
           S_LS(2*j-0,2*i-1) = 0.0d0
           S_LS(2*j-0,2*i-0) = bmat%S(j,i)
        end do
     end do

     call MPI__ZHEGVLS( Base%npao, H_LS, S_LS, bmat%E, iter )

     do i=MPI%smat, MPI%emat
        do j=1, Base%npao
           bmat%H(j,i,1) = H_LS(2*j-1,2*i-1)
           bmat%H(j,i,2) = H_LS(2*j-0,2*i-1)
           bmat%H(j,i,3) = H_LS(2*j-1,2*i-0)
           bmat%H(j,i,4) = H_LS(2*j-0,2*i-0)
        end do
     end do

     deallocate( H_LS, S_LS )
  else
     do spin=1, Param%Option%nspin
        call MPI__ZHEGV( Base%npao, bmat%H(:,:,spin), bmat%S, bmat%E(:,spin), iter )
     end do
  end if

  return
end subroutine BandMatrix__solve

subroutine BandMatrix__calcCDM( bmat )
  use ac_misc_module
  use ac_mpi_module
  implicit none
  type(BandMatrix_type), intent(out) :: bmat

  integer :: m, i, j, spin, ispin, jspin, kspin
  real(8) :: F
  complex(8), allocatable :: tempS(:,:), tempH(:,:)
  complex(8), allocatable :: tempQ(:,:,:),tempR(:,:)
  complex(8), parameter   :: C1 = dcmplx(1.0d0,0.d0), C0 = dcmplx(0.0d0,0.d0)


  allocate( tempS(Param%Data%npao,MPI%smat:MPI%smat+MPI%localsize1-1) )
  allocate( tempH(Param%Data%npao,MPI%smat:MPI%smat+MPI%localsize1-1) )
  allocate( tempQ(Param%Data%npao,MPI%smat:MPI%smat+MPI%localsize1-1, Param%Option%nspin) )
  allocate( tempR(Param%Data%npao,MPI%smat:MPI%smat+MPI%localsize1-1) )

  if( Param%Option%spin_orbit ) then
     do m=1, Base%npao
        call Energy__fermifunc(F,bmat%E(2*m-1,1))
        bmat%F(2*m-1,1) = Param%Data%g * F 
        call Energy__fermifunc(F,bmat%E(2*m-0,1))
        bmat%F(2*m-0,1) = Param%Data%g * F 
     end do

     tempR=0.d0
     do spin=1, Param%Option%nspin

        select case(spin)
        case(1)
           ispin = 1
           jspin = 1
        case(2)
           ispin = 1
           jspin = 2
        case(3)
           ispin = 2
           jspin = 1
        case(4)
           ispin = 2
           jspin = 2
        end select

        tempQ(:,:,spin) = 0.d0
        tempH=0.d0
        tempS=0.d0

        do kspin=0,1
           do m=MPI%smat, MPI%emat
              do j=1, Base%npao
                 tempH(j,m) = dconjg(bmat%H(j,m,jspin+kspin*2))
              end do
              do i=1, Base%npao
                 tempS(i,m) = bmat%H(i,m,ispin+kspin*2) * bmat%F(2*m-1+kspin,1)
              end do
           end do

           call MPI__ZGEMM('N', 'T', Base%npao, &
                C1, tempH, tempS, C1*dfloat(kspin), tempQ(:,:,spin) )

           if( spin == 1 .or. spin == 4 ) then
              do m=MPI%smat, MPI%emat
                 do i=1, Base%npao
                    tempS(i,m) = tempS(i,m) * bmat%E(2*m-1+kspin,1)
                 end do
              end do

              call MPI__ZGEMM('N', 'T', Base%npao, &
                   C1, tempH, tempS, C1, tempR )
           end if
        end do

     end do

     do m=MPI%smat, MPI%emat
        do j=1, Base%npao
           do spin=1, Param%Option%nspin
              bmat%H(j,m,spin) = tempQ(j,m,spin)
           end do
           bmat%S(j,m) = tempR(j,m)
        end do
     end do
  else
     do m=MPI%smat, MPI%emat
        do spin=1, Param%Option%nspin
           call Energy__fermifunc(F,bmat%E(m,spin))
           bmat%F(m,spin) = Param%Data%g*F
        end do
     end do
     tempR=0.d0
     do spin=1, Param%Option%nspin
        tempQ(:,:,spin) = 0.d0
        tempH=0.d0
        tempS=0.d0

        do m=MPI%smat, MPI%emat
           do j=1, Base%npao
              tempH(j,m) = dconjg(bmat%H(j,m,spin))
           end do
           do i=1, Base%npao
              tempS(i,m) = bmat%H(i,m,spin) * bmat%F(m,spin)
           end do
        end do

        call MPI__ZGEMM('N', 'T', Base%npao, &
             C1, tempH, tempS, C0, tempQ(:,:,spin) )

        do m=MPI%smat, MPI%emat
           do i=1, Base%npao
              tempS(i,m) = tempS(i,m) * bmat%E(m,spin)
           end do
        end do

        call MPI__ZGEMM('N', 'T', Base%npao, &
             C1, tempH, tempS, C1, tempR )
     end do

     do m=MPI%smat, MPI%emat
        do j=1, Base%npao
           do spin=1, Param%Option%nspin
              bmat%H(j,m,spin) = tempQ(j,m,spin)
           end do
           bmat%S(j,m) = tempR(j,m)
        end do
     end do

  end if

  deallocate( tempQ, tempR )
  deallocate( tempS, tempH )

  return
end subroutine BandMatrix__calcCDM
