!> @file
!>  Contains files for constrained DFT
!! @author
!!    Copyright (C) 2013-2014 BigDFT group
!!    This file is distributed under the terms of the
!!    GNU General Public License, see ~/COPYING file
!!    or http://www.gnu.org/copyleft/gpl.txt .
!!    For the list of contributors, see ~/AUTHORS



! THIS COMMENT TO BE UPDATED
!> CDFT: calculates the weight matrix w_ab via the expression S^1/2PS^1/2, where S is the overlap of the whole system
!! CDFT: and P is a projector matrix onto the tmbs of the desired fragment
!! CDFT: for standalone CDFT calculations, assuming one charged fragment, for transfer integrals assuming two fragments
!! CDFT: where we constrain the difference - should later generalize this
subroutine calculate_weight_matrix_lowdin_wrapper(cdft,tmb,input,ref_frags,calculate_overlap_matrix,meth_overlap,&
     first_it,calc_gradient)
  use module_bigdft_profiling
  use module_types
  use module_fragments
  use constrained_dft
  implicit none
  integer, intent(in) :: meth_overlap
  type(cdft_data), intent(inout) :: cdft
  type(input_variables),intent(in) :: input
  type(dft_wavefunction), intent(inout) :: tmb
  logical, intent(in) :: calculate_overlap_matrix
  type(system_fragment), dimension(input%frag%nfrag_ref), intent(in) :: ref_frags
  logical, intent(in) :: first_it, calc_gradient

  integer :: i
  !real(kind=gp), dimension(:,:), pointer :: ovrlp_half
  character(len=*),parameter :: subname='calculate_weight_matrix_lowdin'

  call f_routine('calculate_weight_matrix_lowdin_wrapper')

  !ovrlp_half=f_malloc_ptr((/tmb%orbs%norb,tmb%orbs%norb/), id='ovrlp_half')
  do i=1,cdft%num_constraints
     call calculate_weight_matrix_lowdin(cdft%constraints(i)%weight_matrix,cdft%constraints(i)%weight_matrix_,&
          cdft%constraints(i),cdft%input_constraints(i),tmb,input%frag,&
          ref_frags,calculate_overlap_matrix,.true.,meth_overlap,&
          input%lin%cdft_add_w_guess,first_it,cdft%psit_c,cdft%psit_f,calc_gradient)
  end do
  !call f_free_ptr(ovrlp_half)

  call f_release_routine()

end subroutine calculate_weight_matrix_lowdin_wrapper


subroutine calculate_weight_matrix_lowdin(weight_matrix,weight_matrix_,constraint,input_constraint,tmb,input_frag,&
                                          ref_frags,calculate_overlap_matrix,calculate_ovrlp_half,meth_overlap,frac_w_guess,&
                                          first_it,cdft_psit_c,cdft_psit_f,calc_gradient)
  use module_precisions
  use module_types
  use module_fragments
  use constrained_dft
  use public_enums
  use communications_base, only: TRANSPOSE_FULL
  use communications, only: transpose_localized
  use sparsematrix_base, only: matrices, sparse_matrix, sparsematrix_malloc_ptr, &
                               DENSE_FULL, assignment(=), &
                               allocate_matrices, deallocate_matrices
  use sparsematrix, only: compress_matrix, uncompress_matrix, &
                          gather_matrix_from_taskgroups_inplace, uncompress_matrix2
  use transposed_operations, only: calculate_overlap_transposed
  use matrix_operations, only: overlapPowerGeneral
  use module_bigdft_mpi
  use module_bigdft_arrays
  use module_bigdft_profiling
  implicit none
  type(sparse_matrix), intent(in) :: weight_matrix
  type(matrices), intent(inout) :: weight_matrix_
  type(fragmentInputParameters),intent(in) :: input_frag
  type(dft_wavefunction), intent(inout) :: tmb
  logical, intent(in) :: calculate_overlap_matrix, calculate_ovrlp_half
  type(system_fragment), dimension(input_frag%nfrag_ref), intent(in) :: ref_frags
  integer, intent(in) :: meth_overlap
  real(kind=gp), intent(in) :: frac_w_guess
  logical, intent(in) :: first_it, calc_gradient
  real(kind=gp), dimension(tmb%collcom%ndimind_c), intent(inout) :: cdft_psit_c
  real(kind=gp), dimension(7*tmb%collcom%ndimind_f), intent(inout) :: cdft_psit_f
  type(constraint_data), intent(inout) :: constraint
  type(input_constraint_data), intent(inout) :: input_constraint
  !local variables
  integer :: ifrag,iorb,ifrag_ref,isforb,ierr,cdft_iorb1,cdft_iorb2,jorb,ispin,iiorb,iiorbs
  integer,dimension(1) :: power
  real(kind=gp), allocatable, dimension(:,:,:) :: proj_mat
  real(kind=gp), allocatable, dimension(:,:) :: weight_matrixp, proj_ovrlp_half
  character(len=*),parameter :: subname='calculate_weight_matrix_lowdin'
  real(kind=gp) :: max_error, mean_error, sn, hh, ll, hl
  type(matrices),dimension(1) :: inv_ovrlp

  call f_routine(id='calculate_weight_matrix_lowdin')

  call allocate_matrices(tmb%linmat%smat(3), allocate_full=.true., matname='inv_ovrlp', mat=inv_ovrlp(1))
 
  if ((.not. first_it)) stop 'CDFT not yet compatible with SF optimization'

  if (calculate_overlap_matrix .or. first_it) then
     if(.not.tmb%can_use_transposed) then
         !if(.not.associated(tmb%psit_c)) then
         !    tmb%psit_c = f_malloc_ptr(sum(tmb%collcom%nrecvcounts_c),id='tmb%psit_c')
         !end if
         !if(.not.associated(tmb%psit_f)) then
         !    tmb%psit_f = f_malloc_ptr(7*sum(tmb%collcom%nrecvcounts_f),id='tmb%psit_f')
         !end if
         call transpose_localized(bigdft_mpi%iproc, bigdft_mpi%nproc, tmb%npsidim_orbs, tmb%orbs, tmb%collcom, &
              TRANSPOSE_FULL, tmb%psi, tmb%psit_c, tmb%psit_f, tmb%lzd)
         tmb%can_use_transposed=.true.
     end if
  end if  

  ! first iteration, therefore the current SFs belong to the ground state and should be saved
  ! easier to do it here than in cdft_data_init due to the fact that we know we have the transposed form here
  if (first_it) then
     call vcopy(tmb%collcom%ndimind_c, tmb%psit_c(1), 1, cdft_psit_c(1), 1)
     call vcopy(7*tmb%collcom%ndimind_f, tmb%psit_f(1), 1, cdft_psit_f(1), 1)
  end if


  ! this is now the overlap between the current SF basis and the input GS basis
  if (calculate_overlap_matrix .or. calculate_ovrlp_half .or. (.not. first_it)) then
     call calculate_overlap_transposed(bigdft_mpi%iproc, bigdft_mpi%nproc, tmb%orbs, tmb%collcom, tmb%psit_c, &
          cdft_psit_c, tmb%psit_f, cdft_psit_f, tmb%linmat%smat(1), tmb%linmat%auxs, tmb%linmat%ovrlp_)
     !!call gather_matrix_from_taskgroups_inplace(bigdft_mpi%iproc, bigdft_mpi%nproc, tmb%linmat%smat(1), tmb%linmat%ovrlp_)

     if (calculate_ovrlp_half) then
        tmb%linmat%ovrlp_%matrix = sparsematrix_malloc_ptr(tmb%linmat%smat(1), iaction=DENSE_FULL, id='tmb%linmat%ovrlp_%matrix')
     end if

     call uncompress_matrix2(bigdft_mpi%iproc, bigdft_mpi%nproc, bigdft_mpi%mpi_comm, &
          tmb%linmat%smat(1), tmb%linmat%ovrlp_%matrix_compr, tmb%linmat%ovrlp_%matrix)
  end if  

 
  if (calculate_ovrlp_half) then
     if (trim(input_constraint%constraint_type) == 'SPATIAL' .or. trim(input_constraint%constraint_type) == 'CHARGE_TRANSFER') then
        ! Maybe not clean here to use twice tmb%linmat%smat(1), but it should not
        ! matter as dense is used
        power(1)=2
        call overlapPowerGeneral(bigdft_mpi%iproc, bigdft_mpi%nproc,bigdft_mpi%mpi_comm,&
                meth_overlap, 1, power, &
                tmb%orthpar%blocksize_pdsyev, &
                imode=2, ovrlp_smat=tmb%linmat%smat(1), inv_ovrlp_smat=tmb%linmat%smat(3), &
                ovrlp_mat=tmb%linmat%ovrlp_, inv_ovrlp_mat=inv_ovrlp, &
                check_accur=.false., max_error=max_error, mean_error=mean_error)
     end if
  end if

  proj_mat=f_malloc0((/tmb%orbs%norbu,tmb%orbs%norbu,tmb%linmat%smat(1)%nspin/),id='proj_mat')
  call get_proj_mat(constraint, input_constraint, input_frag, ref_frags, tmb, proj_mat)

  ! once we have generated proj_mat, we no longer need the CDFT input data, so we can free it 
  call input_constraint_data_free(input_constraint)

  !!if (bigdft_mpi%iproc==0) then
  !!  write(2728,*) 'proj mat'
  !!  do ispin=1,weight_matrix%nspin
  !!     do iorb=1,tmb%orbs%norbu
  !!        do jorb=1,tmb%orbs%norbu
  !!           write(2728,*) ispin,iorb,jorb,proj_mat(iorb,jorb,ispin)
  !!        end do
  !!     end do
  !!  end do
  !!end if

  proj_ovrlp_half=f_malloc0((/tmb%orbs%norbu,tmb%orbs%norbp/),id='proj_ovrlp_half')
  if (tmb%orbs%norbp>0) then
     do iorb=1,tmb%orbs%norbp
         iiorb=tmb%orbs%isorb+iorb
         if (iiorb<=tmb%orbs%norbu) then
             ispin=1
         else
             ispin=2
         end if
         iiorbs=iiorb-(ispin-1)*tmb%orbs%norbu

         if (trim(input_constraint%constraint_type) == 'SPATIAL' .or. & 
                   trim(input_constraint%constraint_type) == 'CHARGE_TRANSFER') then
            call dgemm('n', 'n', tmb%orbs%norbu, 1, tmb%orbs%norbu, 1.d0, &
                 proj_mat(1,1,ispin), tmb%orbs%norbu, &
                 inv_ovrlp(1)%matrix(1,iiorbs,ispin), tmb%orbs%norbu, &
                 1.d0, proj_ovrlp_half(1,iorb), tmb%orbs%norbu)
         else
            call dgemm('n', 'n', tmb%orbs%norbu, 1, tmb%orbs%norbu, 1.d0, &
                 proj_mat(1,1,ispin), tmb%orbs%norbu, &
                 tmb%linmat%ovrlp_%matrix(1,iiorbs,ispin), tmb%orbs%norbu, &
                 1.d0, proj_ovrlp_half(1,iorb), tmb%orbs%norbu)
         end if
     end do
  end if

  call f_free(proj_mat)

  weight_matrixp=f_malloc0((/tmb%orbs%norbu,tmb%orbs%norbp/), id='weight_matrixp')
  if (tmb%orbs%norbp>0) then
     do iorb=1,tmb%orbs%norbp
         iiorb=tmb%orbs%isorb+iorb
         if (iiorb<=tmb%orbs%norbu) then
             ispin=1
         else
             ispin=2
         end if
         iiorbs=iiorb-(ispin-1)*tmb%orbs%norbu

         ! calculate full weight matrix
         if (.not. calc_gradient) then
           if (trim(input_constraint%constraint_type) == 'SPATIAL' .or.&
                    trim(input_constraint%constraint_type) == 'CHARGE_TRANSFER') then
               call dgemm('t', 'n', tmb%orbs%norbu, 1, tmb%orbs%norbu, 1.d0, &
                    inv_ovrlp(1)%matrix(1,1,ispin), tmb%orbs%norbu, &
                    proj_ovrlp_half(1,iorb), tmb%orbs%norbu, &
                    1.d0, weight_matrixp(1,iorb), tmb%orbs%norbu)
            else
               call dgemm('t', 'n', tmb%orbs%norbu, 1, tmb%orbs%norbu, 1.d0, &
                    tmb%linmat%ovrlp_%matrix(1,1,ispin), tmb%orbs%norbu, &
                    proj_ovrlp_half(1,iorb), tmb%orbs%norbu, &
                    1.d0, weight_matrixp(1,iorb), tmb%orbs%norbu)
            end if

           ! calculate gradient
           else
             if (trim(input_constraint%constraint_type) == 'SPATIAL' .or.&
                     trim(input_constraint%constraint_type) == 'CHARGE_TRANSFER') then
                stop 'CDFT SF gradient not available for spatial constraint'
             else
                call vcopy(tmb%orbs%norbu, proj_ovrlp_half(1,iorb), 1, weight_matrixp(1,iorb), 1)
             end if
           end if

        end do
  end if

  call f_free(proj_ovrlp_half)

  weight_matrix_%matrix = sparsematrix_malloc_ptr(weight_matrix,iaction=DENSE_FULL,id='weight_matrix_%matrix')
  if (bigdft_mpi%nproc>1) then
     call mpi_allgatherv(weight_matrixp, tmb%orbs%norbu*tmb%orbs%norbp, mpi_double_precision, &
          weight_matrix_%matrix, tmb%orbs%norbu*tmb%orbs%norb_par(:,0), tmb%orbs%norbu*tmb%orbs%isorb_par, &
          mpi_double_precision, bigdft_mpi%mpi_comm, ierr)
  else
     call vcopy(tmb%orbs%norbu*tmb%orbs%norbu*weight_matrix%nspin,weight_matrixp(1,1),1,weight_matrix_%matrix(1,1,1),1)
  end if
  call f_free(weight_matrixp)

  !!! actually want the transpose?
  !!if (calc_gradient) then
  !!!!if (bigdft_mpi%iproc==0) then
  !!!!  hh = 0.0d0
  !!!!  write(2727,*) 'weight matrix'
  !!  do ispin=1,weight_matrix%nspin
  !!     do iorb=1,tmb%orbs%norbu
  !!        do jorb=1,iorb-1
  !!          hh = weight_matrix_%matrix(iorb,jorb,ispin)
  !!          weight_matrix_%matrix(iorb,jorb,ispin) = weight_matrix_%matrix(jorb,iorb,ispin)
  !!          weight_matrix_%matrix(jorb,iorb,ispin) = hh 
  !!!!        do jorb=1,tmb%orbs%norbu
  !!!!           write(2727,*) ispin,iorb,jorb,weight_matrix_%matrix(iorb,jorb,ispin)
  !!!!           hh = max(hh, weight_matrix_%matrix(iorb,jorb,ispin) - weight_matrix_%matrix(jorb,iorb,ispin))
  !!        end do
  !!     end do
  !!  end do
  !!!!  write(2727,*) ''
  !!!!  print*,'Max asymmetry of weight matrix',hh
  !!end if

  !LR: could add a check here that we aren't losing too much information by converting from dense -> sparse
  ! ok in gas phase if using large enough rloc_kernel but could eventually be a problem again
  call compress_matrix(bigdft_mpi%iproc,bigdft_mpi%nproc,weight_matrix,weight_matrix_%matrix,weight_matrix_%matrix_compr)
  call f_free_ptr(weight_matrix_%matrix)

  ! add weight matrix to kernel for initial guess - this is ok since for CDFT calculations ham has same sparsity as kernel
  if (frac_w_guess /= 0.0d0 .and. first_it) then
     call daxpy(tmb%linmat%smat(2)%nvctr*tmb%linmat%smat(2)%nspin,frac_w_guess,&
         weight_matrix_%matrix_compr,1,tmb%linmat%kernel_%matrix_compr,1)
  end if

  call deallocate_matrices(inv_ovrlp(1))
  if (calculate_ovrlp_half) call f_free_ptr(tmb%linmat%ovrlp_%matrix)

  ! if this wasn't the first iteration the overlap is no longer the overlap between the current SFs but rather between current and GS SFs
  ! therefore we need to recalculate it - could save it instead but this operation shouldn't be that expensive anyway
  if (.not. first_it) then
     call calculate_overlap_transposed(bigdft_mpi%iproc, bigdft_mpi%nproc, tmb%orbs, tmb%collcom, tmb%psit_c, &
          tmb%psit_c, tmb%psit_f, tmb%psit_f, tmb%linmat%smat(1), tmb%linmat%auxs, tmb%linmat%ovrlp_)
     !!call gather_matrix_from_taskgroups_inplace(bigdft_mpi%iproc, bigdft_mpi%nproc, tmb%linmat%smat(1), tmb%linmat%ovrlp_)
  end if  

  call f_release_routine()

end subroutine calculate_weight_matrix_lowdin


subroutine get_proj_mat(constraint, input_constraint, input_frag, ref_frags, tmb, proj_mat)
  use module_precisions
  use constrained_dft
  use module_types
  use module_fragments
  use public_enums
  use module_bigdft_arrays
  implicit none
  type(constraint_data), intent(inout) :: constraint
  type(input_constraint_data), intent(inout) :: input_constraint
  type(fragmentInputParameters),intent(in) :: input_frag
  type(dft_wavefunction), intent(in) :: tmb
  type(system_fragment), dimension(input_frag%nfrag_ref), intent(in) :: ref_frags
  real(kind=gp), dimension(tmb%orbs%norbu,tmb%orbs%norbu,tmb%linmat%smat(1)%nspin), intent(inout) :: proj_mat

  integer :: isforb, iorb, jorb, ispin, i, ifrag, ifrag_ref
  real(kind=gp) :: fac


  select case(trim(input_constraint%constraint_type))
  case('SPATIAL')

     do i=1,input_constraint%nfrags_spatial
        ifrag = input_constraint%spatial_frag_constraints(i)%frag
        ispin = input_constraint%spatial_frag_constraints(i)%spin
        call get_fragment_isforb(ref_frags,input_frag,ifrag,isforb,ifrag_ref)
        do iorb=1,ref_frags(ifrag_ref)%fbasis%forbs%norbu
           proj_mat(iorb+isforb,iorb+isforb,ispin) = 1.0_gp
        end do
     end do

  case('CHARGE_TRANSFER')

     do i=1,input_constraint%nfrags_source
        ifrag = input_constraint%source_frag_constraints(i)%frag
        ispin = input_constraint%source_frag_constraints(i)%spin
        call get_fragment_isforb(ref_frags,input_frag,ifrag,isforb,ifrag_ref)
        do iorb=1,ref_frags(ifrag_ref)%fbasis%forbs%norbu
           proj_mat(iorb+isforb,iorb+isforb,ispin) = -0.5_gp
        end do
     end do

     do i=1,input_constraint%nfrags_dest
        ifrag = input_constraint%dest_frag_constraints(i)%frag
        ispin = input_constraint%dest_frag_constraints(i)%spin
        call get_fragment_isforb(ref_frags,input_frag,ifrag,isforb,ifrag_ref)
        do iorb=1,ref_frags(ifrag_ref)%fbasis%forbs%norbu
           proj_mat(iorb+isforb,iorb+isforb,ispin) = -0.5_gp
        end do
     end do

  case('ELECTRONIC')

     input_constraint%orbital = f_malloc0_ptr(tmb%orbs%norbu,id='input_constraint%orbital')
     do i=1,input_constraint%nfrags_orbital
        call get_cdft_orbital(tmb%orbs%norbu,input_constraint%orbital,input_constraint%dest_frag_constraints(i),&
                              input_frag,ref_frags)
     end do

     ! what to do if nfrags is > 1 and the spins are different? do we just forbid this from happening?
     ispin = input_constraint%dest_frag_constraints(1)%spin
     do iorb=1,tmb%orbs%norbu
        proj_mat(iorb,jorb,ispin) = input_constraint%orbital(iorb) * input_constraint%orbital(iorb)
     end do
     call f_free_ptr(input_constraint%orbital)

  case('OPTICAL')

     input_constraint%electron = f_malloc0_ptr(tmb%orbs%norbu,id='input_constraint%electron')
     do i=1,input_constraint%nfrags_electron
        call get_cdft_orbital(tmb%orbs%norbu,input_constraint%electron,input_constraint%electron_frag_constraints(i),&
                              input_frag,ref_frags)
     end do

     input_constraint%hole = f_malloc0_ptr(tmb%orbs%norbu,id='input_constraint%hole')
     do i=1,input_constraint%nfrags_hole
        call get_cdft_orbital(tmb%orbs%norbu,input_constraint%hole,input_constraint%hole_frag_constraints(i),&
                              input_frag,ref_frags)
     end do

     ispin = 1
     do iorb=1,tmb%orbs%norbu
       do jorb=1,tmb%orbs%norbu
          proj_mat(iorb,jorb,ispin) = 0.5d0 * (input_constraint%hole(iorb) * input_constraint%electron(jorb) &
                                             + input_constraint%electron(iorb) * input_constraint%hole(jorb))
       end do
     end do

     if (input_constraint%excitation_type == 'TRIPLET') then
        ispin = 2
        do iorb=1,tmb%orbs%norbu
          do jorb=1,tmb%orbs%norbu
             proj_mat(iorb,jorb,ispin) = -0.5d0 * (input_constraint%hole(iorb) * input_constraint%electron(jorb) &
                                                   + input_constraint%electron(iorb) * input_constraint%hole(jorb))
          end do
        end do
     end if

     call f_free_ptr(input_constraint%electron)
     call f_free_ptr(input_constraint%hole)

  end select

end subroutine get_proj_mat


subroutine get_cdft_orbital(norb,orbital,input_frag_constraint,input_frag,ref_frags)
  use module_precisions
  use module_types
  use module_fragments
  use constrained_dft
  use module_bigdft_mpi
  use module_bigdft_arrays
  implicit none
  integer, intent(in) :: norb
  real(gp), dimension(norb), intent(inout) :: orbital
  type(fragment_constraint_components), intent(in) :: input_frag_constraint
  type(fragmentInputParameters), intent(in) :: input_frag
  type(system_fragment), dimension(input_frag%nfrag_ref), intent(in) :: ref_frags

  integer :: icomp, isforb, ifrag, ifrag_ref, ispin, ncomp, iorb
  
  ifrag = input_frag_constraint%frag
  ispin = input_frag_constraint%spin
  ncomp = input_frag_constraint%ncomp
  call get_fragment_isforb(ref_frags,input_frag,ifrag,isforb,ifrag_ref)
  do icomp=1,ncomp
     do iorb=1,ref_frags(ifrag_ref)%fbasis%forbs%norbu
        orbital(isforb + iorb) = orbital(isforb + iorb) + input_frag_constraint%orbital_weights(icomp) &
                               * ref_frags(ifrag_ref)%coeff(iorb, input_frag_constraint%orbital_indices(icomp), ispin)
     end do
  end do

end subroutine get_cdft_orbital



