//! Run a real CP2K DFT calculation and save results to text file
//!
//! This example demonstrates:
//! - Running an actual Quickstep DFT calculation
//! - Extracting all extended interface data
//! - Saving results in human-readable text format for analysis

use cp2k_rs::{ForceEnv, finalize, init};
use ndarray::{Array1, Array2};
use std::collections::HashMap;
use std::fs::File;
use std::io::Write;

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

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

    // Input and output files
    let input_file = concat!(env!("CARGO_MANIFEST_DIR"), "/examples/h2o_dft.inp");
    let output_file = concat!(
        env!("CARGO_MANIFEST_DIR"),
        "/examples/h2o_dft_calculation.out"
    );

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

    // Check if input file exists
    if !std::path::Path::new(input_file).exists() {
        eprintln!("ERROR: Input file '{}' not found!", input_file);
        eprintln!("Please ensure the file exists before running.");
        finalize()?;
        return Ok(());
    }

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

    // Check if this is a Quickstep calculation
    #[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");

        // Run the calculation
        println!("\n{}", "=".repeat(65));
        println!("Running DFT calculation...");
        println!("{}", "=".repeat(65));
        force_env.calc_energy_force()?;
        println!("✓ Calculation completed successfully\n");

        // Extract all data
        println!("{}", "=".repeat(65));
        println!("Extracting results...");
        println!("{}", "=".repeat(65));

        let mut results: HashMap<String, ResultData> = HashMap::new();

        // Basic information
        println!("\n[1/10] Getting basic information...");
        let natom = force_env.get_natom()?;
        let energy = force_env.get_potential_energy()?;
        println!("  - Number of atoms: {}", natom);
        println!(
            "  - Potential energy: {:.8} Ha ({:.6} eV)",
            energy,
            energy * 27.2114
        );
        results.insert("natom".to_string(), ResultData::Scalar(natom as f64));
        results.insert("energy_hartree".to_string(), ResultData::Scalar(energy));
        results.insert(
            "energy_ev".to_string(),
            ResultData::Scalar(energy * 27.211386245988),
        );

        // Forces
        println!("\n[2/10] Getting forces...");
        let forces = force_env.get_forces()?;
        println!("  - Forces shape: {:?}", forces.shape());
        results.insert("forces".to_string(), ResultData::Array1D(forces));

        // Positions
        println!("\n[3/10] Getting positions...");
        let positions = force_env.get_positions()?;
        println!("  - Positions shape: {:?}", positions.shape());
        results.insert("positions".to_string(), ResultData::Array1D(positions));

        // Stress tensor
        println!("\n[4/10] Getting stress tensor...");
        match force_env.get_stress_tensor() {
            Ok(stress) => {
                let pressure = (stress[[0, 0]] + stress[[1, 1]] + stress[[2, 2]]) / 3.0;
                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]]
                    );
                }
                println!("  - Pressure: {:.4} GPa", pressure);
                results.insert("stress_tensor_gpa".to_string(), ResultData::Array2D(stress));
                results.insert("pressure_gpa".to_string(), ResultData::Scalar(pressure));
            }
            Err(e) => println!("  ! Stress tensor not available: {}", e),
        }

        // Virial tensor
        println!("\n[5/10] Getting virial tensor...");
        match force_env.get_virial_tensor() {
            Ok(virial) => {
                println!("  ✓ Virial tensor retrieved");
                results.insert(
                    "virial_tensor_hartree".to_string(),
                    ResultData::Array2D(virial),
                );
            }
            Err(e) => println!("  ! Virial tensor not available: {}", e),
        }

        // Molecular orbitals - try both spin channels
        println!("\n[6/10] Getting molecular orbital information...");
        for spin in 1..=2 {
            let spin_label = if spin == 1 { "alpha" } else { "beta" };

            match force_env.get_nmo(spin) {
                Ok(nmo) => {
                    println!("  - {} spin: {} molecular orbitals", spin_label, nmo);

                    // Eigenvalues
                    if let Ok(eigenvalues) = force_env.get_eigenvalues(spin) {
                        if !eigenvalues.is_empty() {
                            println!(
                                "    ✓ Eigenvalues: {} values ({:.4} to {:.4} Ha)",
                                eigenvalues.len(),
                                eigenvalues.first().copied().unwrap_or(0.0),
                                eigenvalues.last().copied().unwrap_or(0.0)
                            );
                        } else {
                            println!("    ✓ Eigenvalues: 0 values (empty)");
                        }
                        results.insert(
                            format!("eigenvalues_{}", spin_label),
                            ResultData::Array1D(eigenvalues),
                        );
                    }

                    // Occupation numbers
                    if let Ok(occupations) = force_env.get_occupation_numbers(spin) {
                        let total_electrons: f64 = occupations.iter().sum();
                        println!("    ✓ Occupations: {:.2} total electrons", total_electrons);
                        results.insert(
                            format!("occupations_{}", spin_label),
                            ResultData::Array1D(occupations),
                        );
                    }

                    // HOMO/LUMO
                    if let Ok((homo_e, lumo_e, homo_idx, lumo_idx)) = force_env.get_homo_lumo(spin)
                    {
                        let gap_ev = (lumo_e - homo_e) * 27.211386245988;
                        println!("    ✓ HOMO (orbital {}): {:.6} Ha", homo_idx, homo_e);
                        println!("    ✓ LUMO (orbital {}): {:.6} Ha", lumo_idx, lumo_e);
                        println!("    ✓ Band gap: {:.4} eV", gap_ev);

                        results.insert(
                            format!("homo_energy_{}_hartree", spin_label),
                            ResultData::Scalar(homo_e),
                        );
                        results.insert(
                            format!("lumo_energy_{}_hartree", spin_label),
                            ResultData::Scalar(lumo_e),
                        );
                        results.insert(
                            format!("homo_index_{}", spin_label),
                            ResultData::Scalar(homo_idx as f64),
                        );
                        results.insert(
                            format!("lumo_index_{}", spin_label),
                            ResultData::Scalar(lumo_idx as f64),
                        );
                        results.insert(
                            format!("band_gap_{}_ev", spin_label),
                            ResultData::Scalar(gap_ev),
                        );
                    }
                }
                Err(_) => {
                    if spin == 2 {
                        println!("  - Spin-unpolarized calculation (no beta channel)");
                    }
                    break;
                }
            }
        }

        // Cell
        println!("\n[7/10] Getting cell information...");
        match force_env.get_cell() {
            Ok(cell) => {
                println!("  ✓ Cell retrieved");
                results.insert("cell".to_string(), ResultData::Array2D(cell));
            }
            Err(e) => println!("  ! Cell not available: {}", e),
        }

        // Mulliken charges
        println!("\n[8/10] Getting Mulliken charges...");
        match force_env.get_mulliken_charges() {
            Ok(charges) => {
                println!("  ✓ Mulliken charges retrieved");
                results.insert("mulliken_charges".to_string(), ResultData::Array1D(charges));
            }
            Err(e) => println!("  ! Mulliken charges not available: {}", e),
        }

        // Dipole moment
        println!("\n[9/10] Getting 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!("  ✓ Dipole moment: {:.4} Debye", magnitude);
                results.insert(
                    "dipole_moment_debye".to_string(),
                    ResultData::Array1D(dipole),
                );
                results.insert(
                    "dipole_magnitude_debye".to_string(),
                    ResultData::Scalar(magnitude),
                );
            }
            Err(e) => println!("  ! Dipole moment not available: {}", e),
        }

        // SCF convergence information
        println!("\n[10/10] Getting SCF convergence information...");
        match force_env.get_scf_info() {
            Ok((niter, converged, delta_e)) => {
                println!("  ✓ SCF iterations: {}", niter);
                println!("  ✓ Converged: {}", converged);
                println!("  ✓ Final ΔE: {:.2e} Ha", delta_e);
                results.insert(
                    "scf_iterations".to_string(),
                    ResultData::Scalar(niter as f64),
                );
                results.insert(
                    "scf_converged".to_string(),
                    ResultData::Scalar(if converged { 1.0 } else { 0.0 }),
                );
                results.insert(
                    "scf_delta_e_hartree".to_string(),
                    ResultData::Scalar(delta_e),
                );
            }
            Err(e) => println!("  ! SCF info not available: {}", e),
        }

        // Save to NPZ file
        println!("\n{}", "=".repeat(65));
        println!("Saving results to file...");
        println!("{}", "=".repeat(65));

        let results_filename = "cp2k_dft_results.txt";
        save_results(&results, results_filename)?;
        println!("\n✓ Results saved to: {}", results_filename);

        // Also save a summary to text file
        let summary_filename = "cp2k_dft_summary.txt";
        save_summary(&results, summary_filename)?;
        println!("✓ Summary saved to: {}", summary_filename);

        println!("\n{}", "=".repeat(65));
        println!("SUCCESS: All results extracted and saved!");
        println!("{}", "=".repeat(65));
    }

    // 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);

    #[cfg(not(feature = "extended"))]
    {
        println!("ERROR: Extended features not enabled!");
        println!("Build with --features extended to use this example.");
    }

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

    println!("=================================================================");
    println!("  Complete! Check the output files for results.");
    println!("=================================================================");

    Ok(())
}

// Helper enum for different result types
enum ResultData {
    Scalar(f64),
    Array1D(Array1<f64>),
    Array2D(Array2<f64>),
}

// Save results to text file (human-readable format)
fn save_results(
    results: &HashMap<String, ResultData>,
    filename: &str,
) -> Result<(), Box<dyn std::error::Error>> {
    use std::io::BufWriter;

    let file = File::create(filename)?;
    let mut writer = BufWriter::new(file);

    // Write NPZ header
    writeln!(writer, "# NumPy-compatible NPZ file")?;
    writeln!(writer, "# Generated by CP2K-RS extended interface")?;
    writeln!(writer, "# ")?;

    // For simplicity, save as text format that can be loaded by NumPy
    // In production, use the `npyz` or `ndarray-npy` crate for proper NPZ format

    for (key, data) in results.iter() {
        match data {
            ResultData::Scalar(val) => {
                writeln!(writer, "# {}: scalar", key)?;
                writeln!(writer, "{}: {}", key, val)?;
            }
            ResultData::Array1D(arr) => {
                writeln!(writer, "# {}: array1d, shape = {:?}", key, arr.shape())?;
                write!(writer, "{}: [", key)?;
                for (i, val) in arr.iter().enumerate() {
                    if i > 0 {
                        write!(writer, ", ")?;
                    }
                    write!(writer, "{}", val)?;
                }
                writeln!(writer, "]")?;
            }
            ResultData::Array2D(arr) => {
                writeln!(writer, "# {}: array2d, shape = {:?}", key, arr.shape())?;
                writeln!(writer, "{}: [", key)?;
                for i in 0..arr.shape()[0] {
                    write!(writer, "  [")?;
                    for j in 0..arr.shape()[1] {
                        if j > 0 {
                            write!(writer, ", ")?;
                        }
                        write!(writer, "{}", arr[[i, j]])?;
                    }
                    writeln!(writer, "],")?;
                }
                writeln!(writer, "]")?;
            }
        }
    }

    println!("  Data saved in text format (compatible with NumPy)");
    println!("  Note: For binary NPZ, use 'ndarray-npy' crate in production");

    Ok(())
}

// Save human-readable summary
fn save_summary(
    results: &HashMap<String, ResultData>,
    filename: &str,
) -> Result<(), Box<dyn std::error::Error>> {
    let mut file = File::create(filename)?;

    writeln!(file, "CP2K DFT Calculation - Results Summary")?;
    writeln!(file, "======================================")?;
    writeln!(file)?;

    // Scalars
    writeln!(file, "Scalar Values:")?;
    writeln!(file, "--------------")?;
    for (key, data) in results.iter() {
        if let ResultData::Scalar(val) = data {
            writeln!(file, "  {}: {}", key, val)?;
        }
    }
    writeln!(file)?;

    // Arrays
    writeln!(file, "Array Data:")?;
    writeln!(file, "-----------")?;
    for (key, data) in results.iter() {
        match data {
            ResultData::Array1D(arr) => {
                writeln!(file, "  {}: 1D array, shape = {:?}", key, arr.shape())?;
            }
            ResultData::Array2D(arr) => {
                writeln!(file, "  {}: 2D array, shape = {:?}", key, arr.shape())?;
            }
            _ => {}
        }
    }

    Ok(())
}
