//! Comprehensive DFT Analysis Example
//!
//! This example demonstrates all extended interface features:
//! - Basic properties (energy, forces, positions)
//! - Electronic structure (eigenvalues, HOMO/LUMO, band gap)
//! - Stress/pressure analysis
//! - SCF convergence info
//! - Energy components
//! - Electron count and spin
//! - Fermi energy (if applicable)

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

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

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

    // Input and output files
    let input_file = "examples/h2o_dft.inp";
    let output_file = "h2o_comprehensive.out";

    println!("Input:  {}", input_file);
    println!("Output: {}\n", output_file);

    // Create force environment
    println!("Creating force environment...");
    let mut force_env = ForceEnv::new(input_file, output_file)?;
    println!("✓ Force environment created\n");

    // Check if Quickstep
    #[cfg(feature = "extended")]
    {
        if !force_env.is_quickstep() {
            eprintln!("ERROR: Not a Quickstep (DFT) calculation!");
            finalize()?;
            return Err("Not a DFT calculation".into());
        }
        println!("✓ Confirmed Quickstep DFT calculation\n");
    }

    // Run calculation
    println!("=================================================================");
    println!("Running DFT calculation...");
    println!("=================================================================\n");
    force_env.calc_energy_force()?;
    println!("✓ Calculation completed\n");

    // Extract and display all properties
    println!("=================================================================");
    println!("BASIC PROPERTIES");
    println!("=================================================================\n");

    let natom = force_env.get_natom()?;
    let energy = force_env.get_potential_energy()?;
    let forces = force_env.get_forces()?;
    let positions = force_env.get_positions()?;

    println!("Number of atoms:    {}", natom);
    println!("Total energy:       {:.8} Ha", energy);
    println!("                    {:.6} eV", energy * 27.2114);
    println!("Forces shape:       [{}]", forces.len());
    println!("Positions shape:    [{}]\n", positions.len());

    #[cfg(feature = "extended")]
    {
        // Energy components
        println!("=================================================================");
        println!("ENERGY COMPONENTS");
        println!("=================================================================\n");

        if let Ok((e_kin, e_hartree, e_xc, e_core, e_total)) = force_env.get_energy_components() {
            println!("Kinetic energy:     {:.8} Ha", e_kin);
            println!("Hartree energy:     {:.8} Ha", e_hartree);
            println!("XC energy:          {:.8} Ha", e_xc);
            println!("Core energy:        {:.8} Ha", e_core);
            println!("Total energy:       {:.8} Ha", e_total);
            println!("                    {:.6} eV\n", e_total * 27.2114);
        }

        // Electronic properties
        println!("=================================================================");
        println!("ELECTRONIC PROPERTIES");
        println!("=================================================================\n");

        if let Ok(nelec) = force_env.get_nelectron() {
            println!("Number of electrons: {}", nelec);
        }

        if let Ok(spin) = force_env.get_total_spin() {
            println!("Total spin:          {:.4}", spin);
            if spin.abs() < 0.01 {
                println!("System type:         Diamagnetic (closed shell)");
            } else {
                println!("System type:         Paramagnetic (open shell)");
            }
        }

        if let Ok(e_fermi) = force_env.get_fermi_energy() {
            println!("Fermi energy:        {:.6} Ha", e_fermi);
            println!("                     {:.4} eV", e_fermi * 27.2114);
        } else {
            println!("Fermi energy:        N/A (molecular system)");
        }
        println!();

        // Molecular orbitals
        println!("=================================================================");
        println!("MOLECULAR ORBITAL ANALYSIS");
        println!("=================================================================\n");

        for spin in 1..=2 {
            let spin_label = if spin == 1 { "Alpha" } else { "Beta" };

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

                    if let Ok(eigenvals) = force_env.get_eigenvalues(spin) {
                        let min_e = eigenvals.iter().cloned().fold(f64::INFINITY, f64::min);
                        let max_e = eigenvals.iter().cloned().fold(f64::NEG_INFINITY, f64::max);
                        println!("  Energy range:      {:.4} to {:.4} Ha", min_e, max_e);
                        println!(
                            "                     {:.3} to {:.3} eV",
                            min_e * 27.2114,
                            max_e * 27.2114
                        );
                    }

                    if let Ok(occs) = force_env.get_occupation_numbers(spin) {
                        let total_elec: f64 = occs.iter().sum();
                        let n_occ = occs.iter().filter(|&&x| x > 0.5).count();
                        println!("  Occupied MOs:      {}", n_occ);
                        println!("  Electrons:         {:.2}", total_elec);
                    }

                    if let Ok((homo, lumo, homo_idx, lumo_idx)) = force_env.get_homo_lumo(spin) {
                        let gap_ha = lumo - homo;
                        let gap_ev = gap_ha * 27.2114;

                        println!("\n  Frontier Orbitals:");
                        println!(
                            "    HOMO (orbital {}): {:.6} Ha ({:.3} eV)",
                            homo_idx,
                            homo,
                            homo * 27.2114
                        );
                        println!(
                            "    LUMO (orbital {}): {:.6} Ha ({:.3} eV)",
                            lumo_idx,
                            lumo,
                            lumo * 27.2114
                        );
                        println!("    Band gap:          {:.6} Ha ({:.4} eV)", gap_ha, gap_ev);

                        if gap_ev < 0.5 {
                            println!("    Character:         Metallic/Small gap");
                        } else if gap_ev < 3.0 {
                            println!("    Character:         Semiconductor");
                        } else {
                            println!("    Character:         Insulator/Large gap");
                        }
                    }
                    println!();
                }
                Err(_) => {
                    if spin == 2 {
                        println!("Beta spin: N/A (spin-unpolarized calculation)\n");
                    }
                    break;
                }
            }
        }

        // Stress and pressure
        println!("=================================================================");
        println!("STRESS AND PRESSURE ANALYSIS");
        println!("=================================================================\n");

        if let Ok(stress) = force_env.get_stress_tensor() {
            println!("Stress tensor (GPa):");
            for i in 0..3 {
                println!(
                    "  [{:10.6} {:10.6} {:10.6}]",
                    stress[[i, 0]],
                    stress[[i, 1]],
                    stress[[i, 2]]
                );
            }

            let pressure = -(stress[[0, 0]] + stress[[1, 1]] + stress[[2, 2]]) / 3.0;
            println!("\nPressure:            {:.6} GPa", pressure);
            println!("                     {:.4} bar", pressure * 10000.0);

            // Check if stress is nearly diagonal
            let off_diag = stress[[0, 1]].abs()
                + stress[[0, 2]].abs()
                + stress[[1, 0]].abs()
                + stress[[1, 2]].abs()
                + stress[[2, 0]].abs()
                + stress[[2, 1]].abs();
            if off_diag < 0.01 {
                println!("Stress character:    Nearly diagonal (isotropic)");
            } else {
                println!("Stress character:    Non-diagonal (anisotropic)");
            }
        }

        if let Ok(virial) = force_env.get_virial_tensor() {
            let virial_trace = virial[[0, 0]] + virial[[1, 1]] + virial[[2, 2]];
            println!("\nVirial trace:        {:.6} Ha", virial_trace);
        }
        println!();

        // SCF convergence
        println!("=================================================================");
        println!("SCF CONVERGENCE INFORMATION");
        println!("=================================================================\n");

        if let Ok((niter, converged, delta_e)) = force_env.get_scf_info() {
            println!("SCF iterations:      {}", niter);
            println!(
                "Converged:           {}",
                if converged { "Yes" } else { "No" }
            );
            println!("Final ΔE:            {:.2e} Ha", delta_e);

            if converged {
                println!("Status:              ✓ SCF converged successfully");
            } else {
                println!("Status:              ⚠ SCF did not converge");
            }
        }
        println!();

        // Atomic charges
        println!("=================================================================");
        println!("ATOMIC CHARGES");
        println!("=================================================================\n");

        if let Ok(mulliken) = force_env.get_mulliken_charges() {
            let sum: f64 = mulliken.iter().sum();
            println!("Mulliken charges:");
            for (i, &q) in mulliken.iter().enumerate() {
                println!("  Atom {:2}: {:+.6} e", i + 1, q);
            }
            println!("  Total:   {:+.6} e", sum);
        } else {
            println!("Mulliken charges:    Not available (AO density not present)");
        }

        if let Ok(hirshfeld) = force_env.get_hirshfeld_charges() {
            let sum: f64 = hirshfeld.iter().sum();
            println!("\nHirshfeld charges:");
            for (i, &q) in hirshfeld.iter().enumerate() {
                println!("  Atom {:2}: {:+.6} e", i + 1, q);
            }
            println!("  Total:   {:+.6} e", sum);
        } else {
            println!("\nHirshfeld charges:   Not available (real-space density not present)");
        }
        println!();

        // Dipole moment
        println!("=================================================================");
        println!("MOLECULAR PROPERTIES");
        println!("=================================================================\n");

        if let Ok(dipole) = force_env.get_dipole_moment() {
            let magnitude = (dipole[0].powi(2) + dipole[1].powi(2) + dipole[2].powi(2)).sqrt();
            println!("Dipole moment (Debye):");
            println!("  X component:       {:10.6}", dipole[0]);
            println!("  Y component:       {:10.6}", dipole[1]);
            println!("  Z component:       {:10.6}", dipole[2]);
            println!("  Magnitude:         {:10.6}", magnitude);

            if magnitude < 0.1 {
                println!("  Character:         Non-polar");
            } else if magnitude < 2.0 {
                println!("  Character:         Weakly polar");
            } else {
                println!("  Character:         Polar");
            }
        } else {
            println!("Dipole moment:       Not available (AO density not present)");
        }
        println!();
    }

    // Explicitly drop force_env before finalize().  CP2K's f77_interface
    // aborts with "invalid env_id" if cp2k_destroy_force_env is called
    // after the environment has been finalized.
    drop(force_env);

    // Finalize
    println!("=================================================================");
    println!("Finalizing CP2K...");
    finalize()?;
    println!("✓ CP2K finalized\n");

    println!("=================================================================");
    println!(
        "  Analysis complete! See {} for detailed CP2K output.",
        output_file
    );
    println!("=================================================================");

    Ok(())
}
