cp2k-sys 0.3.1

/*----------------------------------------------------------------------------*/
/*  CP2K-RS Extended Interface Header                                        */
/*  Copyright 2026 Nils Holle                                        */
/*                                                                            */
/*  SPDX-License-Identifier: GPL-2.0-or-later                                */
/*----------------------------------------------------------------------------*/

/**
 * @file libcp2k_extended.h
 * @brief Extended C/C++ interface for CP2K providing access to Quickstep DFT internals
 *
 * This header file provides extended CP2K functionality for accessing internal
 * DFT data structures from Quickstep calculations, including eigenvalues,
 * occupation numbers, charges, and other electronic structure information.
 *
 * @author Nils Holle
 * @date 2025-01
 */

#ifndef LIBCP2K_EXTENDED_H
#define LIBCP2K_EXTENDED_H

#ifdef __cplusplus
extern "C" {
#endif

/*============================================================================*/
/* MPI-Aware Initialization                                                   */
/*============================================================================*/

/**
 * @brief Initialize CP2K with automatic MPI detection
 *
 * This function checks if MPI is already initialized (e.g., by mpi4py in Python
 * or mpi crate in Rust) and calls the appropriate CP2K initialization function.
 * If MPI is already initialized, it uses the existing MPI environment.
 * Otherwise, it initializes MPI as part of CP2K initialization.
 *
 * This is the recommended initialization function when using CP2K with MPI from
 * Python (mpi4py) or Rust (mpi crate), as it avoids double MPI initialization.
 *
 * @return 0 on success, non-zero on failure
 *
 * @par Example (Python with mpi4py):
 * @code{.py}
 * from mpi4py import MPI  # MPI already initialized here
 * import cp2k_rs
 * cp2k_rs.init_cp2k()  # Uses existing MPI, no conflict
 * @endcode
 *
 * @par Example (Rust with mpi crate):
 * @code{.rs}
 * let universe = mpi::initialize().unwrap();  // MPI initialized
 * cp2k_rs::init().unwrap();  // Uses existing MPI
 * @endcode
 */
int cp2k_init_with_mpi_check(void);

/**
 * @brief Finalize CP2K with automatic MPI detection
 *
 * This function finalizes CP2K without finalizing MPI if MPI was already
 * initialized by the host application. The host application (mpi4py or mpi crate)
 * is responsible for finalizing MPI.
 *
 * @return 0 on success, non-zero on failure
 *
 * @par Example:
 * @code{.py}
 * cp2k_rs.finalize_cp2k()  # Finalizes CP2K but not MPI
 * # mpi4py will finalize MPI automatically at exit
 * @endcode
 */
int cp2k_finalize_with_mpi_check(void);

/*============================================================================*/
/* Environment Type Checking                                                  */
/*============================================================================*/

/**
 * @brief Check if a force environment is a Quickstep (DFT) calculation
 *
 * Use this function to verify that the force environment is running a
 * Quickstep DFT calculation before attempting to access DFT-specific data.
 *
 * @param env_id The force environment identifier
 * @return 1 if Quickstep/DFT, 0 otherwise
 *
 * @par Example:
 * @code
 * if (cp2k_is_qs_env(force_env_id)) {
 *     // Safe to call DFT-specific functions
 *     int nmo = cp2k_get_nmo(force_env_id, 1);
 * }
 * @endcode
 */
int cp2k_is_qs_env(int env_id);

/*============================================================================*/
/* Stress and Virial Tensors                                                 */
/*============================================================================*/

/**
 * @brief Get the stress tensor from a force environment
 *
 * The stress tensor is computed from the virial tensor and cell volume.
 * Essential for NPT molecular dynamics and materials science applications.
 *
 * @param env_id The force environment identifier
 * @param stress Pointer to a 3x3 array to store the stress tensor (in GPa)
 *
 * @note The stress tensor is returned in GPa (GigaPascals).
 * @note A valid energy/force calculation must have been performed.
 * @note The stress tensor is symmetric: stress[i][j] == stress[j][i]
 *
 * @par Example:
 * @code
 * double stress[3][3];
 * cp2k_get_stress_tensor(force_env_id, &stress[0][0]);
 * double pressure = (stress[0][0] + stress[1][1] + stress[2][2]) / 3.0;
 * printf("Pressure: %.3f GPa\n", pressure);
 * @endcode
 */
void cp2k_get_stress_tensor(int env_id, double stress[3][3]);

/**
 * @brief Get the virial tensor from a force environment
 *
 * The virial tensor is related to the stress tensor and provides information
 * about internal pressure. Used in thermodynamic property calculations.
 *
 * @param env_id The force environment identifier
 * @param virial_tensor Pointer to a 3x3 array to store the virial tensor (in Hartree)
 *
 * @note The virial tensor is returned in atomic units (Hartree).
 * @note A valid energy/force calculation must have been performed.
 * @note Relationship: stress = -virial / volume * conversion_factor
 *
 * @par Example:
 * @code
 * double virial[3][3];
 * cp2k_get_virial_tensor(force_env_id, &virial[0][0]);
 * @endcode
 */
void cp2k_get_virial_tensor(int env_id, double virial_tensor[3][3]);

/*============================================================================*/
/* Molecular Orbital Information                                             */
/*============================================================================*/

/**
 * @brief Get the number of molecular orbitals for a given spin channel
 *
 * Returns the total number of Kohn-Sham orbitals for the specified spin channel.
 *
 * @param env_id The force environment identifier
 * @param spin Spin channel (1 for alpha/up, 2 for beta/down)
 * @return Number of molecular orbitals, or -1 if not available
 *
 * @note For spin-unpolarized calculations, use spin = 1
 * @note For spin-polarized calculations, query both spin = 1 and spin = 2
 *
 * @par Example:
 * @code
 * int nmo_alpha = cp2k_get_nmo(force_env_id, 1);
 * int nmo_beta = cp2k_get_nmo(force_env_id, 2);
 * if (nmo_beta == -1) {
 *     printf("Spin-unpolarized calculation\n");
 * }
 * @endcode
 */
int cp2k_get_nmo(int env_id, int spin);

/**
 * @brief Get Kohn-Sham orbital eigenvalues (energy levels)
 *
 * Retrieves the eigenvalues (orbital energies) from the last SCF calculation.
 * Essential for band structure analysis and DOS calculations.
 *
 * @param env_id The force environment identifier
 * @param spin Spin channel (1 for alpha/up, 2 for beta/down)
 * @param eigenvalues Buffer to store eigenvalues in Hartree (output)
 * @param nmax Maximum number of eigenvalues to retrieve (buffer size)
 * @return Number of eigenvalues retrieved, or -1 on error
 *
 * @note eigenvalues buffer must be pre-allocated with size >= nmax
 * @note Eigenvalues are in atomic units (Hartree). Convert to eV: E_eV = E_Ha * 27.2114
 * @note A valid SCF calculation must have been performed
 *
 * @par Example:
 * @code
 * int nmo = cp2k_get_nmo(force_env_id, 1);
 * double* eigenvalues = malloc(nmo * sizeof(double));
 * int n = cp2k_get_eigenvalues(force_env_id, 1, eigenvalues, nmo);
 * for (int i = 0; i < n; i++) {
 *     printf("Orbital %d: %.6f Ha (%.3f eV)\n", i+1,
 *            eigenvalues[i], eigenvalues[i] * 27.2114);
 * }
 * @endcode
 */
int cp2k_get_eigenvalues(int env_id, int spin, double* eigenvalues, int nmax);

/**
 * @brief Get orbital occupation numbers
 *
 * Retrieves the occupation numbers for molecular orbitals. Values range from
 * 0 (unoccupied) to 2 (fully occupied, spin-unpolarized) or 1 (spin-polarized).
 *
 * @param env_id The force environment identifier
 * @param spin Spin channel (1 for alpha/up, 2 for beta/down)
 * @param occupations Buffer to store occupation numbers (output)
 * @param nmax Maximum number of occupations to retrieve (buffer size)
 * @return Number of occupation numbers retrieved, or -1 on error
 *
 * @note occupations buffer must be pre-allocated with size >= nmax
 * @note For spin-unpolarized: 0 <= occupation <= 2
 * @note For spin-polarized: 0 <= occupation <= 1
 * @note Fractional occupations occur with smearing or finite temperature
 *
 * @par Example:
 * @code
 * int nmo = cp2k_get_nmo(force_env_id, 1);
 * double* occupations = malloc(nmo * sizeof(double));
 * cp2k_get_occupation_numbers(force_env_id, 1, occupations, nmo);
 *
 * int n_occupied = 0;
 * for (int i = 0; i < nmo; i++) {
 *     if (occupations[i] > 0.5) n_occupied++;
 * }
 * printf("Occupied orbitals: %d\n", n_occupied);
 * @endcode
 */
int cp2k_get_occupation_numbers(int env_id, int spin, double* occupations, int nmax);

/**
 * @brief Get HOMO and LUMO energies and indices
 *
 * Retrieves the energies and orbital indices for the Highest Occupied Molecular
 * Orbital (HOMO) and Lowest Unoccupied Molecular Orbital (LUMO). The gap between
 * these is the fundamental band gap.
 *
 * @param env_id The force environment identifier
 * @param spin Spin channel (1 for alpha/up, 2 for beta/down)
 * @param homo_energy HOMO energy in Hartree (output)
 * @param lumo_energy LUMO energy in Hartree (output)
 * @param homo_index Index of HOMO orbital, 1-based (output)
 * @param lumo_index Index of LUMO orbital, 1-based (output)
 * @return 0 on success, -1 on error
 *
 * @note Energies are in atomic units (Hartree)
 * @note Indices are 1-based (Fortran convention)
 * @note For metallic systems, HOMO and LUMO may overlap (gap = 0)
 *
 * @par Example:
 * @code
 * double homo, lumo;
 * int homo_idx, lumo_idx;
 *
 * if (cp2k_get_homo_lumo(force_env_id, 1, &homo, &lumo,
 *                        &homo_idx, &lumo_idx) == 0) {
 *     double gap_eV = (lumo - homo) * 27.2114;
 *     printf("HOMO (orbital %d): %.4f eV\n", homo_idx, homo * 27.2114);
 *     printf("LUMO (orbital %d): %.4f eV\n", lumo_idx, lumo * 27.2114);
 *     printf("Band gap: %.4f eV\n", gap_eV);
 * }
 * @endcode
 */
int cp2k_get_homo_lumo(int env_id, int spin, double* homo_energy,
                       double* lumo_energy, int* homo_index, int* lumo_index);

/*============================================================================*/
/* Atomic Properties                                                          */
/*============================================================================*/

/**
 * @brief Get Mulliken atomic populations and net charges
 *
 * Retrieves Mulliken population analysis for all atoms in the system.
 * `populations` receives the electronic population per atom (sum over spin);
 * it sums to the total number of valence electrons. `net_charges` receives
 * the net charge per atom (zeff - population); it sums to ~0 for a neutral
 * system. Charges are in units of elementary charge (e).
 *
 * @param env_id The force environment identifier
 * @param populations Buffer to store electronic populations per atom (output)
 * @param net_charges Buffer to store net charges per atom in elementary charge units (output)
 * @param natom Number of atoms (must match system size)
 * @return 0 on success, -1 on error
 *
 * @note Both buffers must be pre-allocated with size >= natom
 * @note Positive net charge = electron deficiency, Negative = electron excess
 * @note Mulliken analysis must be enabled in CP2K input
 *
 * @par Example:
 * @code
 * int natom;
 * cp2k_get_natom(force_env_id, &natom);
 *
 * double* populations = malloc(natom * sizeof(double));
 * double* net_charges = malloc(natom * sizeof(double));
 * if (cp2k_get_mulliken_charges(force_env_id, populations, net_charges, natom) == 0) {
 *     for (int i = 0; i < natom; i++) {
 *         printf("Atom %d: population %.4f, net charge %+.4f e\n",
 *                i+1, populations[i], net_charges[i]);
 *     }
 * }
 * @endcode
 *
 */
int cp2k_get_mulliken_charges(int env_id, double* populations, double* net_charges,
                              int natom);

/**
 * @brief Get the electric dipole moment vector
 *
 * Retrieves the total electric dipole moment of the system, including both
 * electronic and nuclear contributions.
 *
 * @param env_id The force environment identifier
 * @param dipole 3-component dipole moment vector in Debye (output)
 * @return 0 on success, -1 on error
 *
 * @note dipole buffer must be pre-allocated with size 3
 * @note Dipole moment is in Debye units (1 Debye = 3.33564e-30 C·m)
 * @note A valid SCF calculation must have been performed
 *
 * @par Example:
 * @code
 * double dipole[3];
 * if (cp2k_get_dipole_moment(force_env_id, dipole) == 0) {
 *     double magnitude = sqrt(dipole[0]*dipole[0] +
 *                            dipole[1]*dipole[1] +
 *                            dipole[2]*dipole[2]);
 *     printf("Dipole moment: [%.4f, %.4f, %.4f] Debye\n",
 *            dipole[0], dipole[1], dipole[2]);
 *     printf("Magnitude: %.4f Debye\n", magnitude);
 * }
 * @endcode
 *
 */
int cp2k_get_dipole_moment(int env_id, double dipole[3]);

/*============================================================================*/
/* SCF Information                                                            */
/*============================================================================*/

/**
 * @brief Get SCF convergence information
 *
 * Retrieves information about the Self-Consistent Field (SCF) convergence
 * from the last calculation, including iteration count and convergence status.
 *
 * @param env_id The force environment identifier
 * @param niter Number of SCF iterations performed (output)
 * @param converged 1 if SCF converged, 0 otherwise (output)
 * @param energy_change Last energy change in Hartree (output)
 * @return 0 on success, -1 on error
 *
 * @note A valid SCF calculation must have been performed
 * @note energy_change is the last delta_E used for convergence checking
 *
 * @par Example:
 * @code
 * int niter, converged;
 * double delta_e;
 *
 * if (cp2k_get_scf_info(force_env_id, &niter, &converged, &delta_e) == 0) {
 *     printf("SCF iterations: %d\n", niter);
 *     printf("Converged: %s\n", converged ? "yes" : "no");
 *     printf("Final energy change: %.2e Ha\n", delta_e);
 * }
 * @endcode
 */
int cp2k_get_scf_info(int env_id, int* niter, int* converged, double* energy_change);

/*============================================================================*/
/* Energy Components and Electronic Properties                               */
/*============================================================================*/

/**
 * @brief Get energy components (kinetic, Hartree, XC, etc.)
 *
 * Retrieves individual energy components from the last SCF calculation.
 * Useful for understanding energy contributions and validating calculations.
 *
 * @param env_id The force environment identifier
 * @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
 *
 * @note All energies are in atomic units (Hartree)
 * @note e_total = e_kinetic + e_hartree + e_xc + e_core + other terms
 *
 * @par Example:
 * @code
 * double e_kin, e_hartree, e_xc, e_core, e_total;
 * if (cp2k_get_energy_components(force_env_id, &e_kin, &e_hartree,
 *                                 &e_xc, &e_core, &e_total) == 0) {
 *     printf("Kinetic energy: %.6f Ha\n", e_kin);
 *     printf("Hartree energy: %.6f Ha\n", e_hartree);
 *     printf("XC energy: %.6f Ha\n", e_xc);
 *     printf("Total energy: %.6f Ha\n", e_total);
 * }
 * @endcode
 */
int cp2k_get_energy_components(int env_id, double* e_kinetic, double* e_hartree,
                                double* e_xc, double* e_core, double* e_total);

/**
 * @brief Get the number of electrons in the system
 *
 * Returns the total number of electrons from the SCF calculation.
 *
 * @param env_id The force environment identifier
 * @param nelectron Number of electrons (output)
 * @return 0 on success, -1 on error
 *
 * @par Example:
 * @code
 * int nelec;
 * if (cp2k_get_nelectron(force_env_id, &nelec) == 0) {
 *     printf("Number of electrons: %d\n", nelec);
 * }
 * @endcode
 */
int cp2k_get_nelectron(int env_id, int* nelectron);

/**
 * @brief Get the Fermi energy (chemical potential)
 *
 * Retrieves the Fermi energy, which is the chemical potential at which
 * electrons are added to or removed from the system. Particularly important
 * for metallic systems and calculations with smeared occupations.
 *
 * @param env_id The force environment identifier
 * @param e_fermi Fermi energy in Hartree (output)
 * @return 0 on success, -1 on error or if not applicable
 *
 * @note Only meaningful for calculations with smearing or metallic systems
 * @note Returns -1 if Fermi energy is not defined (e.g., isolated molecules)
 *
 * @par Example:
 * @code
 * double e_fermi;
 * if (cp2k_get_fermi_energy(force_env_id, &e_fermi) == 0) {
 *     printf("Fermi energy: %.6f Ha (%.3f eV)\n",
 *            e_fermi, e_fermi * 27.2114);
 * }
 * @endcode
 */
int cp2k_get_fermi_energy(int env_id, double* e_fermi);

/**
 * @brief Get Hirshfeld atomic charges
 *
 * Retrieves Hirshfeld population analysis charges for all atoms.
 * Hirshfeld charges are generally more robust than Mulliken charges.
 *
 * @param env_id The force environment identifier
 * @param charges Buffer to store atomic charges in elementary charge units (output)
 * @param natom Number of atoms (must match system size)
 * @return 0 on success, -1 on error
 *
 * @note charges buffer must be pre-allocated with size >= natom
 * @note Hirshfeld analysis must be enabled in CP2K input
 * @note Positive = electron deficiency, Negative = electron excess
 *
 * @par Example:
 * @code
 * int natom;
 * cp2k_get_natom(force_env_id, &natom);
 *
 * double* charges = malloc(natom * sizeof(double));
 * if (cp2k_get_hirshfeld_charges(force_env_id, charges, natom) == 0) {
 *     for (int i = 0; i < natom; i++) {
 *         printf("Atom %d Hirshfeld charge: %+.4f e\n", i+1, charges[i]);
 *     }
 * }
 * @endcode
 *
 */
int cp2k_get_hirshfeld_charges(int env_id, double* charges, int natom);

/**
 * @brief Get total spin (for spin-polarized calculations)
 *
 * Returns the total spin as the difference between alpha and beta electrons.
 * Only meaningful for spin-polarized calculations.
 *
 * @param env_id The force environment identifier
 * @param total_spin Total spin = N_alpha - N_beta (output)
 * @return 0 on success, -1 on error
 *
 * @note For spin-unpolarized calculations, returns 0
 * @note total_spin = (number of alpha electrons) - (number of beta electrons)
 * @note Related to magnetic moment: μ = total_spin * μ_B (Bohr magneton)
 *
 * @par Example:
 * @code
 * double spin;
 * if (cp2k_get_total_spin(force_env_id, &spin) == 0) {
 *     printf("Total spin: %.2f\n", spin);
 *     if (fabs(spin) < 0.01) {
 *         printf("System is diamagnetic (closed shell)\n");
 *     } else {
 *         printf("System has net spin (open shell)\n");
 *     }
 * }
 * @endcode
 */
int cp2k_get_total_spin(int env_id, double* total_spin);

/*============================================================================*/
/* Electron Density Grid                                                     */
/*============================================================================*/

/**
 * @brief Get grid metadata for the electron density of a spin channel.
 *
 * @param env_id  Force environment ID.
 * @param spin    Spin channel: 1 = alpha/total, 2 = beta.
 * @param npts    Output: grid points per dimension [3].
 * @param origin  Output: grid origin in Bohr [3].
 * @param dh      Output: 3x3 cell increment matrix in Bohr, row-major flat [9].
 * @return 0 on success, -1 on error.
 */
int cp2k_get_grid_info(int env_id, int spin, int* npts, double* origin, double* dh);

/**
 * @brief Get the full electron density on the realspace grid.
 *
 * Data is written in Fortran column-major order (first index varies fastest).
 * For MPI-parallel runs the distributed grid is gathered via MPI_ALLREDUCE.
 *
 * @param env_id   Force environment ID.
 * @param spin     Spin channel: 1 = alpha/total, 2 = beta.
 * @param density  Pre-allocated output buffer (>= npts[0]*npts[1]*npts[2] doubles).
 * @param nmax     Buffer capacity in number of doubles.
 * @return Number of grid points written, or -1 on error.
 */
int cp2k_get_electron_density_grid(int env_id, int spin, double* density, int nmax);

/*============================================================================*/
/* MO Coefficients (Wavefunction)                                            */
/*============================================================================*/

/**
 * @brief Get the dimensions of the MO coefficient matrix.
 *
 * @param env_id  Force environment ID.
 * @param spin    Spin channel: 1 = alpha/total, 2 = beta.
 * @param nao     Output: number of atomic orbitals (rows).
 * @param nmo     Output: number of molecular orbitals (columns).
 * @return 0 on success, -1 on error.
 */
int cp2k_get_mo_coeff_info(int env_id, int spin, int* nao, int* nmo);

/**
 * @brief Get the full MO coefficient matrix.
 *
 * The ScaLAPACK-distributed matrix is gathered to all MPI ranks.
 * Data is written in Fortran column-major order: column j holds the j-th MO
 * expressed in the AO basis.
 *
 * @param env_id  Force environment ID.
 * @param spin    Spin channel: 1 = alpha/total, 2 = beta.
 * @param coeffs  Pre-allocated output buffer (>= nao * nmo doubles).
 * @param nmax    Buffer capacity in number of doubles.
 * @return Number of elements written (nao*nmo), or -1 on error.
 */
int cp2k_get_mo_coefficients(int env_id, int spin, double* coeffs, int nmax);

/*============================================================================*/
/* Future Extensions (Placeholders)                                          */
/*============================================================================*/

/**
 * @brief Get the number of k-points in the current calculation
 *
 * @param env_id The force environment identifier
 * @return Number of k-points (>0), 0 for Gamma-point-only calculations, -1 on error
 */
int cp2k_get_nkpoints(int env_id);

/**
 * @brief Get Kohn-Sham eigenvalues for one k-point and one spin channel
 *
 * @param env_id     The force environment identifier
 * @param kpt_idx    1-based k-point index
 * @param spin       Spin channel (1=alpha/total, 2=beta)
 * @param eigenvalues Pre-allocated output buffer in Hartree
 * @param nmax       Buffer capacity (number of doubles)
 * @return Number of eigenvalues written, or -1 on error
 */
int cp2k_get_kpoint_eigenvalues(int env_id, int kpt_idx, int spin,
                                double* eigenvalues, int nmax);

/**
 * @brief Get orbital occupation numbers for one k-point and one spin channel
 *
 * Companion to cp2k_get_kpoint_eigenvalues for k-point calculations.
 *
 * @param env_id     The force environment identifier
 * @param kpt_idx    1-based k-point index
 * @param spin       Spin channel (1=alpha/total, 2=beta)
 * @param occupations Pre-allocated output buffer for occupation numbers
 * @param nmax       Buffer capacity (number of doubles)
 * @return Number of occupation numbers written, or -1 on error
 */
int cp2k_get_kpoint_occupation_numbers(int env_id, int kpt_idx, int spin,
                                       double* occupations, int nmax);

/*============================================================================*/
/* AO matrices (overlap S and Kohn-Sham H)                                    */
/*============================================================================*/

/**
 * @brief Get the number of AO basis functions (dimension of the square S/KS matrices).
 *
 * @param env_id Force environment ID.
 * @return nao (number of atomic orbitals), or -1 on error.
 *
 * Use this to size the buffers for cp2k_get_overlap_matrix / cp2k_get_ks_matrix,
 * which both return nao*nao doubles.
 */
int cp2k_get_ao_nsize(int env_id);

/**
 * @brief Get the full AO overlap matrix S.
 *
 * The (symmetric, distributed) matrix is desymmetrized to a full dense matrix and
 * gathered to all MPI ranks. Data is written in Fortran column-major order.
 * S is geometry-dependent but SCF-independent (available after the first build).
 *
 * @param env_id Force environment ID.
 * @param s_out  Pre-allocated output buffer (>= nao*nao doubles).
 * @param nmax   Buffer capacity in number of doubles.
 * @return Number of elements written (nao*nao), or -1 on error.
 */
int cp2k_get_overlap_matrix(int env_id, double* s_out, int nmax);

/**
 * @brief Get the full Kohn-Sham matrix H for one spin channel.
 *
 * The (symmetric, distributed) matrix is desymmetrized to a full dense matrix and
 * gathered to all MPI ranks. Data is written in Fortran column-major order.
 * Only valid after a converged SCF. With S, C and the eigenvalues this satisfies
 * H C = S C diag(eps) and C^T S C = 1.
 *
 * @param env_id Force environment ID.
 * @param spin   Spin channel: 1 = alpha/total, 2 = beta.
 * @param h_out  Pre-allocated output buffer (>= nao*nao doubles).
 * @param nmax   Buffer capacity in number of doubles.
 * @return Number of elements written (nao*nao), or -1 on error.
 */
int cp2k_get_ks_matrix(int env_id, int spin, double* h_out, int nmax);

/**
 * @brief Number of non-zero entries of the AO overlap matrix S above `threshold`.
 *
 * Sparse counterpart of cp2k_get_overlap_matrix. CP2K holds S as a DBCSR sparse
 * matrix, so the CSR path never materialises the nao*nao dense matrix — which for
 * a DZVP-MOLOPT basis is ~21 GB at 4096 atoms and ~43 GB at 5824 atoms, per rank.
 *
 * Call this first to size the buffers for cp2k_get_overlap_matrix_csr.
 *
 * @param env_id    Force environment ID.
 * @param threshold Entries with |value| < threshold are dropped.
 * @return Global number of non-zeros, or -1 on error.
 */
int cp2k_get_overlap_matrix_csr_nnz(int env_id, double threshold);

/**
 * @brief Get the AO overlap matrix S in CSR form, with 0-based indices.
 *
 * @param env_id    Force environment ID.
 * @param threshold Entries with |value| < threshold are dropped.
 * @param rowptr    Pre-allocated, nao+1 ints.
 * @param colind    Pre-allocated, nnz ints.
 * @param values    Pre-allocated, nnz doubles.
 * @param nnz_max   Capacity of colind/values.
 * @return Number of non-zeros written, or -1 on error / insufficient capacity.
 */
int cp2k_get_overlap_matrix_csr(int env_id, double threshold, int* rowptr,
                                int* colind, double* values, int nnz_max);

/**
 * @brief Number of non-zero entries of the Kohn-Sham matrix H above `threshold`.
 *
 * @param env_id    Force environment ID.
 * @param spin      Spin channel (1 = alpha/total, 2 = beta).
 * @param threshold Entries with |value| < threshold are dropped.
 * @return Global number of non-zeros, or -1 on error.
 */
int cp2k_get_ks_matrix_csr_nnz(int env_id, int spin, double threshold);

/**
 * @brief Get the Kohn-Sham matrix H for one spin channel in CSR form (0-based).
 *
 * Only valid after a converged SCF.
 *
 * @param env_id    Force environment ID.
 * @param spin      Spin channel (1 = alpha/total, 2 = beta).
 * @param threshold Entries with |value| < threshold are dropped.
 * @param rowptr    Pre-allocated, nao+1 ints.
 * @param colind    Pre-allocated, nnz ints.
 * @param values    Pre-allocated, nnz doubles.
 * @param nnz_max   Capacity of colind/values.
 * @return Number of non-zeros written, or -1 on error / insufficient capacity.
 */
int cp2k_get_ks_matrix_csr(int env_id, int spin, double threshold, int* rowptr,
                           int* colind, double* values, int nnz_max);

/**
 * @brief Get the three Cartesian AO position matrices <phi_mu|r_alpha|phi_nu>.
 *
 * The dipole (position) integral matrices of the AO basis, in Bohr, built with
 * CP2K's own moment-matrix machinery. Each of the three matrices is nao x nao
 * and the three are stored contiguously: x, y, z, each in Fortran column-major
 * order.
 *
 * @param env_id Force environment ID.
 * @param x_out  Pre-allocated output buffer (>= 3*nao*nao doubles).
 * @param nmax   Buffer capacity in number of doubles (must be >= 3*nao*nao).
 * @return Number of elements written (3*nao*nao), or -1 on error.
 */
int cp2k_get_position_matrices(int env_id, double* x_out, int nmax);

/**
 * @brief Number of non-zero entries of one AO position matrix above `threshold`.
 *
 * Sparse counterpart of cp2k_get_position_matrices. The three Cartesian
 * position matrices share S's sparsity pattern, so this never materialises the
 * dense matrix. Query per direction, since a component can vanish by symmetry
 * where the others do not.
 *
 * Call this first to size the buffers for cp2k_get_position_matrices_csr.
 *
 * @param env_id    Force environment ID.
 * @param direction Cartesian direction (1=x, 2=y, 3=z).
 * @param threshold Entries with |value| < threshold are dropped.
 * @return Global number of non-zeros, or -1 on error.
 */
int cp2k_get_position_matrices_csr_nnz(int env_id, int direction, double threshold);

/**
 * @brief Get one AO position matrix <phi_mu|r_alpha|phi_nu> in CSR form (0-based).
 *
 * Sparse counterpart of cp2k_get_position_matrices: the three matrices cost
 * 3*nao*nao doubles dense, versus O(nnz) here. Units are Bohr, as in the
 * dense getter.
 *
 * @param env_id    Force environment ID.
 * @param direction Cartesian direction (1=x, 2=y, 3=z).
 * @param threshold Entries with |value| < threshold are dropped.
 * @param rowptr    Pre-allocated, nao+1 ints.
 * @param colind    Pre-allocated, nnz ints (see cp2k_get_position_matrices_csr_nnz).
 * @param values    Pre-allocated, nnz doubles.
 * @param nnz_max   Capacity of colind/values.
 * @return Number of non-zeros written, or -1 on error / insufficient capacity.
 */
int cp2k_get_position_matrices_csr(int env_id, int direction, double threshold,
                                   int* rowptr, int* colind, double* values,
                                   int nnz_max);

/**
 * @brief Get the AO-to-atom map.
 *
 * For each atomic orbital, writes the 1-based index of the atom it belongs to,
 * in the same AO ordering as the overlap/KS matrices and the MO coefficients.
 * Derived from the overlap matrix block structure (consistent by construction).
 *
 * @param env_id  Force environment ID.
 * @param ao_atom Pre-allocated output buffer (>= nao ints).
 * @param nmax    Buffer capacity in number of ints.
 * @return Number of AOs written (nao), or -1 on error.
 */
int cp2k_get_ao_to_atom(int env_id, int* ao_atom, int nmax);

/**
 * @brief Get the angular momentum (l) and magnetic quantum number (m) per AO.
 *
 * Companion to cp2k_get_ao_to_atom: same AO ordering, so the three arrays
 * together say which atom, which shell and which orientation each basis
 * function belongs to — what an orbital-resolved analysis (e.g. COHP) needs.
 *
 * @param env_id Force environment ID.
 * @param ao_l   Pre-allocated output buffer, receives l per AO (output).
 * @param ao_m   Pre-allocated output buffer, receives m per AO, -l..l (output).
 * @param nmax   Buffer capacity in number of ints (must be >= nao).
 * @return Number of AOs written, or -1 on error.
 */
int cp2k_get_ao_angular_momentum(int env_id, int* ao_l, int* ao_m, int nmax);

/*============================================================================*/
/* K-point matrix access                                                      */
/*============================================================================*/

/**
 * @brief Get the number of MOs for a k-point and spin channel.
 *
 * @param env_id   Force environment ID.
 * @param kpt_idx  1-based k-point index.
 * @param spin     Spin channel: 1 = alpha/total, 2 = beta.
 * @return Number of MOs, or -1 on error.
 */
int cp2k_get_kpoint_nmo(int env_id, int kpt_idx, int spin);

/**
 * @brief Get the KS Hamiltonian (real + imaginary) for a k-point and spin.
 *
 * Both buffers are always written. h_imag is zero-filled when the imaginary
 * KS matrix is not allocated (Gamma-point, time-reversal symmetry).
 *
 * @param env_id   Force environment ID.
 * @param kpt_idx  1-based k-point index.
 * @param spin     Spin channel: 1 = alpha/total, 2 = beta.
 * @param h_real   Pre-allocated buffer (>= nao*nao doubles, column-major).
 * @param h_imag   Pre-allocated buffer (>= nao*nao doubles, column-major).
 * @param nmax     Buffer capacity per buffer.
 * @return Number of elements written per buffer (nao*nao), or -1 on error.
 */
int cp2k_get_kpoint_ks_matrix(int env_id, int kpt_idx, int spin,
                               double* h_real, double* h_imag, int nmax);

/**
 * @brief Get the AO overlap matrix S(k) for a k-point.
 *
 * S(R) is real-symmetric for real Gaussian bases, so there is no
 * imaginary real-space counterpart. S(k) = Σ_R S(R)·exp(i·2π·k·R) is
 * nevertheless complex for k ≠ Γ via the phase factor.
 *
 * @param env_id   Force environment ID.
 * @param kpt_idx  1-based k-point index.
 * @param spin     Kept for API symmetry with cp2k_get_kpoint_ks_matrix;
 *                 S is spin-independent; must be 1.
 * @param s_real   Pre-allocated buffer (>= nao*nao doubles, column-major).
 * @param s_imag   Pre-allocated buffer (>= nao*nao doubles, column-major).
 * @param nmax     Buffer capacity per buffer.
 * @return Number of elements written per buffer (nao*nao), or -1 on error.
 */
int cp2k_get_kpoint_overlap_matrix(int env_id, int kpt_idx, int spin,
                                    double* s_real, double* s_imag, int nmax);

/**
 * @brief Get MO coefficient matrix dimensions for a k-point and spin.
 *
 * @param env_id   Force environment ID.
 * @param kpt_idx  1-based k-point index.
 * @param spin     Spin channel: 1 = alpha/total, 2 = beta.
 * @param nao      Output: number of atomic orbitals.
 * @param nmo      Output: number of molecular orbitals.
 * @return 0 on success, -1 on error.
 */
int cp2k_get_kpoint_mo_coeff_info(int env_id, int kpt_idx, int spin,
                                   int* nao, int* nmo);

/**
 * @brief Get MO coefficients (real + imaginary) for a k-point and spin.
 *
 * c_imag is zero-filled when the imaginary MO set is not allocated (Gamma-point).
 *
 * @param env_id   Force environment ID.
 * @param kpt_idx  1-based k-point index.
 * @param spin     Spin channel: 1 = alpha/total, 2 = beta.
 * @param c_real   Pre-allocated buffer (>= nao*nmo doubles, column-major).
 * @param c_imag   Pre-allocated buffer (>= nao*nmo doubles, column-major).
 * @param nmax     Buffer capacity per buffer.
 * @return Number of elements written per buffer (nao*nmo), or -1 on error.
 */
int cp2k_get_kpoint_mo_coefficients(int env_id, int kpt_idx, int spin,
                                     double* c_real, double* c_imag, int nmax);

/*============================================================================*/
/* Arbitrary k-point access (off the SCF mesh)                                */
/*============================================================================*/

/**
 * @brief Get eigenvalues at an arbitrary k-point list, without a new SCF.
 *
 * Evaluates the already converged density at k-points that need not be in the
 * SCF mesh: the converged real-space matrices are diagonalised at each
 * requested k-point against the fixed density (no SCF is re-run).
 *
 * @param env_id      Force environment ID.
 * @param kpts        3*nkpts fractional coordinates, k-major.
 * @param nkpts       Number of k-points.
 * @param spin        Spin channel (1=alpha/total, 2=beta).
 * @param nadd        Extra empty bands (as ADDED_MOS).
 * @param eigenvalues Pre-allocated output buffer; receives nkpts*nmo values,
 *                    k-major, in Hartree.
 * @param nmax        Capacity of `eigenvalues` in number of doubles.
 * @return Number of eigenvalues written, or -1 on error, or -2 if the SCF
 *         scratch could not be restored.
 */
int cp2k_kpoint_nonscf_eigenvalues(int env_id, const double* kpts, int nkpts,
                                   int spin, int nadd, double* eigenvalues,
                                   int nmax);

/**
 * @brief Get the Kohn-Sham H(k) or overlap S(k) matrix at an arbitrary k-point.
 *
 * Companion to cp2k_get_kpoint_ks_matrix / cp2k_get_kpoint_overlap_matrix,
 * which index into the SCF's own k-point list. This is the Fourier sum of the
 * converged real-space image matrices, so any fractional k-point may be
 * requested, not just the ones in the SCF mesh.
 *
 * @param env_id  Force environment ID.
 * @param kpt     3 fractional coordinates.
 * @param which   1 = Kohn-Sham H(k), 2 = overlap S(k).
 * @param spin    Spin channel for H(k); ignored for S(k).
 * @param m_real  Pre-allocated buffer (>= nao*nao doubles, column-major).
 * @param m_imag  Pre-allocated buffer (>= nao*nao doubles, column-major).
 * @param nmax    Buffer capacity per buffer.
 * @return Number of elements written per buffer (nao*nao), or -1 on error.
 */
int cp2k_kpoint_matrix_at(int env_id, const double* kpt, int which, int spin,
                          double* m_real, double* m_imag, int nmax);

/**
 * @brief Number of non-zero entries of Re/Im H(k) or S(k) at an arbitrary k-point.
 *
 * Companion to cp2k_kpoint_matrix_at_csr, in the same two-call shape as the
 * position-matrix getters: ask for the count, allocate, then ask for the data.
 *
 * Real and imaginary parts share the neighbour-list block pattern, so the two
 * parts of one matrix return the same structure and only the values differ.
 *
 * Call this first to size the buffers for cp2k_kpoint_matrix_at_csr.
 *
 * @param env_id    Force environment ID.
 * @param kpt       3 fractional coordinates.
 * @param which     1 = Kohn-Sham H(k), 2 = overlap S(k).
 * @param spin      Spin channel for H(k); ignored for S(k).
 * @param part      1 = real part, 2 = imaginary part.
 * @param threshold Entries with |value| < threshold are dropped.
 * @return Global number of non-zeros, or -1 on error.
 */
int cp2k_kpoint_matrix_at_csr_nnz(int env_id, const double* kpt, int which,
                                  int spin, int part, double threshold);

/**
 * @brief Get Re or Im of H(k)/S(k) at an arbitrary k-point, in CSR form (0-based).
 *
 * Sparse counterpart of cp2k_kpoint_matrix_at: use this for production-sized
 * cells, where the dense nao*nao form does not fit in RAM.
 *
 * @param env_id    Force environment ID.
 * @param kpt       3 fractional coordinates.
 * @param which     1 = Kohn-Sham H(k), 2 = overlap S(k).
 * @param spin      Spin channel for H(k); ignored for S(k).
 * @param part      1 = real part, 2 = imaginary part.
 * @param threshold Entries with |value| < threshold are dropped.
 * @param rowptr    Pre-allocated, nao+1 ints.
 * @param colind    Pre-allocated, nnz ints (see cp2k_kpoint_matrix_at_csr_nnz).
 * @param values    Pre-allocated, nnz doubles.
 * @param nnz_max   Capacity of colind/values.
 * @return Number of non-zeros written, or -1 on error / insufficient capacity.
 */
int cp2k_kpoint_matrix_at_csr(int env_id, const double* kpt, int which, int spin,
                              int part, double threshold, int* rowptr,
                              int* colind, double* values, int nnz_max);

/**
 * @brief Get the MO coefficients C(k) at an arbitrary k-point list.
 *
 * The C(k) counterpart of cp2k_kpoint_nonscf_eigenvalues: diagonalises the
 * converged Kohn-Sham matrices at the requested k-points against the fixed
 * density (no SCF is re-run).
 *
 * Layout is (nao, nmo) column-major per k-point, k-major overall: k-point ik
 * occupies elements ik*nao*nmo .. (ik+1)*nao*nmo-1 of each buffer. `nmo` is
 * not known in advance; the return value is the total written per buffer, so
 * nmo = total/(nkpts*nao).
 *
 * @param env_id  Force environment ID.
 * @param kpts    3*nkpts fractional coordinates, k-major.
 * @param nkpts   Number of k-points.
 * @param spin    Spin channel (1=alpha/total, 2=beta).
 * @param nadd    Extra empty bands (as ADDED_MOS).
 * @param c_real  Pre-allocated output buffer, receives Re C(k).
 * @param c_imag  Pre-allocated output buffer, receives Im C(k).
 * @param nmax    Capacity of each buffer in number of doubles.
 * @return Number of values written per buffer, -1 on failure, -2 if the SCF
 *         scratch could not be restored.
 */
int cp2k_kpoint_nonscf_mo_coefficients(int env_id, const double* kpts, int nkpts,
                                       int spin, int nadd, double* c_real,
                                       double* c_imag, int nmax);

#ifdef __cplusplus
}
#endif

#endif /* LIBCP2K_EXTENDED_H */