/*----------------------------------------------------------------------------*/
/*  CP2K-RS Extended Interface Header                                        */
/*  Copyright 2025 CP2K-RS developers                                        */
/*                                                                            */
/*  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 CP2K-RS Developers
 * @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 charges
 * 
 * Retrieves Mulliken population analysis charges for all atoms in the system.
 * Charges are in units of elementary charge (e).
 * 
 * @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 Charges are in elementary charge units (e)
 * @note Positive = 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* charges = malloc(natom * sizeof(double));
 * if (cp2k_get_mulliken_charges(force_env_id, charges, natom) == 0) {
 *     for (int i = 0; i < natom; i++) {
 *         printf("Atom %d charge: %+.4f e\n", i+1, charges[i]);
 *     }
 * }
 * @endcode
 * 
 */
int cp2k_get_mulliken_charges(int env_id, double* 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 density matrix elements
 * @note Not yet implemented - placeholder for future development
 */
/* int cp2k_get_density_matrix(int env_id, int spin, double* matrix, int nmax); */

/**
 * @brief Get Hamiltonian/Fock matrix elements
 * @note Not yet implemented - placeholder for future development
 */
/* int cp2k_get_hamiltonian_matrix(int env_id, int spin, double* matrix, int nmax); */

/**
 * @brief Get overlap matrix elements
 * @note Not yet implemented - placeholder for future development
 */
/* int cp2k_get_overlap_matrix(int env_id, double* matrix, int nmax); */

#ifdef __cplusplus
}
#endif

#endif /* LIBCP2K_EXTENDED_H */