!--------------------------------------------------------------------------------------------------!
!   CP2K: A general program to perform molecular dynamics simulations                              !
!   Copyright 2000-2025 CP2K developers group <https://cp2k.org>                                   !
!                                                                                                  !
!   SPDX-License-Identifier: GPL-2.0-or-later                                                      !
!--------------------------------------------------------------------------------------------------!

! **************************************************************************************************
!> \brief compute mulliken charges
!>      we (currently) define them as c_i = 1/2 [ (PS)_{ii} + (SP)_{ii} ]
!> \author Joost VandeVondele March 2003
! **************************************************************************************************
MODULE mulliken
   USE atomic_charges,                  ONLY: print_atomic_charges
   USE cp_control_types,                ONLY: mulliken_restraint_type
   USE cp_dbcsr_api,                    ONLY: &
        dbcsr_get_block_p, dbcsr_get_info, dbcsr_iterator_blocks_left, dbcsr_iterator_next_block, &
        dbcsr_iterator_start, dbcsr_iterator_stop, dbcsr_iterator_type, dbcsr_p_type, dbcsr_type
   USE kinds,                           ONLY: dp
   USE message_passing,                 ONLY: mp_para_env_type
   USE particle_types,                  ONLY: particle_type
   USE qs_kind_types,                   ONLY: qs_kind_type
#include "./base/base_uses.f90"

   IMPLICIT NONE

   PRIVATE

   CHARACTER(len=*), PARAMETER, PRIVATE :: moduleN = 'mulliken'

! *** Public subroutines ***

   PUBLIC :: mulliken_charges, ao_charges, mulliken_restraint, compute_bond_order
   PUBLIC :: compute_charges, atom_trace

   INTERFACE mulliken_charges
      MODULE PROCEDURE mulliken_charges_a, mulliken_charges_b, mulliken_charges_c, &
         mulliken_charges_s, &
         mulliken_charges_akp, mulliken_charges_bkp, mulliken_charges_ckp
   END INTERFACE

   INTERFACE ao_charges
      MODULE PROCEDURE ao_charges_1, ao_charges_2, ao_charges_kp, ao_charges_kp_2
   END INTERFACE

CONTAINS

! **************************************************************************************************
!> \brief computes the energy and density matrix derivate of a constraint on the
!>      mulliken charges
!>
!>      optional outputs:
!>      computes energy (added)
!>      contribution to KS matrix (added)
!>      contribution to W  matrix (added)
!> \param mulliken_restraint_control additional parameters needed to control the restraint
!> \param para_env para_env of the matrices
!> \param s_matrix ,p_matrix : containing the respective quantities
!> \param p_matrix ...
!> \param energy ...
!> \param order_p ...
!> \param ks_matrix ...
!> \param w_matrix ...
!> \par History
!>      06.2004 created [Joost VandeVondele]
!> \note
!>      contribution to the KS matrix is derivative wrt P
!>      contribution to the W matrix is derivate wrt S (sign?)
!>      needed for orbital and ionic forces respectively
! **************************************************************************************************
   SUBROUTINE mulliken_restraint(mulliken_restraint_control, para_env, &
                                 s_matrix, p_matrix, energy, order_p, ks_matrix, w_matrix)
      TYPE(mulliken_restraint_type), INTENT(IN)          :: mulliken_restraint_control
      TYPE(mp_para_env_type), POINTER                    :: para_env
      TYPE(dbcsr_type), POINTER                          :: s_matrix
      TYPE(dbcsr_p_type), DIMENSION(:), POINTER          :: p_matrix
      REAL(KIND=dp), OPTIONAL                            :: energy, order_p
      TYPE(dbcsr_p_type), DIMENSION(:), OPTIONAL, &
         POINTER                                         :: ks_matrix, w_matrix

      INTEGER                                            :: blk, iblock_col, iblock_row, ispin, &
                                                            nblock, nspin
      LOGICAL                                            :: found
      REAL(kind=dp)                                      :: mult, my_energy, my_order_p
      REAL(KIND=dp), DIMENSION(:, :), POINTER            :: charges, charges_deriv, ks_block, &
                                                            p_block, s_block, w_block
      TYPE(dbcsr_iterator_type)                          :: iter

! here we get the numbers for charges

      nspin = SIZE(p_matrix)
      CALL dbcsr_get_info(s_matrix, nblkrows_total=nblock)

      ALLOCATE (charges(nblock, nspin))
      ALLOCATE (charges_deriv(nblock, nspin))
      CALL compute_charges(p_matrix, s_matrix, charges, para_env)
      !
      ! this can be used to check the correct implementation of the derivative
      ! CALL rf_deriv_check(mulliken_restraint_control,charges)
      !
      CALL restraint_functional(mulliken_restraint_control, &
                                charges, charges_deriv, my_energy, my_order_p)

      IF (PRESENT(order_p)) THEN
         order_p = my_order_p
      END IF
      IF (PRESENT(energy)) THEN
         energy = my_energy
      END IF

      IF (PRESENT(ks_matrix)) THEN

         DO ispin = 1, nspin
            CALL dbcsr_iterator_start(iter, s_matrix)
            DO WHILE (dbcsr_iterator_blocks_left(iter))
               CALL dbcsr_iterator_next_block(iter, iblock_row, iblock_col, s_block, blk)
               CALL dbcsr_get_block_p(matrix=ks_matrix(ispin)%matrix, &
                                      row=iblock_row, col=iblock_col, BLOCK=ks_block, found=found)

               IF (.NOT. (ASSOCIATED(s_block) .AND. ASSOCIATED(ks_block))) THEN
                  CPABORT("Unexpected s / ks structure")
               END IF
               mult = 0.5_dp*charges_deriv(iblock_row, ispin) + &
                      0.5_dp*charges_deriv(iblock_col, ispin)

               ks_block = ks_block + mult*s_block

            END DO
            CALL dbcsr_iterator_stop(iter)
         END DO

      END IF

      IF (PRESENT(w_matrix)) THEN

         DO ispin = 1, nspin
            CALL dbcsr_iterator_start(iter, p_matrix(ispin)%matrix)
            DO WHILE (dbcsr_iterator_blocks_left(iter))
               CALL dbcsr_iterator_next_block(iter, iblock_row, iblock_col, p_block, blk)
               CALL dbcsr_get_block_p(matrix=w_matrix(ispin)%matrix, &
                                      row=iblock_row, col=iblock_col, BLOCK=w_block, found=found)

               ! we can cycle if a block is not present
               IF (.NOT. (ASSOCIATED(w_block) .AND. ASSOCIATED(p_block))) CYCLE

               ! minus sign relates to convention for W
               mult = -0.5_dp*charges_deriv(iblock_row, ispin) &
                      - 0.5_dp*charges_deriv(iblock_col, ispin)

               w_block = w_block + mult*p_block

            END DO
            CALL dbcsr_iterator_stop(iter)
         END DO

      END IF

      DEALLOCATE (charges)
      DEALLOCATE (charges_deriv)

   END SUBROUTINE mulliken_restraint

! **************************************************************************************************
!> \brief computes energy and derivatives given a set of charges
!>       this implementation uses the spin density on a number of atoms
!>       as a penalty function
!> \param mulliken_restraint_control ...
!> \param charges (nblock,nspin)
!> \param charges_deriv derivate wrt the corresponding charge entry
!> \param energy ...
!> \param order_p ...
!> \par History
!>      06.2004 created [Joost VandeVondele]
!>      02.2005 added more general form [Joost VandeVondele]
!> \note
!>       should be easy to adapt for other specialized cases
! **************************************************************************************************
   SUBROUTINE restraint_functional(mulliken_restraint_control, charges, &
                                   charges_deriv, energy, order_p)
      TYPE(mulliken_restraint_type), INTENT(IN)          :: mulliken_restraint_control
      REAL(KIND=dp), DIMENSION(:, :), POINTER            :: charges, charges_deriv
      REAL(KIND=dp), INTENT(OUT)                         :: energy, order_p

      INTEGER                                            :: I
      REAL(KIND=dp)                                      :: dum

      charges_deriv = 0.0_dp
      order_p = 0.0_dp

      DO I = 1, mulliken_restraint_control%natoms
         order_p = order_p + charges(mulliken_restraint_control%atoms(I), 1) &
                   - charges(mulliken_restraint_control%atoms(I), 2) ! spin density on the relevant atoms
      END DO
      ! energy
      energy = mulliken_restraint_control%strength*(order_p - mulliken_restraint_control%target)**2
      ! derivative
      dum = 2*mulliken_restraint_control%strength*(order_p - mulliken_restraint_control%target)
      DO I = 1, mulliken_restraint_control%natoms
         charges_deriv(mulliken_restraint_control%atoms(I), 1) = dum
         charges_deriv(mulliken_restraint_control%atoms(I), 2) = -dum
      END DO
   END SUBROUTINE restraint_functional

! **************************************************************************************************
!> \brief compute the mulliken charges
!> \param p_matrix , s_matrix, para_env ...
!> \param s_matrix ...
!> \param charges previously allocated with the right size (natom,nspin)
!> \param para_env ...
!> \par History
!>      06.2004 created [Joost VandeVondele]
!> \note
!>      charges are computed per spin in the LSD case
! **************************************************************************************************
   SUBROUTINE compute_charges(p_matrix, s_matrix, charges, para_env)
      TYPE(dbcsr_p_type), DIMENSION(:), POINTER          :: p_matrix
      TYPE(dbcsr_type), POINTER                          :: s_matrix
      REAL(KIND=dp), DIMENSION(:, :), POINTER            :: charges
      TYPE(mp_para_env_type), POINTER                    :: para_env

      INTEGER                                            :: blk, iblock_col, iblock_row, ispin, nspin
      LOGICAL                                            :: found
      REAL(kind=dp)                                      :: mult
      REAL(KIND=dp), DIMENSION(:, :), POINTER            :: p_block, s_block
      TYPE(dbcsr_iterator_type)                          :: iter

      nspin = SIZE(p_matrix)

      charges = 0.0_dp
      DO ispin = 1, nspin
         CALL dbcsr_iterator_start(iter, s_matrix)
         DO WHILE (dbcsr_iterator_blocks_left(iter))
            NULLIFY (s_block, p_block)
            CALL dbcsr_iterator_next_block(iter, iblock_row, iblock_col, s_block, blk)
            CALL dbcsr_get_block_p(matrix=p_matrix(ispin)%matrix, &
                                   row=iblock_row, col=iblock_col, BLOCK=p_block, found=found)

            ! we can cycle if a block is not present
            IF (.NOT. found) CYCLE
            IF (.NOT. (ASSOCIATED(s_block) .AND. ASSOCIATED(p_block))) CYCLE

            IF (iblock_row .EQ. iblock_col) THEN
               mult = 0.5_dp ! avoid double counting of diagonal blocks
            ELSE
               mult = 1.0_dp
            END IF
            charges(iblock_row, ispin) = charges(iblock_row, ispin) + &
                                         mult*SUM(p_block*s_block)
            charges(iblock_col, ispin) = charges(iblock_col, ispin) + &
                                         mult*SUM(p_block*s_block)

         END DO
         CALL dbcsr_iterator_stop(iter)
      END DO
      CALL para_env%sum(charges)

   END SUBROUTINE compute_charges

! **************************************************************************************************
!> \brief compute the mulliken charges for single matrices
!> \param p_matrix , s_matrix, para_env ...
!> \param s_matrix ...
!> \param charges previously allocated with the right size (natom,nspin)
!> \param para_env ...
!> \par History
!>      06.2004 created [Joost VandeVondele]
! **************************************************************************************************
   SUBROUTINE compute_charges_single(p_matrix, s_matrix, charges, para_env)
      TYPE(dbcsr_type)                                   :: p_matrix, s_matrix
      REAL(KIND=dp), DIMENSION(:)                        :: charges
      TYPE(mp_para_env_type), POINTER                    :: para_env

      INTEGER                                            :: blk, iblock_col, iblock_row
      LOGICAL                                            :: found
      REAL(kind=dp)                                      :: mult
      REAL(KIND=dp), DIMENSION(:, :), POINTER            :: p_block, s_block
      TYPE(dbcsr_iterator_type)                          :: iter

      charges = 0.0_dp
      CALL dbcsr_iterator_start(iter, s_matrix)
      DO WHILE (dbcsr_iterator_blocks_left(iter))
         NULLIFY (s_block, p_block)
         CALL dbcsr_iterator_next_block(iter, iblock_row, iblock_col, s_block, blk)
         CALL dbcsr_get_block_p(matrix=p_matrix, &
                                row=iblock_row, col=iblock_col, BLOCK=p_block, found=found)

         ! we can cycle if a block is not present
         IF (.NOT. found) CYCLE
         IF (.NOT. (ASSOCIATED(s_block) .AND. ASSOCIATED(p_block))) CYCLE

         IF (iblock_row .EQ. iblock_col) THEN
            mult = 0.5_dp ! avoid double counting of diagonal blocks
         ELSE
            mult = 1.0_dp
         END IF
         charges(iblock_row) = charges(iblock_row) + mult*SUM(p_block*s_block)
         charges(iblock_col) = charges(iblock_col) + mult*SUM(p_block*s_block)
      END DO
      CALL dbcsr_iterator_stop(iter)

      CALL para_env%sum(charges)

   END SUBROUTINE compute_charges_single

! **************************************************************************************************
!> \brief compute the mulliken charge derivatives
!> \param p_matrix , s_matrix, para_env ...
!> \param s_matrix ...
!> \param charges ...
!> \param dcharges previously allocated with the right size (natom,3)
!> \param para_env ...
!> \par History
!>      01.2012 created [JHU]
! **************************************************************************************************
   SUBROUTINE compute_dcharges(p_matrix, s_matrix, charges, dcharges, para_env)
      TYPE(dbcsr_p_type), DIMENSION(:), POINTER          :: p_matrix, s_matrix
      REAL(KIND=dp), DIMENSION(:, :), POINTER            :: charges, dcharges
      TYPE(mp_para_env_type), POINTER                    :: para_env

      INTEGER                                            :: blk, iblock_col, iblock_row, ider, &
                                                            ispin, nspin
      LOGICAL                                            :: found
      REAL(kind=dp)                                      :: mult
      REAL(KIND=dp), DIMENSION(:, :), POINTER            :: ds_block, p_block, s_block
      TYPE(dbcsr_iterator_type)                          :: iter

      nspin = SIZE(p_matrix)

      charges = 0.0_dp
      dcharges = 0.0_dp
      DO ispin = 1, nspin
         CALL dbcsr_iterator_start(iter, s_matrix(1)%matrix)
         DO WHILE (dbcsr_iterator_blocks_left(iter))
            NULLIFY (s_block, p_block)
            CALL dbcsr_iterator_next_block(iter, iblock_row, iblock_col, s_block, blk)
            CALL dbcsr_get_block_p(matrix=p_matrix(ispin)%matrix, &
                                   row=iblock_row, col=iblock_col, BLOCK=p_block, found=found)

            ! we can cycle if a block is not present
            IF (.NOT. found) CYCLE
            IF (.NOT. (ASSOCIATED(s_block) .AND. ASSOCIATED(p_block))) CYCLE

            IF (iblock_row .EQ. iblock_col) THEN
               mult = 0.5_dp ! avoid double counting of diagonal blocks
            ELSE
               mult = 1.0_dp
            END IF
            charges(iblock_row, ispin) = charges(iblock_row, ispin) + mult*SUM(p_block*s_block)
            charges(iblock_col, ispin) = charges(iblock_col, ispin) + mult*SUM(p_block*s_block)
            DO ider = 1, 3
               CALL dbcsr_get_block_p(matrix=s_matrix(ider + 1)%matrix, &
                                      row=iblock_row, col=iblock_col, BLOCK=ds_block, found=found)
               dcharges(iblock_row, ider) = dcharges(iblock_row, ider) + mult*SUM(p_block*ds_block)
               dcharges(iblock_col, ider) = dcharges(iblock_col, ider) + mult*SUM(p_block*ds_block)
            END DO

         END DO
         CALL dbcsr_iterator_stop(iter)
      END DO
      CALL para_env%sum(charges)
      CALL para_env%sum(dcharges)

   END SUBROUTINE compute_dcharges

! **************************************************************************************************
!> \brief print the mulliken charges to scr on ionode
!> \param p_matrix , s_matrix, para_env ...
!> \param s_matrix ...
!> \param para_env ...
!> \param particle_set (needed for Z)
!> \param qs_kind_set ...
!> \param scr unit for output
!> \param title ...
!> \par History
!>      06.2004 adapted to remove explicit matrix multiply [Joost VandeVondele]
! **************************************************************************************************
   SUBROUTINE mulliken_charges_a(p_matrix, s_matrix, para_env, particle_set, &
                                 qs_kind_set, scr, title)

      TYPE(dbcsr_p_type), DIMENSION(:), POINTER          :: p_matrix
      TYPE(dbcsr_type), POINTER                          :: s_matrix
      TYPE(mp_para_env_type), POINTER                    :: para_env
      TYPE(particle_type), DIMENSION(:), POINTER         :: particle_set
      TYPE(qs_kind_type), DIMENSION(:), POINTER          :: qs_kind_set
      INTEGER                                            :: scr
      CHARACTER(LEN=*)                                   :: title

      CHARACTER(len=*), PARAMETER :: routineN = 'mulliken_charges_a'

      INTEGER                                            :: handle, nblock, nspin
      REAL(KIND=dp), DIMENSION(:, :), POINTER            :: charges

      CALL timeset(routineN, handle)

      CPASSERT(ASSOCIATED(p_matrix))
      CPASSERT(ASSOCIATED(s_matrix))
      ! here we get the numbers for charges
      nspin = SIZE(p_matrix)
      CALL dbcsr_get_info(s_matrix, nblkrows_total=nblock)
      ALLOCATE (charges(nblock, nspin))

      CALL compute_charges(p_matrix, s_matrix, charges, para_env)

      CALL print_atomic_charges(particle_set, qs_kind_set, scr, title, electronic_charges=charges)

      DEALLOCATE (charges)

      CALL timestop(handle)

   END SUBROUTINE mulliken_charges_a

! **************************************************************************************************
!> \brief ...
!> \param p_matrix ...
!> \param s_matrix ...
!> \param para_env ...
!> \param mcharge ...
! **************************************************************************************************
   SUBROUTINE mulliken_charges_b(p_matrix, s_matrix, para_env, mcharge)

      TYPE(dbcsr_p_type), DIMENSION(:), POINTER          :: p_matrix
      TYPE(dbcsr_type), POINTER                          :: s_matrix
      TYPE(mp_para_env_type), POINTER                    :: para_env
      REAL(KIND=dp), DIMENSION(:, :), POINTER            :: mcharge

      CHARACTER(len=*), PARAMETER :: routineN = 'mulliken_charges_b'

      INTEGER                                            :: handle

      CALL timeset(routineN, handle)

      IF (ASSOCIATED(p_matrix) .AND. ASSOCIATED(s_matrix)) THEN
         CALL compute_charges(p_matrix, s_matrix, mcharge, para_env)
      END IF

      CALL timestop(handle)

   END SUBROUTINE mulliken_charges_b

! **************************************************************************************************
!> \brief ...
!> \param p_matrix ...
!> \param s_matrix ...
!> \param para_env ...
!> \param mcharge ...
!> \param dmcharge ...
! **************************************************************************************************
   SUBROUTINE mulliken_charges_c(p_matrix, s_matrix, para_env, mcharge, dmcharge)

      TYPE(dbcsr_p_type), DIMENSION(:), POINTER          :: p_matrix, s_matrix
      TYPE(mp_para_env_type), POINTER                    :: para_env
      REAL(KIND=dp), DIMENSION(:, :), POINTER            :: mcharge, dmcharge

      CHARACTER(len=*), PARAMETER :: routineN = 'mulliken_charges_c'

      INTEGER                                            :: handle

      CALL timeset(routineN, handle)

      IF (ASSOCIATED(p_matrix) .AND. ASSOCIATED(s_matrix)) THEN
         CALL compute_dcharges(p_matrix, s_matrix, mcharge, dmcharge, para_env)
      END IF

      CALL timestop(handle)

   END SUBROUTINE mulliken_charges_c

! **************************************************************************************************
!> \brief ...
!> \param p_matrix ...
!> \param s_matrix ...
!> \param para_env ...
!> \param mcharge ...
! **************************************************************************************************
   SUBROUTINE mulliken_charges_s(p_matrix, s_matrix, para_env, mcharge)

      TYPE(dbcsr_type)                                   :: p_matrix, s_matrix
      TYPE(mp_para_env_type), POINTER                    :: para_env
      REAL(KIND=dp), DIMENSION(:)                        :: mcharge

      CHARACTER(len=*), PARAMETER :: routineN = 'mulliken_charges_s'

      INTEGER                                            :: handle

      CALL timeset(routineN, handle)

      CALL compute_charges_single(p_matrix, s_matrix, mcharge, para_env)

      CALL timestop(handle)

   END SUBROUTINE mulliken_charges_s

! **************************************************************************************************
!> \brief print the mulliken charges to scr on ionode
!> \param p_matrix_kp ...
!> \param s_matrix_kp ...
!> \param para_env ...
!> \param particle_set (needed for Z)
!> \param qs_kind_set ...
!> \param scr unit for output
!> \param title ...
!> \par History
!>      06.2004 adapted to remove explicit matrix multiply [Joost VandeVondele]
! **************************************************************************************************
   SUBROUTINE mulliken_charges_akp(p_matrix_kp, s_matrix_kp, para_env, particle_set, &
                                   qs_kind_set, scr, title)

      TYPE(dbcsr_p_type), DIMENSION(:, :), POINTER       :: p_matrix_kp, s_matrix_kp
      TYPE(mp_para_env_type), POINTER                    :: para_env
      TYPE(particle_type), DIMENSION(:), POINTER         :: particle_set
      TYPE(qs_kind_type), DIMENSION(:), POINTER          :: qs_kind_set
      INTEGER                                            :: scr
      CHARACTER(LEN=*)                                   :: title

      CHARACTER(len=*), PARAMETER :: routineN = 'mulliken_charges_akp'

      INTEGER                                            :: handle, ic, nblock, nspin
      REAL(KIND=dp), DIMENSION(:, :), POINTER            :: charges, charges_im
      TYPE(dbcsr_p_type), DIMENSION(:), POINTER          :: p_matrix
      TYPE(dbcsr_type), POINTER                          :: s_matrix

      CALL timeset(routineN, handle)

      CPASSERT(ASSOCIATED(p_matrix_kp))
      CPASSERT(ASSOCIATED(s_matrix_kp))

      nspin = SIZE(p_matrix_kp, 1)
      CALL dbcsr_get_info(s_matrix_kp(1, 1)%matrix, nblkrows_total=nblock)
      ALLOCATE (charges(nblock, nspin), charges_im(nblock, nspin))
      charges = 0.0_dp

      DO ic = 1, SIZE(s_matrix_kp, 2)
         NULLIFY (p_matrix, s_matrix)
         p_matrix => p_matrix_kp(:, ic)
         s_matrix => s_matrix_kp(1, ic)%matrix
         charges_im = 0.0_dp
         CALL compute_charges(p_matrix, s_matrix, charges_im, para_env)
         charges(:, :) = charges(:, :) + charges_im(:, :)
      END DO

      CALL print_atomic_charges(particle_set, qs_kind_set, scr, title, electronic_charges=charges)

      DEALLOCATE (charges, charges_im)

      CALL timestop(handle)

   END SUBROUTINE mulliken_charges_akp

! **************************************************************************************************
!> \brief ...
!> \param p_matrix_kp ...
!> \param s_matrix_kp ...
!> \param para_env ...
!> \param mcharge ...
! **************************************************************************************************
   SUBROUTINE mulliken_charges_bkp(p_matrix_kp, s_matrix_kp, para_env, mcharge)

      TYPE(dbcsr_p_type), DIMENSION(:, :), POINTER       :: p_matrix_kp, s_matrix_kp
      TYPE(mp_para_env_type), POINTER                    :: para_env
      REAL(KIND=dp), DIMENSION(:, :), POINTER            :: mcharge

      CHARACTER(len=*), PARAMETER :: routineN = 'mulliken_charges_bkp'

      INTEGER                                            :: handle, ic, natom, nspin
      REAL(KIND=dp), DIMENSION(:, :), POINTER            :: mcharge_im
      TYPE(dbcsr_p_type), DIMENSION(:), POINTER          :: p_matrix
      TYPE(dbcsr_type), POINTER                          :: s_matrix

      CALL timeset(routineN, handle)

      IF (ASSOCIATED(p_matrix_kp) .AND. ASSOCIATED(s_matrix_kp)) THEN

         mcharge = 0.0_dp
         natom = SIZE(mcharge, 1)
         nspin = SIZE(mcharge, 2)
         ALLOCATE (mcharge_im(natom, nspin))

         DO ic = 1, SIZE(s_matrix_kp, 2)
            NULLIFY (p_matrix, s_matrix)
            p_matrix => p_matrix_kp(:, ic)
            s_matrix => s_matrix_kp(1, ic)%matrix
            IF (ASSOCIATED(p_matrix) .AND. ASSOCIATED(s_matrix)) THEN
               CALL compute_charges(p_matrix, s_matrix, mcharge_im, para_env)
               mcharge(:, :) = mcharge(:, :) + mcharge_im(:, :)
            END IF
         END DO

         DEALLOCATE (mcharge_im)

      END IF

      CALL timestop(handle)

   END SUBROUTINE mulliken_charges_bkp

! **************************************************************************************************
!> \brief ...
!> \param p_matrix_kp ...
!> \param s_matrix_kp ...
!> \param para_env ...
!> \param mcharge ...
!> \param dmcharge ...
! **************************************************************************************************
   SUBROUTINE mulliken_charges_ckp(p_matrix_kp, s_matrix_kp, para_env, &
                                   mcharge, dmcharge)

      TYPE(dbcsr_p_type), DIMENSION(:, :), POINTER       :: p_matrix_kp, s_matrix_kp
      TYPE(mp_para_env_type), POINTER                    :: para_env
      REAL(KIND=dp), DIMENSION(:, :), POINTER            :: mcharge, dmcharge

      CHARACTER(len=*), PARAMETER :: routineN = 'mulliken_charges_ckp'

      INTEGER                                            :: handle, ic, natom, nder, nspin
      REAL(KIND=dp), DIMENSION(:, :), POINTER            :: dmcharge_im, mcharge_im
      TYPE(dbcsr_p_type), DIMENSION(:), POINTER          :: p_matrix, s_matrix

      CALL timeset(routineN, handle)

      IF (ASSOCIATED(p_matrix_kp) .AND. ASSOCIATED(s_matrix_kp)) THEN

         mcharge = 0.0_dp
         dmcharge = 0.0_dp
         natom = SIZE(mcharge, 1)
         nspin = SIZE(mcharge, 2)
         nder = SIZE(dmcharge, 2)
         ALLOCATE (mcharge_im(natom, nspin), dmcharge_im(natom, nder))

         DO ic = 1, SIZE(s_matrix_kp, 2)
            NULLIFY (p_matrix, s_matrix)
            p_matrix => p_matrix_kp(:, ic)
            s_matrix => s_matrix_kp(:, ic)
            IF (ASSOCIATED(p_matrix) .AND. ASSOCIATED(s_matrix)) THEN
               CALL compute_dcharges(p_matrix, s_matrix, mcharge_im, dmcharge_im, para_env)
               mcharge(:, :) = mcharge(:, :) + mcharge_im(:, :)
               dmcharge(:, :) = dmcharge(:, :) + dmcharge_im(:, :)
            END IF
         END DO

         DEALLOCATE (mcharge_im, dmcharge_im)

      END IF

      CALL timestop(handle)

   END SUBROUTINE mulliken_charges_ckp

! **************************************************************************************************
!> \brief compute Mayer bond orders for a single spin channel
!>        for complete result sum up over all spins and multiply by Nspin
!> \param psmat ...
!> \param spmat ...
!> \param bond_order ...
!> \par History
!>      12.2016 created [JGH]
! **************************************************************************************************
   SUBROUTINE compute_bond_order(psmat, spmat, bond_order)
      TYPE(dbcsr_type)                                   :: psmat, spmat
      REAL(KIND=dp), DIMENSION(:, :)                     :: bond_order

      INTEGER                                            :: blk, iat, jat
      LOGICAL                                            :: found
      REAL(KIND=dp), DIMENSION(:, :), POINTER            :: ps, sp
      TYPE(dbcsr_iterator_type)                          :: iter

      CALL dbcsr_iterator_start(iter, psmat)
      DO WHILE (dbcsr_iterator_blocks_left(iter))
         NULLIFY (ps, sp)
         CALL dbcsr_iterator_next_block(iter, iat, jat, ps, blk)
         CALL dbcsr_get_block_p(matrix=spmat, &
                                row=iat, col=jat, BLOCK=sp, found=found)
         IF (.NOT. found) CYCLE
         IF (.NOT. (ASSOCIATED(sp) .AND. ASSOCIATED(ps))) CYCLE

         bond_order(iat, jat) = bond_order(iat, jat) + SUM(ps*sp)
      END DO
      CALL dbcsr_iterator_stop(iter)

   END SUBROUTINE compute_bond_order

! **************************************************************************************************
!> \brief compute the mulliken charges for a single atom (AO resolved)
!> \param p_matrix , s_matrix, para_env ...
!> \param s_matrix ...
!> \param charges ...
!> \param iatom ...
!> \param para_env ...
!> \par History
!>      06.2004 created [Joost VandeVondele]
!>      10.2018 adapted [JGH]
!> \note
! **************************************************************************************************
   SUBROUTINE ao_charges_1(p_matrix, s_matrix, charges, iatom, para_env)
      TYPE(dbcsr_p_type), DIMENSION(:), POINTER          :: p_matrix
      TYPE(dbcsr_type), POINTER                          :: s_matrix
      REAL(KIND=dp), DIMENSION(:), INTENT(INOUT)         :: charges
      INTEGER, INTENT(IN)                                :: iatom
      TYPE(mp_para_env_type), POINTER                    :: para_env

      CHARACTER(len=*), PARAMETER                        :: routineN = 'ao_charges_1'

      INTEGER                                            :: blk, handle, i, iblock_col, iblock_row, &
                                                            ispin, j, nspin
      LOGICAL                                            :: found
      REAL(KIND=dp), DIMENSION(:, :), POINTER            :: p_block, s_block
      TYPE(dbcsr_iterator_type)                          :: iter

      CALL timeset(routineN, handle)

      nspin = SIZE(p_matrix)
      charges = 0.0_dp
      DO ispin = 1, nspin
         CALL dbcsr_iterator_start(iter, s_matrix)
         DO WHILE (dbcsr_iterator_blocks_left(iter))
            NULLIFY (s_block, p_block)
            CALL dbcsr_iterator_next_block(iter, iblock_row, iblock_col, s_block, blk)
            CALL dbcsr_get_block_p(matrix=p_matrix(ispin)%matrix, &
                                   row=iblock_row, col=iblock_col, BLOCK=p_block, found=found)

            ! we can cycle if a block is not present
            IF (.NOT. found) CYCLE
            IF (.NOT. (ASSOCIATED(s_block) .AND. ASSOCIATED(p_block))) CYCLE

            IF (iblock_row == iatom) THEN
               DO j = 1, SIZE(p_block, 2)
                  DO i = 1, SIZE(p_block, 1)
                     charges(i) = charges(i) + p_block(i, j)*s_block(i, j)
                  END DO
               END DO
            ELSEIF (iblock_col == iatom) THEN
               DO j = 1, SIZE(p_block, 2)
                  DO i = 1, SIZE(p_block, 1)
                     charges(j) = charges(j) + p_block(i, j)*s_block(i, j)
                  END DO
               END DO
            END IF
         END DO
         CALL dbcsr_iterator_stop(iter)
      END DO
      CALL para_env%sum(charges)

      CALL timestop(handle)

   END SUBROUTINE ao_charges_1

! **************************************************************************************************
!> \brief compute the mulliken charges (AO resolved)
!> \param p_matrix , s_matrix, para_env ...
!> \param s_matrix ...
!> \param charges ...
!> \param para_env ...
!> \par History
!>      06.2004 created [Joost VandeVondele]
!>      10.2018 adapted [JGH]
!> \note
! **************************************************************************************************
   SUBROUTINE ao_charges_2(p_matrix, s_matrix, charges, para_env)
      TYPE(dbcsr_p_type), DIMENSION(:), POINTER          :: p_matrix
      TYPE(dbcsr_type), POINTER                          :: s_matrix
      REAL(KIND=dp), DIMENSION(:, :), INTENT(INOUT)      :: charges
      TYPE(mp_para_env_type), POINTER                    :: para_env

      CHARACTER(len=*), PARAMETER                        :: routineN = 'ao_charges_2'

      INTEGER                                            :: blk, handle, i, iblock_col, iblock_row, &
                                                            ispin, j, nspin
      LOGICAL                                            :: found
      REAL(KIND=dp), DIMENSION(:, :), POINTER            :: p_block, s_block
      TYPE(dbcsr_iterator_type)                          :: iter

      CALL timeset(routineN, handle)

      nspin = SIZE(p_matrix)
      charges = 0.0_dp
      DO ispin = 1, nspin
         CALL dbcsr_iterator_start(iter, s_matrix)
         DO WHILE (dbcsr_iterator_blocks_left(iter))
            NULLIFY (s_block, p_block)
            CALL dbcsr_iterator_next_block(iter, iblock_row, iblock_col, s_block, blk)
            CALL dbcsr_get_block_p(matrix=p_matrix(ispin)%matrix, &
                                   row=iblock_row, col=iblock_col, BLOCK=p_block, found=found)

            ! we can cycle if a block is not present
            IF (.NOT. found) CYCLE
            IF (.NOT. (ASSOCIATED(s_block) .AND. ASSOCIATED(p_block))) CYCLE

            DO j = 1, SIZE(p_block, 2)
               DO i = 1, SIZE(p_block, 1)
                  charges(i, iblock_row) = charges(i, iblock_row) + p_block(i, j)*s_block(i, j)
               END DO
            END DO
            IF (iblock_col /= iblock_row) THEN
               DO j = 1, SIZE(p_block, 2)
                  DO i = 1, SIZE(p_block, 1)
                     charges(j, iblock_col) = charges(j, iblock_col) + p_block(i, j)*s_block(i, j)
                  END DO
               END DO
            END IF

         END DO
         CALL dbcsr_iterator_stop(iter)
      END DO
      CALL para_env%sum(charges)

      CALL timestop(handle)

   END SUBROUTINE ao_charges_2

! **************************************************************************************************
!> \brief ...
!> \param p_matrix_kp ...
!> \param s_matrix_kp ...
!> \param charges ...
!> \param iatom ...
!> \param para_env ...
! **************************************************************************************************
   SUBROUTINE ao_charges_kp(p_matrix_kp, s_matrix_kp, charges, iatom, para_env)

      TYPE(dbcsr_p_type), DIMENSION(:, :), POINTER       :: p_matrix_kp, s_matrix_kp
      REAL(KIND=dp), DIMENSION(:), INTENT(INOUT)         :: charges
      INTEGER, INTENT(IN)                                :: iatom
      TYPE(mp_para_env_type), POINTER                    :: para_env

      CHARACTER(len=*), PARAMETER                        :: routineN = 'ao_charges_kp'

      INTEGER                                            :: handle, ic, na
      REAL(KIND=dp), ALLOCATABLE, DIMENSION(:)           :: charge_im
      TYPE(dbcsr_p_type), DIMENSION(:), POINTER          :: p_matrix
      TYPE(dbcsr_type), POINTER                          :: s_matrix

      CALL timeset(routineN, handle)

      IF (ASSOCIATED(p_matrix_kp) .AND. ASSOCIATED(s_matrix_kp)) THEN

         charges = 0.0_dp
         na = SIZE(charges, 1)
         ALLOCATE (charge_im(na))

         DO ic = 1, SIZE(s_matrix_kp, 2)
            NULLIFY (p_matrix, s_matrix)
            p_matrix => p_matrix_kp(:, ic)
            s_matrix => s_matrix_kp(1, ic)%matrix
            IF (ASSOCIATED(p_matrix) .AND. ASSOCIATED(s_matrix)) THEN
               CALL ao_charges_1(p_matrix, s_matrix, charge_im, iatom, para_env)
               charges(:) = charges(:) + charge_im(:)
            END IF
         END DO

         DEALLOCATE (charge_im)

      END IF

      CALL timestop(handle)

   END SUBROUTINE ao_charges_kp

! **************************************************************************************************
!> \brief ...
!> \param p_matrix_kp ...
!> \param s_matrix_kp ...
!> \param charges ...
!> \param para_env ...
! **************************************************************************************************
   SUBROUTINE ao_charges_kp_2(p_matrix_kp, s_matrix_kp, charges, para_env)

      TYPE(dbcsr_p_type), DIMENSION(:, :), POINTER       :: p_matrix_kp, s_matrix_kp
      REAL(KIND=dp), DIMENSION(:, :), INTENT(INOUT)      :: charges
      TYPE(mp_para_env_type), POINTER                    :: para_env

      CHARACTER(len=*), PARAMETER                        :: routineN = 'ao_charges_kp_2'

      INTEGER                                            :: handle, ic, na, nb
      REAL(KIND=dp), ALLOCATABLE, DIMENSION(:, :)        :: charge_im
      TYPE(dbcsr_p_type), DIMENSION(:), POINTER          :: p_matrix
      TYPE(dbcsr_type), POINTER                          :: s_matrix

      CALL timeset(routineN, handle)

      IF (ASSOCIATED(p_matrix_kp) .AND. ASSOCIATED(s_matrix_kp)) THEN

         charges = 0.0_dp
         na = SIZE(charges, 1)
         nb = SIZE(charges, 2)
         ALLOCATE (charge_im(na, nb))

         DO ic = 1, SIZE(s_matrix_kp, 2)
            NULLIFY (p_matrix, s_matrix)
            p_matrix => p_matrix_kp(:, ic)
            s_matrix => s_matrix_kp(1, ic)%matrix
            IF (ASSOCIATED(p_matrix) .AND. ASSOCIATED(s_matrix)) THEN
               CALL ao_charges_2(p_matrix, s_matrix, charge_im, para_env)
               charges(:, :) = charges(:, :) + charge_im(:, :)
            END IF
         END DO

         DEALLOCATE (charge_im)

      END IF

      CALL timestop(handle)

   END SUBROUTINE ao_charges_kp_2

! **************************************************************************************************
!> \brief Compute partial trace of product of two matrices
!> \param amat ...
!> \param bmat ...
!> \param factor ...
!> \param atrace ...
!> \par History
!>      06.2004 created [Joost VandeVondele]
!> \note
!>      charges are computed per spin in the LSD case
! **************************************************************************************************
   SUBROUTINE atom_trace(amat, bmat, factor, atrace)
      TYPE(dbcsr_type), POINTER                          :: amat, bmat
      REAL(kind=dp), INTENT(IN)                          :: factor
      REAL(KIND=dp), DIMENSION(:), POINTER               :: atrace

      INTEGER                                            :: blk, iblock_col, iblock_row, nblock
      LOGICAL                                            :: found
      REAL(kind=dp)                                      :: btr, mult
      REAL(KIND=dp), DIMENSION(:, :), POINTER            :: a_block, b_block
      TYPE(dbcsr_iterator_type)                          :: iter

      CALL dbcsr_get_info(bmat, nblkrows_total=nblock)
      CPASSERT(nblock == SIZE(atrace))

      CALL dbcsr_iterator_start(iter, bmat)
      DO WHILE (dbcsr_iterator_blocks_left(iter))
         CALL dbcsr_iterator_next_block(iter, iblock_row, iblock_col, b_block, blk)
         CALL dbcsr_get_block_p(matrix=amat, &
                                row=iblock_row, col=iblock_col, BLOCK=a_block, found=found)

         ! we can cycle if a block is not present
         IF (.NOT. (ASSOCIATED(b_block) .AND. ASSOCIATED(a_block))) CYCLE

         IF (iblock_row .EQ. iblock_col) THEN
            mult = 0.5_dp ! avoid double counting of diagonal blocks
         ELSE
            mult = 1.0_dp
         END IF
         btr = factor*mult*SUM(a_block*b_block)
         atrace(iblock_row) = atrace(iblock_row) + btr
         atrace(iblock_col) = atrace(iblock_col) + btr

      END DO
      CALL dbcsr_iterator_stop(iter)

   END SUBROUTINE atom_trace

! **************************************************************************************************

END MODULE mulliken
