!--------------------------------------------------------------------------------------------------!
!   CP2K-RS Extended Interface: Main Module for Quickstep DFT Access                              !
!   Copyright 2025 CP2K-RS developers                                                             !
!                                                                                                  !
!   SPDX-License-Identifier: GPL-2.0-or-later                                                      !
!--------------------------------------------------------------------------------------------------!

! **************************************************************************************************
!> \brief Extended CP2K C/C++ interface for accessing Quickstep DFT internals
!> \par History
!>       01.2025 created as part of cp2k-rs extensions
!> \author CP2K-RS Developers
!> \note This module extends the standard libcp2k interface with functionality to access
!>       internal DFT data structures from CP2K's Quickstep code, including eigenvalues,
!>       matrices, charges, and other electronic structure information.
! **************************************************************************************************
MODULE libcp2k_extended
    USE ISO_C_BINDING,                   ONLY: C_BOOL,&
                                               C_CHAR,&
                                               C_DOUBLE,&
                                               C_INT,&
                                               C_NULL_CHAR,&
                                               C_PTR
   USE kinds,                           ONLY: dp,&
                                              default_string_length
   USE f77_interface,                   ONLY: f_env_add_defaults,&
                                              f_env_rm_defaults,&
                                              f_env_type
   USE force_env_types,                 ONLY: force_env_get,&
                                              use_qs_force
   USE qs_environment_types,            ONLY: get_qs_env,&
                                              qs_environment_type
   USE qs_kind_types,                   ONLY: get_qs_kind,&
                                              qs_kind_type
   USE qs_mo_types,                     ONLY: get_mo_set,&
                                              mo_set_type
   USE atomic_kind_types,               ONLY: atomic_kind_type,&
                                              get_atomic_kind
   USE particle_types,                  ONLY: particle_type
   USE cell_types,                      ONLY: cell_type
   USE virial_types,                    ONLY: virial_type
   USE qs_scf_types,                    ONLY: qs_scf_env_type
   USE scf_control_types,               ONLY: scf_control_type
   USE qs_energy_types,                 ONLY: qs_energy_type
   USE qs_rho_types,                    ONLY: qs_rho_type,&
                                              qs_rho_get
   USE pw_types,                        ONLY: pw_r3d_rs_type
   USE pw_grid_types,                   ONLY: pw_grid_type
   USE cp_fm_types,                     ONLY: cp_fm_type,&
                                              cp_fm_get_submatrix,&
                                              cp_fm_get_info
    USE cp_dbcsr_api,                    ONLY: dbcsr_p_type,&
                                               dbcsr_type
   USE mulliken,                        ONLY: mulliken_charges
   USE message_passing,                 ONLY: mp_para_env_type
    USE kpoint_types,                    ONLY: kpoint_type,&
                                               kpoint_env_p_type,&
                                               get_kpoint_info,&
                                               get_kpoint_env
    USE hirshfeld_types,                 ONLY: hirshfeld_type,&
                                               create_hirshfeld_type,&
                                               release_hirshfeld_type,&
                                               set_hirshfeld_info
    USE hirshfeld_methods,               ONLY: create_shape_function,&
                                               comp_hirshfeld_charges
    USE input_constants,                 ONLY: shape_function_gaussian,&
                                               ref_charge_atomic

    IMPLICIT NONE

   PRIVATE

   ! Public interface
   PUBLIC :: cp2k_get_stress_tensor
   PUBLIC :: cp2k_get_virial_tensor
   PUBLIC :: cp2k_get_nmo
   PUBLIC :: cp2k_get_eigenvalues
   PUBLIC :: cp2k_get_occupation_numbers
   PUBLIC :: cp2k_get_homo_lumo
   PUBLIC :: cp2k_get_mulliken_charges
   PUBLIC :: cp2k_get_dipole_moment
   PUBLIC :: cp2k_get_scf_info
   PUBLIC :: cp2k_is_qs_env
   PUBLIC :: cp2k_get_energy_components
   PUBLIC :: cp2k_get_nelectron
   PUBLIC :: cp2k_get_fermi_energy
   PUBLIC :: cp2k_get_hirshfeld_charges
   PUBLIC :: cp2k_get_total_spin
   PUBLIC :: cp2k_get_grid_info
   PUBLIC :: cp2k_get_electron_density_grid
   PUBLIC :: cp2k_get_mo_coeff_info
   PUBLIC :: cp2k_get_mo_coefficients
   PUBLIC :: cp2k_get_nkpoints
   PUBLIC :: cp2k_get_kpoint_eigenvalues

CONTAINS

! **************************************************************************************************
!> \brief Check if the force environment is a Quickstep (DFT) calculation
!> \param env_id The force environment ID
!> \return 1 if Quickstep, 0 otherwise
!> \author CP2K-RS Developers
! **************************************************************************************************
   INTEGER(C_INT) FUNCTION cp2k_is_qs_env(env_id) BIND(C, name='cp2k_is_qs_env')
      INTEGER(C_INT), VALUE, INTENT(IN)                  :: env_id

      INTEGER                                            :: ierr, in_use
      TYPE(f_env_type), POINTER                          :: f_env

      cp2k_is_qs_env = 0
      NULLIFY (f_env)

      CALL f_env_add_defaults(env_id, f_env)
      IF (.NOT. ASSOCIATED(f_env)) RETURN

      CALL force_env_get(f_env%force_env, in_use=in_use)

      IF (in_use == use_qs_force) THEN
         cp2k_is_qs_env = 1
      END IF

      IF (ASSOCIATED(f_env)) CALL f_env_rm_defaults(f_env, ierr)

   END FUNCTION cp2k_is_qs_env

! **************************************************************************************************
!> \brief Get the stress tensor from a force environment
!> \param env_id The force environment ID
!> \param stress 3x3 stress tensor in GPa (output)
!> \author CP2K-RS Developers
! **************************************************************************************************
   SUBROUTINE cp2k_get_stress_tensor(env_id, stress) BIND(C, name='cp2k_get_stress_tensor')
      INTEGER(C_INT), VALUE, INTENT(IN)                  :: env_id
      REAL(C_DOUBLE), DIMENSION(3, 3), INTENT(OUT)       :: stress

      INTEGER                                            :: ierr
      REAL(KIND=dp)                                      :: volume
      TYPE(cell_type), POINTER                           :: cell
      TYPE(f_env_type), POINTER                          :: f_env
      TYPE(qs_environment_type), POINTER                 :: qs_env
      TYPE(virial_type), POINTER                         :: virial

      stress = 0.0_C_DOUBLE
      NULLIFY (f_env, qs_env, virial, cell)

      CALL f_env_add_defaults(env_id, f_env)
      IF (.NOT. ASSOCIATED(f_env)) RETURN

      CALL force_env_get(f_env%force_env, qs_env=qs_env, cell=cell)

      IF (ASSOCIATED(qs_env)) THEN
         CALL get_qs_env(qs_env, virial=virial)
         IF (ASSOCIATED(virial) .AND. ASSOCIATED(cell)) THEN
            volume = cell%deth
            IF (ABS(volume) > 1.0E-10_dp) THEN
               ! stress = -virial / volume, convert to GPa
               ! 1 a.u. = 2.9421912e13 Pa = 29421.912 GPa
               stress(1:3, 1:3) = -virial%pv_virial(1:3, 1:3)/volume*29421.912_dp
            END IF
         END IF
      END IF

      IF (ASSOCIATED(f_env)) CALL f_env_rm_defaults(f_env, ierr)

   END SUBROUTINE cp2k_get_stress_tensor

! **************************************************************************************************
!> \brief Get the virial tensor from a force environment
!> \param env_id The force environment ID
!> \param virial_tensor 3x3 virial tensor in Hartree (output)
!> \author CP2K-RS Developers
! **************************************************************************************************
   SUBROUTINE cp2k_get_virial_tensor(env_id, virial_tensor) BIND(C, name='cp2k_get_virial_tensor')
      INTEGER(C_INT), VALUE, INTENT(IN)                  :: env_id
      REAL(C_DOUBLE), DIMENSION(3, 3), INTENT(OUT)       :: virial_tensor

      INTEGER                                            :: ierr
      TYPE(f_env_type), POINTER                          :: f_env
      TYPE(qs_environment_type), POINTER                 :: qs_env
      TYPE(virial_type), POINTER                         :: virial

      virial_tensor = 0.0_C_DOUBLE
      NULLIFY (f_env, qs_env, virial)

      CALL f_env_add_defaults(env_id, f_env)
      IF (.NOT. ASSOCIATED(f_env)) RETURN

      CALL force_env_get(f_env%force_env, qs_env=qs_env)

      IF (ASSOCIATED(qs_env)) THEN
         CALL get_qs_env(qs_env, virial=virial)
         IF (ASSOCIATED(virial)) THEN
            virial_tensor(1:3, 1:3) = virial%pv_virial(1:3, 1:3)
         END IF
      END IF

      IF (ASSOCIATED(f_env)) CALL f_env_rm_defaults(f_env, ierr)

   END SUBROUTINE cp2k_get_virial_tensor

! **************************************************************************************************
!> \brief Get the number of molecular orbitals for a given spin channel
!> \param env_id The force environment ID
!> \param spin Spin channel (1 or 2)
!> \return Number of MOs, or -1 if not available
!> \author CP2K-RS Developers
! **************************************************************************************************
   INTEGER(C_INT) FUNCTION cp2k_get_nmo(env_id, spin) BIND(C, name='cp2k_get_nmo')
      INTEGER(C_INT), VALUE, INTENT(IN)                  :: env_id, spin

      INTEGER                                            :: ierr, in_use, nmo
      TYPE(f_env_type), POINTER                          :: f_env
      TYPE(mo_set_type), DIMENSION(:), POINTER           :: mos
      TYPE(qs_environment_type), POINTER                 :: qs_env

      cp2k_get_nmo = -1
      NULLIFY (f_env, qs_env, mos)

      CALL f_env_add_defaults(env_id, f_env)
      IF (.NOT. ASSOCIATED(f_env)) RETURN

      CALL force_env_get(f_env%force_env, in_use=in_use)

      IF (in_use == use_qs_force) THEN
         qs_env => f_env%force_env%qs_env
         CALL get_qs_env(qs_env, mos=mos)

         IF (ASSOCIATED(mos)) THEN
            IF (spin >= 1 .AND. spin <= SIZE(mos)) THEN
               CALL get_mo_set(mos(spin), nmo=nmo)
               cp2k_get_nmo = nmo
            END IF
         END IF
      END IF

      IF (ASSOCIATED(f_env)) CALL f_env_rm_defaults(f_env, ierr)

   END FUNCTION cp2k_get_nmo

! **************************************************************************************************
!> \brief Get Kohn-Sham eigenvalues for a spin channel
!> \param env_id The force environment ID
!> \param spin Spin channel (1 or 2)
!> \param eigenvalues Buffer for eigenvalues in Hartree
!> \param nmax Maximum number of eigenvalues to retrieve
!> \return Number of eigenvalues retrieved, or -1 on error
!> \author CP2K-RS Developers
! **************************************************************************************************
   INTEGER(C_INT) FUNCTION cp2k_get_eigenvalues(env_id, spin, eigenvalues, nmax) &
      BIND(C, name='cp2k_get_eigenvalues')
      INTEGER(C_INT), VALUE, INTENT(IN)                  :: env_id, spin, nmax
      REAL(C_DOUBLE), DIMENSION(nmax), INTENT(OUT)       :: eigenvalues

      INTEGER                                            :: i, ierr, in_use, nmo
      REAL(KIND=dp), DIMENSION(:), POINTER               :: eigvals
      TYPE(f_env_type), POINTER                          :: f_env
      TYPE(mo_set_type), DIMENSION(:), POINTER           :: mos
      TYPE(qs_environment_type), POINTER                 :: qs_env

      cp2k_get_eigenvalues = -1
      eigenvalues = 0.0_C_DOUBLE
      NULLIFY (f_env, qs_env, mos, eigvals)

      CALL f_env_add_defaults(env_id, f_env)
      IF (.NOT. ASSOCIATED(f_env)) RETURN

      CALL force_env_get(f_env%force_env, in_use=in_use)

      IF (in_use == use_qs_force) THEN
         qs_env => f_env%force_env%qs_env
         CALL get_qs_env(qs_env, mos=mos)

         IF (ASSOCIATED(mos)) THEN
            IF (spin >= 1 .AND. spin <= SIZE(mos)) THEN
               ! First get nmo
               CALL get_mo_set(mos(spin), nmo=nmo)
               ! Then get eigenvalues pointer
               CALL get_mo_set(mos(spin), eigenvalues=eigvals)

               IF (ASSOCIATED(eigvals)) THEN
                  nmo = MIN(nmo, nmax)
                  DO i = 1, nmo
                     eigenvalues(i) = eigvals(i)
                  END DO
                  cp2k_get_eigenvalues = nmo
               END IF
            END IF
         END IF
      END IF

      IF (ASSOCIATED(f_env)) CALL f_env_rm_defaults(f_env, ierr)

   END FUNCTION cp2k_get_eigenvalues

! **************************************************************************************************
!> \brief Get occupation numbers for a spin channel
!> \param env_id The force environment ID
!> \param spin Spin channel (1 or 2)
!> \param occupations Buffer for occupation numbers
!> \param nmax Maximum number of occupations to retrieve
!> \return Number of occupations retrieved, or -1 on error
!> \author CP2K-RS Developers
! **************************************************************************************************
   INTEGER(C_INT) FUNCTION cp2k_get_occupation_numbers(env_id, spin, occupations, nmax) &
      BIND(C, name='cp2k_get_occupation_numbers')
      INTEGER(C_INT), VALUE, INTENT(IN)                  :: env_id, spin, nmax
      REAL(C_DOUBLE), DIMENSION(nmax), INTENT(OUT)       :: occupations

      INTEGER                                            :: i, ierr, in_use, nmo
      REAL(KIND=dp), DIMENSION(:), POINTER               :: occup
      TYPE(f_env_type), POINTER                          :: f_env
      TYPE(mo_set_type), DIMENSION(:), POINTER           :: mos
      TYPE(qs_environment_type), POINTER                 :: qs_env

      cp2k_get_occupation_numbers = -1
      occupations = 0.0_C_DOUBLE
      NULLIFY (f_env, qs_env, mos, occup)

      CALL f_env_add_defaults(env_id, f_env)
      IF (.NOT. ASSOCIATED(f_env)) RETURN

      CALL force_env_get(f_env%force_env, in_use=in_use)

      IF (in_use == use_qs_force) THEN
         qs_env => f_env%force_env%qs_env
         CALL get_qs_env(qs_env, mos=mos)

         IF (ASSOCIATED(mos)) THEN
            IF (spin >= 1 .AND. spin <= SIZE(mos)) THEN
               ! First get nmo
               CALL get_mo_set(mos(spin), nmo=nmo)
               ! Then get occupation numbers pointer
               CALL get_mo_set(mos(spin), occupation_numbers=occup)

               IF (ASSOCIATED(occup)) THEN
                  nmo = MIN(nmo, nmax)
                  DO i = 1, nmo
                     occupations(i) = occup(i)
                  END DO
                  cp2k_get_occupation_numbers = nmo
               END IF
            END IF
         END IF
      END IF

      IF (ASSOCIATED(f_env)) CALL f_env_rm_defaults(f_env, ierr)

   END FUNCTION cp2k_get_occupation_numbers

! **************************************************************************************************
!> \brief Get HOMO and LUMO energies and indices
!> \param env_id The force environment ID
!> \param spin Spin channel (1 or 2)
!> \param homo_energy HOMO energy in Hartree (output)
!> \param lumo_energy LUMO energy in Hartree (output)
!> \param homo_index Index of HOMO (1-based, output)
!> \param lumo_index Index of LUMO (1-based, output)
!> \return 0 on success, -1 on error
!> \author CP2K-RS Developers
! **************************************************************************************************
   INTEGER(C_INT) FUNCTION cp2k_get_homo_lumo(env_id, spin, homo_energy, lumo_energy, &
                                               homo_index, lumo_index) BIND(C, name='cp2k_get_homo_lumo')
      INTEGER(C_INT), VALUE, INTENT(IN)                  :: env_id, spin
      REAL(C_DOUBLE), INTENT(OUT)                        :: homo_energy, lumo_energy
      INTEGER(C_INT), INTENT(OUT)                        :: homo_index, lumo_index

      INTEGER                                            :: homo, ierr, in_use, lumo
      REAL(KIND=dp), DIMENSION(:), POINTER               :: eigvals
      TYPE(f_env_type), POINTER                          :: f_env
      TYPE(mo_set_type), DIMENSION(:), POINTER           :: mos
      TYPE(qs_environment_type), POINTER                 :: qs_env

      cp2k_get_homo_lumo = -1
      homo_energy = 0.0_C_DOUBLE
      lumo_energy = 0.0_C_DOUBLE
      homo_index = 0
      lumo_index = 0
      NULLIFY (f_env, qs_env, mos, eigvals)

      CALL f_env_add_defaults(env_id, f_env)
      IF (.NOT. ASSOCIATED(f_env)) RETURN

      CALL force_env_get(f_env%force_env, in_use=in_use)

      IF (in_use == use_qs_force) THEN
         qs_env => f_env%force_env%qs_env
         CALL get_qs_env(qs_env, mos=mos)

         IF (ASSOCIATED(mos)) THEN
            IF (spin >= 1 .AND. spin <= SIZE(mos)) THEN
               ! Get HOMO and LUMO indices
               CALL get_mo_set(mos(spin), homo=homo, lfomo=lumo)
               ! Get eigenvalues array
               CALL get_mo_set(mos(spin), eigenvalues=eigvals)

                IF (ASSOCIATED(eigvals)) THEN
                   IF (homo > 0 .AND. homo <= SIZE(eigvals)) THEN
                      homo_energy = eigvals(homo)
                      homo_index = homo
                   END IF

                   IF (lumo > 0 .AND. lumo <= SIZE(eigvals)) THEN
                      lumo_energy = eigvals(lumo)
                      lumo_index = lumo
                   END IF

                   IF (homo_index > 0 .AND. lumo_index > 0) THEN
                      cp2k_get_homo_lumo = 0
                   END IF
                END IF
            END IF
         END IF
      END IF

      IF (ASSOCIATED(f_env)) CALL f_env_rm_defaults(f_env, ierr)

   END FUNCTION cp2k_get_homo_lumo

! **************************************************************************************************
!> \brief Get Mulliken atomic charges (electronic population per atom)
!> \param env_id The force environment ID
!> \param charges Buffer for atomic charges (electronic population, not net charge)
!> \param natom Number of atoms
!> \return 0 on success, -1 on error
!> \note Returns the Mulliken population (sum over all spin channels). Net charge = zeff - charges.
!>       Requires rho_ao (density in AO basis) and matrix_s (overlap matrix) to be available,
!>       i.e., a completed AO-based DFT calculation.
!> \author CP2K-RS Developers
! **************************************************************************************************
    INTEGER(C_INT) FUNCTION cp2k_get_mulliken_charges(env_id, charges, natom) &
       BIND(C, name='cp2k_get_mulliken_charges')
       INTEGER(C_INT), VALUE, INTENT(IN)                  :: env_id, natom
       REAL(C_DOUBLE), DIMENSION(natom), INTENT(OUT)      :: charges

       INTEGER                                            :: i, ierr, in_use, nspin, nat
        REAL(KIND=dp), DIMENSION(:, :), POINTER            :: pop2
        TYPE(dbcsr_type), POINTER                          :: s_matrix
        TYPE(dbcsr_p_type), DIMENSION(:), POINTER          :: matrix_s, rho_ao
       TYPE(f_env_type), POINTER                          :: f_env
       TYPE(mp_para_env_type), POINTER                    :: para_env
       TYPE(particle_type), DIMENSION(:), POINTER         :: particle_set
       TYPE(qs_environment_type), POINTER                 :: qs_env
       TYPE(qs_rho_type), POINTER                         :: rho

       cp2k_get_mulliken_charges = -1
       charges = 0.0_C_DOUBLE
       NULLIFY (f_env, qs_env, particle_set, rho, rho_ao, matrix_s, para_env, s_matrix, pop2)

       CALL f_env_add_defaults(env_id, f_env)
      IF (.NOT. ASSOCIATED(f_env)) RETURN

       CALL force_env_get(f_env%force_env, in_use=in_use)

       IF (in_use == use_qs_force) THEN
          qs_env => f_env%force_env%qs_env
          CALL get_qs_env(qs_env, particle_set=particle_set, matrix_s=matrix_s, &
                          para_env=para_env, rho=rho)

           IF (ASSOCIATED(particle_set) .AND. ASSOCIATED(matrix_s) .AND. &
               ASSOCIATED(para_env) .AND. ASSOCIATED(rho) .AND. SIZE(matrix_s) >= 1) THEN
              nat = SIZE(particle_set)
              s_matrix => matrix_s(1)%matrix
              CALL qs_rho_get(rho, rho_ao=rho_ao)

               IF (ASSOCIATED(rho_ao) .AND. ASSOCIATED(s_matrix)) THEN
                  nspin = SIZE(rho_ao)
                  ALLOCATE (pop2(nat, nspin))
                  pop2 = 0.0_dp

                  CALL mulliken_charges(rho_ao, s_matrix, para_env, pop2)

                  DO i = 1, nat
                     charges(i) = SUM(pop2(i, :))
                  END DO

                  DEALLOCATE (pop2)
                  NULLIFY (pop2)
                  cp2k_get_mulliken_charges = 0
              END IF
          END IF
       END IF

       IF (ASSOCIATED(f_env)) CALL f_env_rm_defaults(f_env, ierr)

    END FUNCTION cp2k_get_mulliken_charges

! **************************************************************************************************
!> \brief Get the dipole moment vector from Mulliken net charges
!> \param env_id The force environment ID
!> \param dipole 3-component dipole moment vector in Debye (output)
!> \return 0 on success, -1 on error
!> \note Computes dipole as sum(q_i^net * R_i) with q_i^net = zeff_i - Mulliken_population_i.
!>       Requires rho_ao and matrix_s, same as cp2k_get_mulliken_charges.
!>       Dipole magnitude in Debye (1 a.u. = 2.541746 D).
!> \author CP2K-RS Developers
! **************************************************************************************************
    INTEGER(C_INT) FUNCTION cp2k_get_dipole_moment(env_id, dipole) &
       BIND(C, name='cp2k_get_dipole_moment')
       INTEGER(C_INT), VALUE, INTENT(IN)                  :: env_id
       REAL(C_DOUBLE), DIMENSION(3), INTENT(OUT)          :: dipole

        INTEGER                                            :: iatom, ikind, ierr, in_use, nat, nkind, nspin
       REAL(KIND=dp), PARAMETER                           :: au2debye = 2.541746_dp
       REAL(KIND=dp)                                      :: zeff
        REAL(KIND=dp), DIMENSION(:, :), POINTER            :: pop2
        TYPE(atomic_kind_type), POINTER                    :: atomic_kind
       TYPE(dbcsr_type), POINTER                          :: s_matrix
       TYPE(dbcsr_p_type), DIMENSION(:), POINTER          :: matrix_s, rho_ao
       TYPE(f_env_type), POINTER                          :: f_env
       TYPE(mp_para_env_type), POINTER                    :: para_env
       TYPE(particle_type), DIMENSION(:), POINTER         :: particle_set
       TYPE(qs_kind_type), DIMENSION(:), POINTER          :: qs_kind_set
       TYPE(qs_environment_type), POINTER                 :: qs_env
       TYPE(qs_rho_type), POINTER                         :: rho

       cp2k_get_dipole_moment = -1
       dipole = 0.0_C_DOUBLE
       NULLIFY (f_env, qs_env, particle_set, qs_kind_set, rho, rho_ao, matrix_s, para_env, s_matrix, pop2)

       CALL f_env_add_defaults(env_id, f_env)
      IF (.NOT. ASSOCIATED(f_env)) RETURN

       CALL force_env_get(f_env%force_env, in_use=in_use)

       IF (in_use == use_qs_force) THEN
          qs_env => f_env%force_env%qs_env
          CALL get_qs_env(qs_env, particle_set=particle_set, qs_kind_set=qs_kind_set, &
                          matrix_s=matrix_s, para_env=para_env, rho=rho)

           IF (ASSOCIATED(particle_set) .AND. ASSOCIATED(qs_kind_set) .AND. &
               ASSOCIATED(matrix_s) .AND. ASSOCIATED(para_env) .AND. ASSOCIATED(rho) .AND. &
               SIZE(matrix_s) >= 1) THEN
              nat = SIZE(particle_set)
              s_matrix => matrix_s(1)%matrix
              CALL qs_rho_get(rho, rho_ao=rho_ao)

               IF (ASSOCIATED(rho_ao) .AND. ASSOCIATED(s_matrix)) THEN
                  nspin = SIZE(rho_ao)
                  ALLOCATE (pop2(nat, nspin))
                  pop2 = 0.0_dp

                  CALL mulliken_charges(rho_ao, s_matrix, para_env, pop2)

                  ! Compute dipole from Mulliken net charges at atomic positions:
                  ! mu = sum_i (q_i^net * R_i), with q_i^net = zeff_i - population_i
                  nkind = SIZE(qs_kind_set)
                  DO iatom = 1, nat
                     atomic_kind => particle_set(iatom)%atomic_kind
                     CALL get_atomic_kind(atomic_kind, kind_number=ikind)
                     IF (ikind >= 1 .AND. ikind <= nkind) THEN
                        CALL get_qs_kind(qs_kind_set(ikind), zeff=zeff)
                        dipole(1:3) = dipole(1:3) + (zeff - SUM(pop2(iatom, :))) * &
                                      particle_set(iatom)%r(1:3)
                     END IF
                  END DO

                  DEALLOCATE (pop2)
                  NULLIFY (pop2)
                  dipole(1:3) = dipole(1:3) * au2debye
                  cp2k_get_dipole_moment = 0
              END IF
          END IF
       END IF

       IF (ASSOCIATED(f_env)) CALL f_env_rm_defaults(f_env, ierr)

    END FUNCTION cp2k_get_dipole_moment

! **************************************************************************************************
!> \brief Get SCF convergence information
!> \param env_id The force environment ID
!> \param niter Number of SCF iterations (output)
!> \param converged 1 if converged, 0 otherwise (output)
!> \param energy_change Last energy change in Hartree (output)
!> \return 0 on success, -1 on error
!> \note Extracts iteration count and energy delta from scf_env.
!>       Convergence is determined by comparing energy change to threshold (1E-6 Ha).
!> \author CP2K-RS Developers
! **************************************************************************************************
    INTEGER(C_INT) FUNCTION cp2k_get_scf_info(env_id, niter, converged, energy_change) &
       BIND(C, name='cp2k_get_scf_info')
       INTEGER(C_INT), VALUE, INTENT(IN)                  :: env_id
       INTEGER(C_INT), INTENT(OUT)                        :: niter, converged
       REAL(C_DOUBLE), INTENT(OUT)                        :: energy_change

       INTEGER                                            :: ierr, in_use
       REAL(KIND=dp)                                      :: eps_scf
       TYPE(f_env_type), POINTER                          :: f_env
       TYPE(qs_environment_type), POINTER                 :: qs_env
       TYPE(qs_scf_env_type), POINTER                     :: scf_env
       TYPE(scf_control_type), POINTER                    :: scf_control

       cp2k_get_scf_info = -1
       niter = 0
       converged = 0
       energy_change = 0.0_C_DOUBLE
       NULLIFY (f_env, qs_env, scf_env, scf_control)

       CALL f_env_add_defaults(env_id, f_env)
      IF (.NOT. ASSOCIATED(f_env)) RETURN

       CALL force_env_get(f_env%force_env, in_use=in_use)

       IF (in_use == use_qs_force) THEN
          qs_env => f_env%force_env%qs_env
          CALL get_qs_env(qs_env, scf_env=scf_env)

          IF (ASSOCIATED(scf_env)) THEN
             ! Get SCF iteration count
             niter = scf_env%iter_count

             ! Get last energy change (delta)
             energy_change = scf_env%iter_delta

             ! Get SCF convergence threshold from scf_control
             eps_scf = 1.0E-6_dp
             CALL get_qs_env(qs_env, scf_control=scf_control)
             IF (ASSOCIATED(scf_control)) THEN
                IF (scf_control%eps_scf > 0.0_dp) eps_scf = scf_control%eps_scf
             END IF

             ! Check convergence using the actual eps_scf threshold
             IF (ABS(scf_env%iter_delta) < eps_scf) THEN
                converged = 1
             ELSE
                converged = 0
             END IF

             cp2k_get_scf_info = 0
          END IF
       END IF

       IF (ASSOCIATED(f_env)) CALL f_env_rm_defaults(f_env, ierr)

    END FUNCTION cp2k_get_scf_info

! **************************************************************************************************
!> \brief Get energy components (kinetic, Hartree, XC, etc.)
!> \param env_id The force environment ID
!> \param e_kinetic Kinetic energy in Hartree (output)
!> \param e_hartree Hartree (electron-electron Coulomb) energy in Hartree (output)
!> \param e_xc Exchange-correlation energy in Hartree (output)
!> \param e_core Core Hamiltonian energy in Hartree (output)
!> \param e_total Total energy in Hartree (output)
!> \return 0 on success, -1 on error
!> \author CP2K-RS Developers
! **************************************************************************************************
   INTEGER(C_INT) FUNCTION cp2k_get_energy_components(env_id, e_kinetic, e_hartree, &
                                                       e_xc, e_core, e_total) &
      BIND(C, name='cp2k_get_energy_components')
      INTEGER(C_INT), VALUE, INTENT(IN)                  :: env_id
      REAL(C_DOUBLE), INTENT(OUT)                        :: e_kinetic, e_hartree, e_xc, e_core, e_total

       INTEGER                                            :: ierr, in_use
       TYPE(f_env_type), POINTER                          :: f_env
       TYPE(qs_environment_type), POINTER                 :: qs_env
       TYPE(qs_energy_type), POINTER                      :: energy

       cp2k_get_energy_components = -1
       e_kinetic = 0.0_C_DOUBLE
       e_hartree = 0.0_C_DOUBLE
       e_xc = 0.0_C_DOUBLE
       e_core = 0.0_C_DOUBLE
       e_total = 0.0_C_DOUBLE
       NULLIFY (f_env, qs_env, energy)

       CALL f_env_add_defaults(env_id, f_env)
      IF (.NOT. ASSOCIATED(f_env)) RETURN

       CALL force_env_get(f_env%force_env, in_use=in_use)

      IF (in_use == use_qs_force) THEN
         qs_env => f_env%force_env%qs_env
         IF (ASSOCIATED(qs_env)) THEN
            CALL get_qs_env(qs_env, energy=energy)
            IF (ASSOCIATED(energy)) THEN
               e_kinetic = energy%kinetic
               e_hartree = energy%hartree
               e_xc      = energy%exc
               e_core    = energy%core
               e_total   = energy%total
               cp2k_get_energy_components = 0
            END IF
         END IF
      END IF

      IF (ASSOCIATED(f_env)) CALL f_env_rm_defaults(f_env, ierr)

   END FUNCTION cp2k_get_energy_components

! **************************************************************************************************
!> \brief Get the number of electrons
!> \param env_id The force environment ID
!> \param nelectron Number of electrons (output)
!> \return 0 on success, -1 on error
!> \author CP2K-RS Developers
! **************************************************************************************************
   INTEGER(C_INT) FUNCTION cp2k_get_nelectron(env_id, nelectron) &
      BIND(C, name='cp2k_get_nelectron')
      INTEGER(C_INT), VALUE, INTENT(IN)                  :: env_id
      INTEGER(C_INT), INTENT(OUT)                        :: nelectron

      INTEGER                                            :: ierr, in_use, nelec
      TYPE(f_env_type), POINTER                          :: f_env
      TYPE(qs_environment_type), POINTER                 :: qs_env
      TYPE(qs_scf_env_type), POINTER                     :: scf_env

      cp2k_get_nelectron = -1
      nelectron = 0
      NULLIFY (f_env, qs_env, scf_env)

      CALL f_env_add_defaults(env_id, f_env)
      IF (.NOT. ASSOCIATED(f_env)) RETURN

      CALL force_env_get(f_env%force_env, in_use=in_use)

      IF (in_use == use_qs_force) THEN
         qs_env => f_env%force_env%qs_env
         CALL get_qs_env(qs_env, scf_env=scf_env)

         IF (ASSOCIATED(scf_env)) THEN
            nelec = scf_env%nelectron
            nelectron = nelec
            cp2k_get_nelectron = 0
         END IF
      END IF

      IF (ASSOCIATED(f_env)) CALL f_env_rm_defaults(f_env, ierr)

   END FUNCTION cp2k_get_nelectron

! **************************************************************************************************
!> \brief Get the Fermi energy (for smeared occupations)
!> \param env_id The force environment ID
!> \param e_fermi Fermi energy in Hartree (output)
!> \return 0 on success, -1 on error or if not available
!> \author CP2K-RS Developers
! **************************************************************************************************
   INTEGER(C_INT) FUNCTION cp2k_get_fermi_energy(env_id, e_fermi) &
      BIND(C, name='cp2k_get_fermi_energy')
      INTEGER(C_INT), VALUE, INTENT(IN)                  :: env_id
      REAL(C_DOUBLE), INTENT(OUT)                        :: e_fermi

      INTEGER                                            :: ierr, in_use
      TYPE(f_env_type), POINTER                          :: f_env
      TYPE(qs_environment_type), POINTER                 :: qs_env
      TYPE(qs_energy_type), POINTER                      :: energy

      cp2k_get_fermi_energy = -1
      e_fermi = 0.0_C_DOUBLE
      NULLIFY (f_env, qs_env, energy)

      CALL f_env_add_defaults(env_id, f_env)
      IF (.NOT. ASSOCIATED(f_env)) RETURN

      CALL force_env_get(f_env%force_env, in_use=in_use)

      IF (in_use == use_qs_force) THEN
         qs_env => f_env%force_env%qs_env

         ! CP2K prints "Fermi energy:" from the Quickstep energy object as energy%efermi.
         ! Returning this value is robust across spin/k-point setups and matches CP2K output.
         NULLIFY (energy)
         CALL get_qs_env(qs_env, energy=energy)

         IF (ASSOCIATED(energy)) THEN
            e_fermi = energy%efermi
            cp2k_get_fermi_energy = 0
         END IF
      END IF

      IF (ASSOCIATED(f_env)) CALL f_env_rm_defaults(f_env, ierr)

   END FUNCTION cp2k_get_fermi_energy

! **************************************************************************************************
!> \brief Get Hirshfeld atomic charges
!> \param env_id The force environment ID
!> \param charges Buffer for atomic charges (output)
!> \param natom Number of atoms
!> \return 0 on success, -1 on error
!> \note Requires a GPW/GAPW calculation with real-space density available.
!>       Uses atomic reference charges with default Gaussian shape functions.
!> \author CP2K-RS Developers
! **************************************************************************************************
    INTEGER(C_INT) FUNCTION cp2k_get_hirshfeld_charges(env_id, charges, natom) &
       BIND(C, name='cp2k_get_hirshfeld_charges')
       INTEGER(C_INT), VALUE, INTENT(IN)                  :: env_id, natom
       REAL(C_DOUBLE), DIMENSION(natom), INTENT(OUT)      :: charges

        INTEGER                                            :: iatom, ikind, ierr, in_use, nspin, nat, nkind
       LOGICAL                                            :: rho_r_valid
       REAL(KIND=dp)                                      :: zeff
        REAL(KIND=dp), DIMENSION(:, :), ALLOCATABLE        :: pop2
        TYPE(atomic_kind_type), DIMENSION(:), POINTER      :: atomic_kind_set
       TYPE(atomic_kind_type), POINTER                    :: atomic_kind
       TYPE(f_env_type), POINTER                          :: f_env
       TYPE(hirshfeld_type), POINTER                      :: hirshfeld_env
       TYPE(particle_type), DIMENSION(:), POINTER         :: particle_set
       TYPE(pw_r3d_rs_type), DIMENSION(:), POINTER        :: rho_r
       TYPE(qs_kind_type), DIMENSION(:), POINTER          :: qs_kind_set
       TYPE(qs_environment_type), POINTER                 :: qs_env
       TYPE(qs_rho_type), POINTER                         :: rho

       cp2k_get_hirshfeld_charges = -1
       charges = 0.0_C_DOUBLE
       NULLIFY (f_env, qs_env, particle_set, qs_kind_set, atomic_kind_set, &
                atomic_kind, rho, rho_r, hirshfeld_env)

       CALL f_env_add_defaults(env_id, f_env)
      IF (.NOT. ASSOCIATED(f_env)) RETURN

       CALL force_env_get(f_env%force_env, in_use=in_use)

       IF (in_use == use_qs_force) THEN
          qs_env => f_env%force_env%qs_env
          CALL get_qs_env(qs_env, particle_set=particle_set, qs_kind_set=qs_kind_set, &
                          atomic_kind_set=atomic_kind_set, rho=rho)

           IF (.NOT. (ASSOCIATED(particle_set) .AND. ASSOCIATED(qs_kind_set) .AND. &
               ASSOCIATED(atomic_kind_set) .AND. ASSOCIATED(rho))) &
              RETURN

            nat = SIZE(particle_set)
            IF (nat /= natom) RETURN

           CALL qs_rho_get(rho, rho_r=rho_r, rho_r_valid=rho_r_valid)
           IF (.NOT. rho_r_valid .OR. .NOT. ASSOCIATED(rho_r)) RETURN

           nspin = SIZE(rho_r)
           NULLIFY (hirshfeld_env)
           CALL create_hirshfeld_type(hirshfeld_env)

           CALL set_hirshfeld_info(hirshfeld_env, shape_function_type=shape_function_gaussian, &
                                   iterative=.FALSE., ref_charge=ref_charge_atomic)
           CALL create_shape_function(hirshfeld_env, qs_kind_set, atomic_kind_set)

            ! Set reference charges to atomic nuclear charges (zeff).
            ! NOTE: Direct allocation of hirshfeld_env%charges is a fragile
            ! dependency on CP2K's internal hirshfeld_type layout. This must
            ! be re-verified after every CP2K version upgrade.
            ALLOCATE (hirshfeld_env%charges(natom))
           nkind = SIZE(qs_kind_set)
           DO iatom = 1, natom
              atomic_kind => particle_set(iatom)%atomic_kind
              CALL get_atomic_kind(atomic_kind, kind_number=ikind)
              IF (ikind >= 1 .AND. ikind <= nkind) THEN
                 CALL get_qs_kind(qs_kind_set(ikind), zeff=zeff)
                 hirshfeld_env%charges(iatom) = zeff
              ELSE
                 hirshfeld_env%charges(iatom) = 0.0_dp
              END IF
           END DO

            ALLOCATE (pop2(natom, nspin))
           pop2 = 0.0_dp

           CALL comp_hirshfeld_charges(qs_env, hirshfeld_env, pop2)

           DO iatom = 1, natom
              charges(iatom) = hirshfeld_env%charges(iatom) - SUM(pop2(iatom, :))
           END DO

            DEALLOCATE (pop2)
            CALL release_hirshfeld_type(hirshfeld_env)
            NULLIFY (hirshfeld_env)

            cp2k_get_hirshfeld_charges = 0
       END IF

       IF (ASSOCIATED(f_env)) CALL f_env_rm_defaults(f_env, ierr)

   END FUNCTION cp2k_get_hirshfeld_charges

! **************************************************************************************************
!> \brief Get total spin (for spin-polarized calculations)
!> \param env_id The force environment ID
!> \param total_spin Total spin (N_alpha - N_beta) (output)
!> \return 0 on success, -1 on error
!> \author CP2K-RS Developers
! **************************************************************************************************
   INTEGER(C_INT) FUNCTION cp2k_get_total_spin(env_id, total_spin) &
      BIND(C, name='cp2k_get_total_spin')
      INTEGER(C_INT), VALUE, INTENT(IN)                  :: env_id
      REAL(C_DOUBLE), INTENT(OUT)                        :: total_spin

      INTEGER                                            :: i, ierr, in_use, nmo
      REAL(KIND=dp)                                      :: occ_sum, spin_val
      REAL(KIND=dp), DIMENSION(:), POINTER               :: occup
      TYPE(f_env_type), POINTER                          :: f_env
      TYPE(mo_set_type), DIMENSION(:), POINTER           :: mos
      TYPE(qs_environment_type), POINTER                 :: qs_env

      cp2k_get_total_spin = -1
      total_spin = 0.0_C_DOUBLE
      NULLIFY (f_env, qs_env, mos, occup)

      CALL f_env_add_defaults(env_id, f_env)
      IF (.NOT. ASSOCIATED(f_env)) RETURN

      CALL force_env_get(f_env%force_env, in_use=in_use)

      IF (in_use == use_qs_force) THEN
         qs_env => f_env%force_env%qs_env
         CALL get_qs_env(qs_env, mos=mos)

         IF (ASSOCIATED(mos)) THEN
            spin_val = 0.0_dp

            ! For spin-polarized: SIZE(mos) == 2
            ! For spin-unpolarized: SIZE(mos) == 1
            IF (SIZE(mos) == 2) THEN
               ! Alpha electrons
               CALL get_mo_set(mos(1), nmo=nmo)
               CALL get_mo_set(mos(1), occupation_numbers=occup)
               IF (ASSOCIATED(occup)) THEN
                  occ_sum = 0.0_dp
                  DO i = 1, nmo
                     occ_sum = occ_sum + occup(i)
                  END DO
                  spin_val = spin_val + occ_sum
               END IF

               ! Beta electrons (subtract)
               CALL get_mo_set(mos(2), nmo=nmo)
               CALL get_mo_set(mos(2), occupation_numbers=occup)
               IF (ASSOCIATED(occup)) THEN
                  occ_sum = 0.0_dp
                  DO i = 1, nmo
                     occ_sum = occ_sum + occup(i)
                  END DO
                  spin_val = spin_val - occ_sum
               END IF
            END IF

            total_spin = spin_val
            cp2k_get_total_spin = 0
         END IF
      END IF

      IF (ASSOCIATED(f_env)) CALL f_env_rm_defaults(f_env, ierr)

   END FUNCTION cp2k_get_total_spin

! **************************************************************************************************
!> \brief Get grid metadata for the electron density of a spin channel
!> \param env_id The force environment ID
!> \param spin Spin channel (1=alpha/total, 2=beta)
!> \param npts Grid points per dimension (output, 3 elements)
!> \param origin Grid origin in Bohr (output, 3 elements)
!> \param dh 3x3 cell increment matrix in Bohr, row-major flat (output, 9 elements)
!> \return 0 on success, -1 on error
!> \author CP2K-RS Developers
! **************************************************************************************************
   INTEGER(C_INT) FUNCTION cp2k_get_grid_info(env_id, spin, npts, origin, dh) &
      BIND(C, name='cp2k_get_grid_info')
      INTEGER(C_INT), VALUE, INTENT(IN)                  :: env_id, spin
      INTEGER(C_INT), DIMENSION(3), INTENT(OUT)          :: npts
      REAL(C_DOUBLE), DIMENSION(3), INTENT(OUT)          :: origin
      REAL(C_DOUBLE), DIMENSION(9), INTENT(OUT)          :: dh

      INTEGER                                            :: i, ierr, in_use, j
      LOGICAL                                            :: rho_r_valid
      TYPE(f_env_type), POINTER                          :: f_env
      TYPE(pw_r3d_rs_type), DIMENSION(:), POINTER        :: rho_r
      TYPE(qs_environment_type), POINTER                 :: qs_env
      TYPE(qs_rho_type), POINTER                         :: rho

      cp2k_get_grid_info = -1
      npts = 0
      origin = 0.0_C_DOUBLE
      dh = 0.0_C_DOUBLE
      NULLIFY (f_env, qs_env, rho, rho_r)

      CALL f_env_add_defaults(env_id, f_env)
      IF (.NOT. ASSOCIATED(f_env)) RETURN

      CALL force_env_get(f_env%force_env, in_use=in_use)

      IF (in_use == use_qs_force) THEN
         qs_env => f_env%force_env%qs_env
         CALL get_qs_env(qs_env, rho=rho)

         IF (ASSOCIATED(rho)) THEN
            CALL qs_rho_get(rho, rho_r=rho_r, rho_r_valid=rho_r_valid)

            IF (rho_r_valid .AND. ASSOCIATED(rho_r)) THEN
               IF (spin >= 1 .AND. spin <= SIZE(rho_r)) THEN
                  npts(1:3) = rho_r(spin)%pw_grid%npts(1:3)

                  ! Origin is at the lower-left corner of the grid
                  ! Compute from bounds and dh matrix
                  DO i = 1, 3
                     origin(i) = 0.0_dp
                     DO j = 1, 3
                        origin(i) = origin(i) + rho_r(spin)%pw_grid%dh(i, j) * &
                                    rho_r(spin)%pw_grid%bounds(1, j)
                     END DO
                  END DO

                  ! dh matrix: row-major flat (C convention)
                  ! dh_flat[(i-1)*3 + j] = dh(i, j) in Fortran
                  DO i = 1, 3
                     DO j = 1, 3
                        dh((i - 1)*3 + j) = rho_r(spin)%pw_grid%dh(i, j)
                     END DO
                  END DO

                  cp2k_get_grid_info = 0
               END IF
            END IF
         END IF
      END IF

      IF (ASSOCIATED(f_env)) CALL f_env_rm_defaults(f_env, ierr)

   END FUNCTION cp2k_get_grid_info

! **************************************************************************************************
!> \brief Get the full electron density on the realspace grid
!> \param env_id The force environment ID
!> \param spin Spin channel (1=alpha/total, 2=beta)
!> \param density Pre-allocated output buffer (Fortran column-major order)
!> \param nmax Buffer capacity (number of doubles)
!> \return Number of grid points written, or -1 on error
!> \note Data is in Fortran column-major order: density[(i1-1) + n1*((i2-1) + n2*(i3-1)) + 1]
!>       where n1=npts(1), n2=npts(2). For MPI-parallel runs, the distributed grid is
!>       gathered to all ranks via MPI_ALLREDUCE.
!> \author CP2K-RS Developers
! **************************************************************************************************
   INTEGER(C_INT) FUNCTION cp2k_get_electron_density_grid(env_id, spin, density, nmax) &
      BIND(C, name='cp2k_get_electron_density_grid')
      INTEGER(C_INT), VALUE, INTENT(IN)                  :: env_id, spin, nmax
      REAL(C_DOUBLE), DIMENSION(nmax), INTENT(OUT)       :: density

      INTEGER                                            :: GL1, GL2, GL3, GU1, GU2, GU3, &
                                                            i1, i2, i3, idx, ierr, in_use, &
                                                            LL1, LL2, LL3, LU1, LU2, LU3, &
                                                            n1, n2, n3, n_total
      LOGICAL                                            :: rho_r_valid
      REAL(KIND=dp), DIMENSION(:), ALLOCATABLE           :: temp
      TYPE(f_env_type), POINTER                          :: f_env
      TYPE(pw_r3d_rs_type), DIMENSION(:), POINTER        :: rho_r
      TYPE(qs_environment_type), POINTER                 :: qs_env
      TYPE(qs_rho_type), POINTER                         :: rho

      cp2k_get_electron_density_grid = -1
      density = 0.0_C_DOUBLE
      NULLIFY (f_env, qs_env, rho, rho_r)

      CALL f_env_add_defaults(env_id, f_env)
      IF (.NOT. ASSOCIATED(f_env)) RETURN

      CALL force_env_get(f_env%force_env, in_use=in_use)

      IF (in_use == use_qs_force) THEN
         qs_env => f_env%force_env%qs_env
         CALL get_qs_env(qs_env, rho=rho)

         IF (ASSOCIATED(rho)) THEN
            CALL qs_rho_get(rho, rho_r=rho_r, rho_r_valid=rho_r_valid)

            IF (rho_r_valid .AND. ASSOCIATED(rho_r)) THEN
               IF (spin >= 1 .AND. spin <= SIZE(rho_r)) THEN
                  ! Global grid dimensions
                  n1 = rho_r(spin)%pw_grid%npts(1)
                  n2 = rho_r(spin)%pw_grid%npts(2)
                  n3 = rho_r(spin)%pw_grid%npts(3)
                  n_total = n1*n2*n3

                  IF (nmax >= n_total) THEN
                     ! Global bounds
                     GL1 = rho_r(spin)%pw_grid%bounds(1, 1)
                     GL2 = rho_r(spin)%pw_grid%bounds(1, 2)
                     GL3 = rho_r(spin)%pw_grid%bounds(1, 3)
                     GU1 = rho_r(spin)%pw_grid%bounds(2, 1)
                     GU2 = rho_r(spin)%pw_grid%bounds(2, 2)
                     GU3 = rho_r(spin)%pw_grid%bounds(2, 3)

                     ! Local bounds on this MPI rank
                     LL1 = rho_r(spin)%pw_grid%bounds_local(1, 1)
                     LL2 = rho_r(spin)%pw_grid%bounds_local(1, 2)
                     LL3 = rho_r(spin)%pw_grid%bounds_local(1, 3)
                     LU1 = rho_r(spin)%pw_grid%bounds_local(2, 1)
                     LU2 = rho_r(spin)%pw_grid%bounds_local(2, 2)
                     LU3 = rho_r(spin)%pw_grid%bounds_local(2, 3)

                     ! Allocate temp buffer, zero it, copy local portion
                     ALLOCATE (temp(n_total))
                     temp = 0.0_dp

                     DO i3 = LL3, LU3
                        DO i2 = LL2, LU2
                           DO i1 = LL1, LU1
                              ! Fortran column-major: i1 varies fastest
                              idx = (i1 - GL1) + n1*((i2 - GL2) + n2*(i3 - GL3)) + 1
                              temp(idx) = rho_r(spin)%array(i1, i2, i3)
                           END DO
                        END DO
                     END DO

                     ! MPI gather: sum across all ranks (non-overlapping => sum = gather)
                     CALL rho_r(spin)%pw_grid%para%group%sum(temp)

                     ! Copy to output
                     density(1:n_total) = temp(1:n_total)

                     DEALLOCATE (temp)
                     cp2k_get_electron_density_grid = n_total
                  END IF
               END IF
            END IF
         END IF
      END IF

      IF (ASSOCIATED(f_env)) CALL f_env_rm_defaults(f_env, ierr)

   END FUNCTION cp2k_get_electron_density_grid

! **************************************************************************************************
!> \brief Get dimensions of the MO coefficient matrix
!> \param env_id The force environment ID
!> \param spin Spin channel (1=alpha/total, 2=beta)
!> \param nao Number of atomic orbitals / basis functions (output)
!> \param nmo Number of molecular orbitals (output)
!> \return 0 on success, -1 on error
!> \author CP2K-RS Developers
! **************************************************************************************************
   INTEGER(C_INT) FUNCTION cp2k_get_mo_coeff_info(env_id, spin, nao, nmo) &
      BIND(C, name='cp2k_get_mo_coeff_info')
      INTEGER(C_INT), VALUE, INTENT(IN)                  :: env_id, spin
      INTEGER(C_INT), INTENT(OUT)                        :: nao, nmo

      INTEGER                                            :: ierr, in_use, nao_f, nmo_f
      TYPE(cp_fm_type), POINTER                          :: mo_coeff
      TYPE(f_env_type), POINTER                          :: f_env
      TYPE(mo_set_type), DIMENSION(:), POINTER           :: mos
      TYPE(qs_environment_type), POINTER                 :: qs_env

      cp2k_get_mo_coeff_info = -1
      nao = 0
      nmo = 0
      NULLIFY (f_env, qs_env, mos, mo_coeff)

      CALL f_env_add_defaults(env_id, f_env)
      IF (.NOT. ASSOCIATED(f_env)) RETURN

      CALL force_env_get(f_env%force_env, in_use=in_use)

      IF (in_use == use_qs_force) THEN
         qs_env => f_env%force_env%qs_env
         CALL get_qs_env(qs_env, mos=mos)

         IF (ASSOCIATED(mos)) THEN
            IF (spin >= 1 .AND. spin <= SIZE(mos)) THEN
               CALL get_mo_set(mos(spin), mo_coeff=mo_coeff)
               IF (ASSOCIATED(mo_coeff)) THEN
                  CALL cp_fm_get_info(mo_coeff, nrow_global=nao_f, ncol_global=nmo_f)
                  nao = nao_f
                  nmo = nmo_f
                  cp2k_get_mo_coeff_info = 0
               END IF
            END IF
         END IF
      END IF

      IF (ASSOCIATED(f_env)) CALL f_env_rm_defaults(f_env, ierr)

   END FUNCTION cp2k_get_mo_coeff_info

! **************************************************************************************************
!> \brief Get the full MO coefficient matrix (gathered from ScaLAPACK distribution)
!> \param env_id The force environment ID
!> \param spin Spin channel (1=alpha/total, 2=beta)
!> \param coeffs Pre-allocated output buffer for nao*nmo doubles (column-major)
!> \param nmax Buffer capacity (number of doubles)
!> \return Number of elements written (nao*nmo), or -1 on error
!> \note Data is in Fortran column-major order: coeffs[(iao-1) + nao*(imo-1) + 1]
!>       Column j contains the j-th MO expressed in the AO basis.
!>       Uses cp_fm_get_submatrix which gathers the ScaLAPACK-distributed matrix
!>       to all MPI ranks via MPI_ALLREDUCE.
!> \author CP2K-RS Developers
! **************************************************************************************************
   INTEGER(C_INT) FUNCTION cp2k_get_mo_coefficients(env_id, spin, coeffs, nmax) &
      BIND(C, name='cp2k_get_mo_coefficients')
      INTEGER(C_INT), VALUE, INTENT(IN)                  :: env_id, spin, nmax
      REAL(C_DOUBLE), DIMENSION(nmax), INTENT(OUT)       :: coeffs

      INTEGER                                            :: i, ierr, in_use, j, nao, nmo, n_total
      REAL(KIND=dp), DIMENSION(:, :), ALLOCATABLE        :: local_matrix
      TYPE(cp_fm_type), POINTER                          :: mo_coeff
      TYPE(f_env_type), POINTER                          :: f_env
      TYPE(mo_set_type), DIMENSION(:), POINTER           :: mos
      TYPE(qs_environment_type), POINTER                 :: qs_env

      cp2k_get_mo_coefficients = -1
      coeffs = 0.0_C_DOUBLE
      NULLIFY (f_env, qs_env, mos, mo_coeff)

      CALL f_env_add_defaults(env_id, f_env)
      IF (.NOT. ASSOCIATED(f_env)) RETURN

      CALL force_env_get(f_env%force_env, in_use=in_use)

      IF (in_use == use_qs_force) THEN
         qs_env => f_env%force_env%qs_env
         CALL get_qs_env(qs_env, mos=mos)

         IF (ASSOCIATED(mos)) THEN
            IF (spin >= 1 .AND. spin <= SIZE(mos)) THEN
               CALL get_mo_set(mos(spin), mo_coeff=mo_coeff)
               IF (ASSOCIATED(mo_coeff)) THEN
                  CALL cp_fm_get_info(mo_coeff, nrow_global=nao, ncol_global=nmo)
                  n_total = nao*nmo

                  IF (nmax >= n_total) THEN
                     ! Gather the distributed matrix to all ranks
                     ALLOCATE (local_matrix(nao, nmo))
                     local_matrix = 0.0_dp
                     CALL cp_fm_get_submatrix(mo_coeff, local_matrix, 1, 1, nao, nmo)

                     ! Copy to flat output buffer in column-major order
                     DO j = 1, nmo
                        DO i = 1, nao
                           coeffs((i - 1) + nao*(j - 1) + 1) = local_matrix(i, j)
                        END DO
                     END DO

                     DEALLOCATE (local_matrix)
                     cp2k_get_mo_coefficients = n_total
                  END IF
               END IF
            END IF
         END IF
      END IF

      IF (ASSOCIATED(f_env)) CALL f_env_rm_defaults(f_env, ierr)

   END FUNCTION cp2k_get_mo_coefficients

! **************************************************************************************************
!> \brief Get the number of k-points in the current calculation
!> \param env_id The force environment ID
!> \return Number of k-points, 0 for Gamma-point-only calculations, -1 on error
!> \author CP2K-RS Developers
! **************************************************************************************************
   INTEGER(C_INT) FUNCTION cp2k_get_nkpoints(env_id) BIND(C, name='cp2k_get_nkpoints')
      INTEGER(C_INT), VALUE, INTENT(IN)                  :: env_id

      INTEGER                                            :: ierr, in_use, nkp
      TYPE(f_env_type), POINTER                          :: f_env
      TYPE(kpoint_type), POINTER                         :: kpoints
      TYPE(qs_environment_type), POINTER                 :: qs_env

      cp2k_get_nkpoints = -1
      NULLIFY (f_env, qs_env, kpoints)

      CALL f_env_add_defaults(env_id, f_env)
      IF (.NOT. ASSOCIATED(f_env)) RETURN

      CALL force_env_get(f_env%force_env, in_use=in_use)

      IF (in_use == use_qs_force) THEN
         qs_env => f_env%force_env%qs_env
         CALL get_qs_env(qs_env, kpoints=kpoints)

         IF (ASSOCIATED(kpoints)) THEN
            CALL get_kpoint_info(kpoints, nkp=nkp)
            cp2k_get_nkpoints = nkp
         ELSE
            ! Gamma-point calculation — no explicit k-point set
            cp2k_get_nkpoints = 0
         END IF
      END IF

      IF (ASSOCIATED(f_env)) CALL f_env_rm_defaults(f_env, ierr)

   END FUNCTION cp2k_get_nkpoints

! **************************************************************************************************
!> \brief Get Kohn-Sham eigenvalues for one k-point and one spin channel
!> \param env_id  The force environment ID
!> \param kpt_idx 1-based k-point index
!> \param spin    Spin channel (1=alpha/total, 2=beta)
!> \param eigenvalues Buffer for eigenvalues in Hartree
!> \param nmax    Buffer capacity
!> \return Number of eigenvalues written, or -1 on error
!> \author CP2K-RS Developers
! **************************************************************************************************
   INTEGER(C_INT) FUNCTION cp2k_get_kpoint_eigenvalues(env_id, kpt_idx, spin, eigenvalues, nmax) &
      BIND(C, name='cp2k_get_kpoint_eigenvalues')
      INTEGER(C_INT), VALUE, INTENT(IN)                  :: env_id, kpt_idx, spin, nmax
      REAL(C_DOUBLE), DIMENSION(nmax), INTENT(OUT)       :: eigenvalues

      INTEGER                                            :: i, ierr, in_use, nmo
      REAL(KIND=dp), DIMENSION(:), POINTER               :: eigvals
      TYPE(f_env_type), POINTER                          :: f_env
      TYPE(kpoint_env_p_type), DIMENSION(:), POINTER     :: kp_env
      TYPE(kpoint_type), POINTER                         :: kpoints
      TYPE(mo_set_type), DIMENSION(:, :), POINTER        :: mos
      TYPE(qs_environment_type), POINTER                 :: qs_env

      cp2k_get_kpoint_eigenvalues = -1
      eigenvalues = 0.0_C_DOUBLE
      NULLIFY (f_env, qs_env, kpoints, kp_env, mos, eigvals)

      CALL f_env_add_defaults(env_id, f_env)
      IF (.NOT. ASSOCIATED(f_env)) RETURN

      CALL force_env_get(f_env%force_env, in_use=in_use)

      IF (in_use == use_qs_force) THEN
         qs_env => f_env%force_env%qs_env
         CALL get_qs_env(qs_env, kpoints=kpoints)

         IF (ASSOCIATED(kpoints)) THEN
            CALL get_kpoint_info(kpoints, kp_env=kp_env)

            IF (ASSOCIATED(kp_env)) THEN
               IF (kpt_idx >= 1 .AND. kpt_idx <= SIZE(kp_env)) THEN
                  IF (ASSOCIATED(kp_env(kpt_idx)%kpoint_env)) THEN
                     CALL get_kpoint_env(kp_env(kpt_idx)%kpoint_env, mos=mos)

                     IF (ASSOCIATED(mos)) THEN
                        ! First dimension: 1=real, 2=imaginary; second: spin
                        IF (spin >= 1 .AND. spin <= SIZE(mos, 2)) THEN
                           CALL get_mo_set(mos(1, spin), nmo=nmo, eigenvalues=eigvals)

                           IF (ASSOCIATED(eigvals)) THEN
                              nmo = MIN(nmo, nmax)
                              DO i = 1, nmo
                                 eigenvalues(i) = eigvals(i)
                              END DO
                              cp2k_get_kpoint_eigenvalues = nmo
                           END IF
                        END IF
                     END IF
                  END IF
               END IF
            END IF
         END IF
      END IF

      IF (ASSOCIATED(f_env)) CALL f_env_rm_defaults(f_env, ierr)

   END FUNCTION cp2k_get_kpoint_eigenvalues

END MODULE libcp2k_extended
