cp2k-rs 0.2.3

//! Example demonstrating the extended DFT interface for CP2K-RS
//!
//! This example shows how to use the extended CP2K bindings to access
//! internal DFT data structures from Quickstep calculations, including:
//! - Stress and virial tensors
//! - Kohn-Sham eigenvalues
//! - Occupation numbers
//! - HOMO/LUMO energies and band gaps
//! - Mulliken charges
//! - Dipole moments
//! - SCF convergence information

use cp2k_rs::{ForceEnv, finalize, init};

fn main() -> Result<(), Box<dyn std::error::Error>> {
    println!("=======================================================");
    println!("  CP2K-RS Extended DFT Interface Example");
    println!("=======================================================\n");

    // Initialize CP2K
    println!("Initializing CP2K...");
    init()?;
    println!("✓ CP2K initialized\n");

    // Note: This example requires an actual DFT calculation input file
    // For demonstration, we'll show the API usage even if the file doesn't exist
    let input_file = "examples/h2o_dft.inp";
    let output_file = "h2o_dft_analysis.out";

    println!("Creating force environment...");
    println!("Input:  {}", input_file);
    println!("Output: {}", output_file);

    // Check if input file exists
    if !std::path::Path::new(input_file).exists() {
        println!("\n⚠ Input file not found!");
        println!("This example demonstrates the API but requires a proper DFT input file.");
        println!("\nExample input file (h2o_dft.inp):");
        println!("-----------------------------------");
        print_example_input();
        println!("-----------------------------------\n");

        finalize()?;
        println!("To run this example, create the input file above and rerun.");
        return Ok(());
    }

    match ForceEnv::new(input_file, output_file) {
        Ok(mut force_env) => {
            println!("✓ Force environment created\n");

            // Check if this is a Quickstep (DFT) calculation
            #[cfg(feature = "extended")]
            {
                if force_env.is_quickstep() {
                    println!("✓ This is a Quickstep DFT calculation\n");

                    // Run energy and force calculation
                    println!("Running energy and force calculation...");
                    match force_env.calc_energy_force() {
                        Ok(_) => {
                            println!("✓ Calculation completed\n");

                            // Analyze the results
                            analyze_basic_results(&force_env)?;
                            analyze_stress(&force_env)?;
                            analyze_electronic_structure(&force_env)?;
                            analyze_properties(&force_env)?;
                            analyze_scf(&force_env)?;
                        }
                        Err(e) => {
                            println!("✗ Calculation failed: {}", e);
                        }
                    }
                } else {
                    println!("✗ Not a DFT calculation - extended features not available");
                }
            }

            #[cfg(not(feature = "extended"))]
            {
                println!("⚠ Extended features not enabled!");
                println!("Build with --features extended to use DFT-specific functions.");
            }
        }
        Err(e) => {
            println!("✗ Failed to create force environment: {}", e);
        }
    }

    // Finalize CP2K
    println!("\nFinalizing CP2K...");
    finalize()?;
    println!("✓ CP2K finalized");

    println!("\n=======================================================");
    println!("  Example completed!");
    println!("=======================================================");

    Ok(())
}

#[cfg(feature = "extended")]
fn analyze_basic_results(force_env: &ForceEnv) -> Result<(), Box<dyn std::error::Error>> {
    println!("─────────────────────────────────────────────────────");
    println!("  Basic Results");
    println!("─────────────────────────────────────────────────────");

    // Get potential energy
    match force_env.get_potential_energy() {
        Ok(energy) => {
            const HARTREE_TO_EV: f64 = 27.211386245988;
            println!("Potential Energy:");
            println!("  {:.8} Hartree", energy);
            println!("  {:.6} eV", energy * HARTREE_TO_EV);
        }
        Err(e) => println!("✗ Failed to get energy: {}", e),
    }

    // Get number of atoms
    match force_env.get_natom() {
        Ok(natom) => {
            println!("\nNumber of atoms: {}", natom);
        }
        Err(e) => println!("✗ Failed to get natom: {}", e),
    }

    println!();
    Ok(())
}

#[cfg(feature = "extended")]
fn analyze_stress(force_env: &ForceEnv) -> Result<(), Box<dyn std::error::Error>> {
    println!("─────────────────────────────────────────────────────");
    println!("  Stress and Virial Tensors");
    println!("─────────────────────────────────────────────────────");

    // Get stress tensor
    match force_env.get_stress_tensor() {
        Ok(stress) => {
            println!("Stress Tensor (GPa):");
            for i in 0..3 {
                println!(
                    "  [{:>10.4} {:>10.4} {:>10.4}]",
                    stress[[i, 0]],
                    stress[[i, 1]],
                    stress[[i, 2]]
                );
            }

            // Calculate pressure (average of diagonal elements)
            let pressure = (stress[[0, 0]] + stress[[1, 1]] + stress[[2, 2]]) / 3.0;
            println!("\nPressure (isotropic): {:.4} GPa", pressure);
        }
        Err(e) => println!("✗ Failed to get stress tensor: {}", e),
    }

    // Get virial tensor
    match force_env.get_virial_tensor() {
        Ok(virial) => {
            println!("\nVirial Tensor (Hartree):");
            for i in 0..3 {
                println!(
                    "  [{:>12.6e} {:>12.6e} {:>12.6e}]",
                    virial[[i, 0]],
                    virial[[i, 1]],
                    virial[[i, 2]]
                );
            }
        }
        Err(e) => println!("✗ Failed to get virial tensor: {}", e),
    }

    println!();
    Ok(())
}

#[cfg(feature = "extended")]
fn analyze_electronic_structure(force_env: &ForceEnv) -> Result<(), Box<dyn std::error::Error>> {
    println!("─────────────────────────────────────────────────────");
    println!("  Electronic Structure");
    println!("─────────────────────────────────────────────────────");

    const HARTREE_TO_EV: f64 = 27.211386245988;

    // Analyze alpha spin (spin 1)
    for spin in 1..=2 {
        let spin_label = if spin == 1 { "Alpha" } else { "Beta" };

        match force_env.get_nmo(spin) {
            Ok(nmo) => {
                println!("\n{} Spin Channel:", spin_label);
                println!("  Number of MOs: {}", nmo);

                // Get eigenvalues
                match force_env.get_eigenvalues(spin) {
                    Ok(eigenvalues) => {
                        println!("  Eigenvalues (first 5 and last 5):");
                        let n_show = 5.min(eigenvalues.len());

                        for i in 0..n_show {
                            println!(
                                "    Orbital {:>3}: {:>10.6} Ha ({:>8.3} eV)",
                                i + 1,
                                eigenvalues[i],
                                eigenvalues[i] * HARTREE_TO_EV
                            );
                        }

                        if eigenvalues.len() > 10 {
                            println!("    ...");
                            for i in (eigenvalues.len() - n_show)..eigenvalues.len() {
                                println!(
                                    "    Orbital {:>3}: {:>10.6} Ha ({:>8.3} eV)",
                                    i + 1,
                                    eigenvalues[i],
                                    eigenvalues[i] * HARTREE_TO_EV
                                );
                            }
                        }
                    }
                    Err(e) => println!("  ✗ Failed to get eigenvalues: {}", e),
                }

                // Get occupation numbers
                match force_env.get_occupation_numbers(spin) {
                    Ok(occupations) => {
                        let total_electrons: f64 = occupations.iter().sum();
                        println!("  Total electrons: {:.2}", total_electrons);

                        // Count occupied orbitals
                        let n_occupied = occupations.iter().filter(|&&o| o > 0.5).count();
                        println!("  Occupied orbitals: {}", n_occupied);
                    }
                    Err(e) => println!("  ✗ Failed to get occupations: {}", e),
                }

                // Get HOMO/LUMO
                match force_env.get_homo_lumo(spin) {
                    Ok((homo_e, lumo_e, homo_idx, lumo_idx)) => {
                        println!("\n  HOMO/LUMO Information:");
                        println!(
                            "    HOMO (orbital {}): {:.6} Ha ({:.3} eV)",
                            homo_idx,
                            homo_e,
                            homo_e * HARTREE_TO_EV
                        );
                        println!(
                            "    LUMO (orbital {}): {:.6} Ha ({:.3} eV)",
                            lumo_idx,
                            lumo_e,
                            lumo_e * HARTREE_TO_EV
                        );

                        let gap = (lumo_e - homo_e) * HARTREE_TO_EV;
                        println!("    Band gap: {:.4} eV", gap);
                    }
                    Err(e) => println!("  ✗ Failed to get HOMO/LUMO: {}", e),
                }

                // Alternative: use the convenience method
                if let Ok(gap_ev) = force_env.get_band_gap(spin) {
                    println!("  Band gap (from convenience method): {:.4} eV", gap_ev);
                }
            }
            Err(_) => {
                if spin == 2 {
                    println!("\n(Spin-unpolarized calculation - no beta channel)");
                }
                break;
            }
        }
    }

    println!();
    Ok(())
}

#[cfg(feature = "extended")]
fn analyze_properties(force_env: &ForceEnv) -> Result<(), Box<dyn std::error::Error>> {
    println!("─────────────────────────────────────────────────────");
    println!("  Atomic Properties");
    println!("─────────────────────────────────────────────────────");

    // Get Mulliken charges
    match force_env.get_mulliken_charges() {
        Ok(charges) => {
            println!("Mulliken Charges (elementary charge):");
            for (i, charge) in charges.iter().enumerate() {
                println!("  Atom {:>3}: {:>8.4} e", i + 1, charge);
            }

            let total_charge: f64 = charges.iter().sum();
            println!("  Total: {:>8.4} e", total_charge);
        }
        Err(e) => println!("✗ Mulliken charges not available: {}", e),
    }

    // Get dipole moment
    match force_env.get_dipole_moment() {
        Ok(dipole) => {
            let magnitude = (dipole[0].powi(2) + dipole[1].powi(2) + dipole[2].powi(2)).sqrt();
            println!("\nDipole Moment:");
            println!("  X: {:>10.4} Debye", dipole[0]);
            println!("  Y: {:>10.4} Debye", dipole[1]);
            println!("  Z: {:>10.4} Debye", dipole[2]);
            println!("  |μ|: {:>8.4} Debye", magnitude);
        }
        Err(e) => println!("✗ Dipole moment not available: {}", e),
    }

    println!();
    Ok(())
}

#[cfg(feature = "extended")]
fn analyze_scf(force_env: &ForceEnv) -> Result<(), Box<dyn std::error::Error>> {
    println!("─────────────────────────────────────────────────────");
    println!("  SCF Convergence Information");
    println!("─────────────────────────────────────────────────────");

    match force_env.get_scf_info() {
        Ok((niter, converged, delta_e)) => {
            println!("SCF iterations: {}", niter);
            println!("Converged: {}", if converged { "Yes" } else { "No" });
            println!("Final energy change: {:.2e} Hartree", delta_e);
        }
        Err(e) => println!("✗ SCF info not available: {}", e),
    }

    println!();
    Ok(())
}

#[cfg(not(feature = "extended"))]
fn analyze_basic_results(_force_env: &ForceEnv) -> Result<(), Box<dyn std::error::Error>> {
    Ok(())
}

#[cfg(not(feature = "extended"))]
fn analyze_stress(_force_env: &ForceEnv) -> Result<(), Box<dyn std::error::Error>> {
    Ok(())
}

#[cfg(not(feature = "extended"))]
fn analyze_electronic_structure(_force_env: &ForceEnv) -> Result<(), Box<dyn std::error::Error>> {
    Ok(())
}

#[cfg(not(feature = "extended"))]
fn analyze_properties(_force_env: &ForceEnv) -> Result<(), Box<dyn std::error::Error>> {
    Ok(())
}

#[cfg(not(feature = "extended"))]
fn analyze_scf(_force_env: &ForceEnv) -> Result<(), Box<dyn std::error::Error>> {
    Ok(())
}

fn print_example_input() {
    println!(
        r#"
&GLOBAL
  PROJECT H2O_DFT
  RUN_TYPE ENERGY_FORCE
  PRINT_LEVEL MEDIUM
&END GLOBAL

&FORCE_EVAL
  METHOD Quickstep
  STRESS_TENSOR ANALYTICAL

  &DFT
    BASIS_SET_FILE_NAME BASIS_MOLOPT
    POTENTIAL_FILE_NAME GTH_POTENTIALS

    &QS
      EPS_DEFAULT 1.0E-10
    &END QS

    &MGRID
      CUTOFF 400
      REL_CUTOFF 50
    &END MGRID

    &SCF
      SCF_GUESS ATOMIC
      EPS_SCF 1.0E-6
      MAX_SCF 50
      &OT
        MINIMIZER DIIS
        PRECONDITIONER FULL_SINGLE_INVERSE
      &END OT
    &END SCF

    &XC
      &XC_FUNCTIONAL PBE
      &END XC_FUNCTIONAL
    &END XC

    &PRINT
      &MULLIKEN ON
      &END MULLIKEN
    &END PRINT
  &END DFT

  &SUBSYS
    &CELL
      ABC 10.0 10.0 10.0
      PERIODIC NONE
    &END CELL

    &COORD
      O   0.000000    0.000000    0.000000
      H   0.757136    0.586047    0.000000
      H  -0.757136    0.586047    0.000000
    &END COORD

    &KIND H
      BASIS_SET DZVP-MOLOPT-GTH
      POTENTIAL GTH-PBE-q1
    &END KIND

    &KIND O
      BASIS_SET DZVP-MOLOPT-GTH
      POTENTIAL GTH-PBE-q6
    &END KIND
  &END SUBSYS
&END FORCE_EVAL
"#
    );
}