chemx-ext 0.3.0

//! Subprocess interface to the external `xtb` binary (GFN2-xTB / GFN-FF).
//!
//! The binary is detected on `PATH` or via the `CHEMX_XTB_PATH` environment
//! variable ([`find_xtb`]); when it is absent every entry point returns
//! [`ExtError::BinaryNotFound`] with install guidance. The input writers
//! ([`xtb_args`], [`crate::xyz::write_xyz`]) and the output parsers
//! ([`parse_turbomole_gradient`], [`parse_xtbopt_xyz`], [`parse_energy_stdout`],
//! [`parse_json_energy`]) are pure functions, unit-tested against the vendored
//! fixtures in this file (see `PROVENANCE.md` for the xtb version and format
//! sources). The actual subprocess run ([`run`]) is exercised only by an
//! `#[ignore]`d, binary-gated test.

use std::path::{Path, PathBuf};

use chemx_core::Molecule;

use crate::ExtError;
use crate::xyz::write_xyz;

/// Override environment variable for the xtb executable path.
pub const XTB_PATH_ENV: &str = "CHEMX_XTB_PATH";

/// xtb method selection.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum XtbMethod {
    /// GFN2-xTB (`--gfn 2`): the default semiempirical tight-binding method.
    Gfn2Xtb,
    /// GFN-FF (`--gfnff`): the generic force field.
    GfnFf,
}

impl XtbMethod {
    /// The chemx CLI keyword (`gfn2-xtb` / `gfn-ff`).
    pub fn keyword(self) -> &'static str {
        match self {
            XtbMethod::Gfn2Xtb => "gfn2-xtb",
            XtbMethod::GfnFf => "gfn-ff",
        }
    }

    /// Parse the chemx CLI keyword.
    pub fn from_keyword(s: &str) -> Option<Self> {
        match s.to_ascii_lowercase().as_str() {
            "gfn2-xtb" | "gfn2" => Some(XtbMethod::Gfn2Xtb),
            "gfn-ff" | "gfnff" => Some(XtbMethod::GfnFf),
            _ => None,
        }
    }

    /// The xtb command-line flag(s) selecting this method.
    fn method_flags(self) -> Vec<String> {
        match self {
            XtbMethod::Gfn2Xtb => vec!["--gfn".into(), "2".into()],
            XtbMethod::GfnFf => vec!["--gfnff".into()],
        }
    }
}

/// What kind of xtb run to perform.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum XtbRun {
    /// Energy only.
    Energy,
    /// Energy + gradient (`--grad`, writes the TURBOMOLE `gradient` file).
    Gradient,
    /// Geometry optimization (`--opt`, writes `xtbopt.xyz`).
    Opt,
}

/// Specification of an xtb calculation (independent of the input geometry,
/// which is written separately as an XYZ file).
#[derive(Debug, Clone)]
pub struct XtbInput {
    /// Method.
    pub method: XtbMethod,
    /// Net molecular charge (`--chrg`).
    pub charge: i32,
    /// Number of **unpaired** electrons (`--uhf`), i.e. multiplicity − 1.
    pub n_unpaired: u32,
    /// Optional ALPB implicit solvent (`--alpb <name>`).
    pub alpb: Option<String>,
}

impl XtbInput {
    /// Build an [`XtbInput`] from a molecule's charge/multiplicity.
    pub fn from_molecule(method: XtbMethod, molecule: &Molecule, alpb: Option<String>) -> Self {
        Self {
            method,
            charge: molecule.charge,
            n_unpaired: molecule.multiplicity.saturating_sub(1),
            alpb,
        }
    }
}

/// Assemble the full xtb argument vector for a run on `xyz_file`.
///
/// xtb is invoked as `xtb <xyz_file> --gfn 2 --chrg C --uhf N [--alpb S]
/// [--grad|--opt] --json`. `--json` makes xtb additionally write `xtbout.json`
/// (the total energy and properties). Geometry/gradient still come from the
/// `gradient` / `xtbopt.xyz` files, which `--json` does not replace.
pub fn xtb_args(input: &XtbInput, xyz_file: &str, run: XtbRun) -> Vec<String> {
    let mut args = vec![xyz_file.to_string()];
    args.extend(input.method.method_flags());
    args.push("--chrg".into());
    args.push(input.charge.to_string());
    args.push("--uhf".into());
    args.push(input.n_unpaired.to_string());
    if let Some(s) = &input.alpb {
        args.push("--alpb".into());
        args.push(s.clone());
    }
    match run {
        XtbRun::Energy => {}
        XtbRun::Gradient => args.push("--grad".into()),
        XtbRun::Opt => args.push("--opt".into()),
    }
    args.push("--json".into());
    args
}

/// Locate the xtb binary: `CHEMX_XTB_PATH` if set and pointing at an existing
/// file, otherwise `xtb` resolved on `PATH`. Returns [`ExtError::BinaryNotFound`]
/// when neither is available.
pub fn find_xtb() -> Result<PathBuf, ExtError> {
    find_binary(
        "xtb",
        XTB_PATH_ENV,
        "https://github.com/grimme-lab/xtb",
        "xtb",
    )
}

/// Shared binary-detection logic for xtb and crest.
pub(crate) fn find_binary(
    program: &'static str,
    env_var: &'static str,
    install_hint: &'static str,
    conda_pkg: &'static str,
) -> Result<PathBuf, ExtError> {
    if let Some(p) = std::env::var_os(env_var) {
        let path = PathBuf::from(p);
        if path.is_file() {
            return Ok(path);
        }
    }
    if let Some(p) = which_on_path(program) {
        return Ok(p);
    }
    Err(ExtError::BinaryNotFound {
        program,
        env_var,
        install_hint,
        conda_pkg,
    })
}

/// Minimal `which`: scan `PATH` for an executable named `program` (with the
/// platform's executable extensions on Windows).
fn which_on_path(program: &str) -> Option<PathBuf> {
    let path_var = std::env::var_os("PATH")?;
    let exts: Vec<String> = if cfg!(windows) {
        std::env::var("PATHEXT")
            .unwrap_or_else(|_| ".EXE;.BAT;.CMD".into())
            .split(';')
            .map(|s| s.to_string())
            .collect()
    } else {
        vec![String::new()]
    };
    for dir in std::env::split_paths(&path_var) {
        for ext in &exts {
            let cand = dir.join(format!("{program}{ext}"));
            if cand.is_file() {
                return Some(cand);
            }
        }
    }
    None
}

/// Result of an xtb single-point/gradient run.
#[derive(Debug, Clone)]
pub struct XtbResult {
    /// Total energy (hartree).
    pub energy: f64,
    /// Cartesian gradient (hartree/bohr), one `[gx, gy, gz]` per atom — present
    /// for [`XtbRun::Gradient`] runs.
    pub gradient: Option<Vec<[f64; 3]>>,
    /// Optimized geometry (a new molecule) — present for [`XtbRun::Opt`] runs.
    pub optimized: Option<Molecule>,
}

/// Run xtb on `molecule` in a fresh working directory `workdir`. Requires the
/// binary; returns [`ExtError::BinaryNotFound`] if absent.
///
/// This is the only function that spawns a process; it is covered by an
/// `#[ignore]`d, binary-gated test.
pub fn run(
    molecule: &Molecule,
    input: &XtbInput,
    run_kind: XtbRun,
    workdir: &Path,
) -> Result<XtbResult, ExtError> {
    let exe = find_xtb()?;
    std::fs::create_dir_all(workdir).map_err(|e| ExtError::io("creating xtb workdir", e))?;
    let xyz_name = "chemx_xtb_input.xyz";
    let xyz_path = workdir.join(xyz_name);
    std::fs::write(&xyz_path, write_xyz(molecule, "chemx")?)
        .map_err(|e| ExtError::io("writing xtb input xyz", e))?;

    let args = xtb_args(input, xyz_name, run_kind);
    let output = std::process::Command::new(&exe)
        .args(&args)
        .current_dir(workdir)
        .output()
        .map_err(|e| ExtError::io(format!("spawning {}", exe.display()), e))?;

    let stdout = String::from_utf8_lossy(&output.stdout).into_owned();
    if !output.status.success() {
        let stderr = String::from_utf8_lossy(&output.stderr);
        let tail: String = stderr
            .chars()
            .rev()
            .take(800)
            .collect::<String>()
            .chars()
            .rev()
            .collect();
        return Err(ExtError::SubprocessFailed {
            program: "xtb",
            status: output.status.to_string(),
            stderr_tail: if tail.is_empty() { stdout } else { tail },
        });
    }

    // Energy: prefer xtbout.json, fall back to stdout.
    let energy = match std::fs::read_to_string(workdir.join("xtbout.json")) {
        Ok(j) => parse_json_energy(&j)?,
        Err(_) => parse_energy_stdout(&stdout)?,
    };

    let gradient = if run_kind == XtbRun::Gradient {
        let grad_path = workdir.join("gradient");
        let text = std::fs::read_to_string(&grad_path).map_err(|_| ExtError::MissingOutput {
            program: "xtb",
            path: grad_path,
        })?;
        Some(parse_turbomole_gradient(&text, molecule.len())?.gradient)
    } else {
        None
    };

    let optimized = if run_kind == XtbRun::Opt {
        let opt_path = workdir.join("xtbopt.xyz");
        let text = std::fs::read_to_string(&opt_path).map_err(|_| ExtError::MissingOutput {
            program: "xtb",
            path: opt_path,
        })?;
        let (mol, _e) = parse_xtbopt_xyz(&text, molecule.charge, molecule.multiplicity)?;
        Some(mol)
    } else {
        None
    };

    Ok(XtbResult {
        energy,
        gradient,
        optimized,
    })
}

/// Parsed TURBOMOLE `gradient` file.
#[derive(Debug, Clone)]
pub struct TurbomoleGradient {
    /// SCF energy from the `cycle = … SCF energy = …` header (hartree).
    pub energy: f64,
    /// Cartesian gradient (hartree/bohr), one per atom.
    pub gradient: Vec<[f64; 3]>,
}

/// Parse xtb's TURBOMOLE-format `gradient` file (written with `--grad`).
///
/// Layout (xtb 6.x):
/// ```text
/// $grad
///  cycle =      1    SCF energy =    -5.0703873… |dE/dxyz| =  0.0001…
///   <N coordinate lines: x y z element>   (coordinates in BOHR)
///   <N gradient lines: gx gy gz>          (hartree/bohr, Fortran D/E exponent)
/// $end
/// ```
/// `n_atoms` is needed to split coordinate lines from gradient lines.
pub fn parse_turbomole_gradient(text: &str, n_atoms: usize) -> Result<TurbomoleGradient, ExtError> {
    let perr = |m: String| ExtError::Parse {
        what: "TURBOMOLE gradient file",
        message: m,
    };
    let mut lines = text.lines();
    // Find $grad.
    let mut saw_grad = false;
    for l in lines.by_ref() {
        if l.trim_start().starts_with("$grad") {
            saw_grad = true;
            break;
        }
    }
    if !saw_grad {
        return Err(perr("no $grad block".into()));
    }
    // Header line with SCF energy.
    let header = lines
        .next()
        .ok_or_else(|| perr("missing cycle header".into()))?;
    let energy = parse_scf_energy_from_header(header).ok_or_else(|| {
        perr(format!(
            "could not parse SCF energy from header: {header:?}"
        ))
    })?;

    // Next n_atoms lines: coordinates (skip). Following n_atoms lines: gradient.
    let body: Vec<&str> = lines
        .take_while(|l| !l.trim_start().starts_with("$end"))
        .filter(|l| !l.trim().is_empty())
        .collect();
    if body.len() < 2 * n_atoms {
        return Err(perr(format!(
            "expected {} coordinate+gradient lines, found {}",
            2 * n_atoms,
            body.len()
        )));
    }
    let mut gradient = Vec::with_capacity(n_atoms);
    for line in &body[n_atoms..2 * n_atoms] {
        let vals =
            parse_three_floats(line).ok_or_else(|| perr(format!("bad gradient line: {line:?}")))?;
        gradient.push(vals);
    }
    Ok(TurbomoleGradient { energy, gradient })
}

/// Pull the `SCF energy = <float>` value out of a TURBOMOLE gradient header.
fn parse_scf_energy_from_header(header: &str) -> Option<f64> {
    let idx = header.find("SCF energy")?;
    let after = &header[idx + "SCF energy".len()..];
    let after = after.trim_start().strip_prefix('=')?;
    // Take the first whitespace-delimited token, normalizing Fortran D exponents.
    let tok = after.split_whitespace().next()?;
    parse_fortran_float(tok)
}

/// Parse three floats (Fortran `D`/`E` exponents accepted) from a line,
/// ignoring any trailing element symbol.
fn parse_three_floats(line: &str) -> Option<[f64; 3]> {
    let mut it = line.split_whitespace();
    let x = parse_fortran_float(it.next()?)?;
    let y = parse_fortran_float(it.next()?)?;
    let z = parse_fortran_float(it.next()?)?;
    Some([x, y, z])
}

/// Parse a float that may use a Fortran `D` exponent (e.g. `1.23D-02`).
fn parse_fortran_float(tok: &str) -> Option<f64> {
    tok.replace(['D', 'd'], "E").parse().ok()
}

/// Parse `xtbopt.xyz` (written by `--opt`): a standard XYZ whose comment line
/// carries `energy: <Eh> gnorm: <…> xtb: <version>`. Returns the optimized
/// molecule (with the supplied charge/multiplicity) and the energy if present.
pub fn parse_xtbopt_xyz(
    text: &str,
    charge: i32,
    multiplicity: u32,
) -> Result<(Molecule, Option<f64>), ExtError> {
    let mol = Molecule::from_xyz(text)
        .map_err(|e| ExtError::Parse {
            what: "xtbopt.xyz",
            message: e.to_string(),
        })?
        .with_charge(charge)
        .with_multiplicity(multiplicity);
    // The comment line is the second line.
    let comment = text.lines().nth(1).unwrap_or("");
    let energy = comment
        .find("energy:")
        .and_then(|i| comment[i + "energy:".len()..].split_whitespace().next())
        .and_then(parse_fortran_float);
    Ok((mol, energy))
}

/// Parse the total energy from xtb's `xtbout.json` (`--json`), key
/// `"total energy"` (hartree).
pub fn parse_json_energy(json: &str) -> Result<f64, ExtError> {
    let v: serde_json::Value = serde_json::from_str(json).map_err(|e| ExtError::Parse {
        what: "xtbout.json",
        message: e.to_string(),
    })?;
    v.get("total energy")
        .and_then(|x| x.as_f64())
        .ok_or_else(|| ExtError::Parse {
            what: "xtbout.json",
            message: "missing numeric \"total energy\" key".into(),
        })
}

/// Parse the total energy from xtb stdout: the boxed
/// `| TOTAL ENERGY   <Eh> Eh |` line.
pub fn parse_energy_stdout(stdout: &str) -> Result<f64, ExtError> {
    for line in stdout.lines() {
        let t = line.trim_start_matches(['|', ' ']);
        if let Some(rest) = t.strip_prefix("TOTAL ENERGY") {
            // rest = "   -5.070387306651 Eh   |"
            for tok in rest.split_whitespace() {
                if let Some(f) = parse_fortran_float(tok) {
                    return Ok(f);
                }
            }
        }
    }
    Err(ExtError::Parse {
        what: "xtb stdout",
        message: "no '| TOTAL ENERGY … Eh |' line found".into(),
    })
}

#[cfg(test)]
mod tests {
    use super::*;

    // ── Vendored fixtures (xtb 6.6.1; see PROVENANCE.md) ────────────────────

    /// `gradient` file from `xtb h2.xyz --gfn 2 --grad` (H₂, GFN2-xTB).
    /// Format reproduced from xtb's TURBOMOLE writer (src/io/writer/*).
    const GRADIENT_FIXTURE: &str = "\
$grad
 cycle =      1    SCF energy =    -1.0599932109   |dE/dxyz| =  0.012345
      0.00000000000000      0.00000000000000     -0.69853235004775      h
      0.00000000000000      0.00000000000000      0.69853235004775      h
      0.00000000000000      0.00000000000000     -0.12000000000000E-01
      0.00000000000000      0.00000000000000      0.12000000000000E-01
$end
";

    /// `xtbopt.xyz` from `xtb h2.xyz --gfn 2 --opt`.
    const XTBOPT_FIXTURE: &str = "\
2
 energy: -1.067568856769 gnorm: 0.000034567890 xtb: 6.6.1 (8d0f1dd)
H         0.00000000000000    0.00000000000000   -0.37094772930000
H         0.00000000000000    0.00000000000000    0.37094772930000
";

    /// Minimal `xtbout.json` (the keys chemx reads).
    const XTBOUT_JSON_FIXTURE: &str = r#"{
  "total energy": -5.070387306651,
  "HOMO-LUMO gap/eV": 13.486,
  "electronic energy": -5.184,
  "dipole": [0.0, 0.0, 0.0]
}"#;

    /// Boxed total-energy line from xtb stdout.
    const STDOUT_FIXTURE: &str = "\
          :::::::::::::::::::::::::::::::::::::::::::::::::::::
          ::                     SUMMARY                     ::
          :::::::::::::::::::::::::::::::::::::::::::::::::::::
          | TOTAL ENERGY               -5.070387306651 Eh   |
          | GRADIENT NORM               0.000339208090 Eh/α |
normal termination of xtb
";

    #[test]
    fn args_gfn2_with_charge_uhf_alpb() {
        let input = XtbInput {
            method: XtbMethod::Gfn2Xtb,
            charge: -1,
            n_unpaired: 2,
            alpb: Some("water".into()),
        };
        let args = xtb_args(&input, "mol.xyz", XtbRun::Gradient);
        assert_eq!(
            args,
            vec![
                "mol.xyz", "--gfn", "2", "--chrg", "-1", "--uhf", "2", "--alpb", "water", "--grad",
                "--json"
            ]
        );
    }

    #[test]
    fn args_gfnff_opt() {
        let input = XtbInput {
            method: XtbMethod::GfnFf,
            charge: 0,
            n_unpaired: 0,
            alpb: None,
        };
        let args = xtb_args(&input, "m.xyz", XtbRun::Opt);
        assert_eq!(
            args,
            vec![
                "m.xyz", "--gfnff", "--chrg", "0", "--uhf", "0", "--opt", "--json"
            ]
        );
    }

    #[test]
    fn method_keyword_round_trip() {
        assert_eq!(
            XtbMethod::from_keyword("gfn2-xtb"),
            Some(XtbMethod::Gfn2Xtb)
        );
        assert_eq!(XtbMethod::from_keyword("gfnff"), Some(XtbMethod::GfnFf));
        assert_eq!(XtbMethod::from_keyword("b3lyp"), None);
    }

    #[test]
    fn parse_gradient_fixture() {
        let g = parse_turbomole_gradient(GRADIENT_FIXTURE, 2).unwrap();
        assert!((g.energy - (-1.0599932109)).abs() < 1e-10);
        assert_eq!(g.gradient.len(), 2);
        // Fortran E-exponent gradient values.
        assert!((g.gradient[0][2] - (-0.012)).abs() < 1e-12);
        assert!((g.gradient[1][2] - 0.012).abs() < 1e-12);
        assert!(g.gradient[0][0].abs() < 1e-15);
    }

    #[test]
    fn parse_gradient_wrong_atom_count_errors() {
        assert!(parse_turbomole_gradient(GRADIENT_FIXTURE, 3).is_err());
        assert!(parse_turbomole_gradient("no block here", 1).is_err());
    }

    #[test]
    fn parse_xtbopt_fixture() {
        let (mol, e) = parse_xtbopt_xyz(XTBOPT_FIXTURE, 0, 1).unwrap();
        assert_eq!(mol.len(), 2);
        assert_eq!(mol.atoms[0].element.symbol(), "H");
        assert!((e.unwrap() - (-1.067568856769)).abs() < 1e-12);
    }

    #[test]
    fn parse_json_fixture() {
        let e = parse_json_energy(XTBOUT_JSON_FIXTURE).unwrap();
        assert!((e - (-5.070387306651)).abs() < 1e-12);
        assert!(parse_json_energy("{}").is_err());
    }

    #[test]
    fn parse_stdout_fixture() {
        let e = parse_energy_stdout(STDOUT_FIXTURE).unwrap();
        assert!((e - (-5.070387306651)).abs() < 1e-12);
        assert!(parse_energy_stdout("nothing here").is_err());
    }

    #[test]
    fn fortran_float_d_exponent() {
        assert!((parse_fortran_float("1.5D-02").unwrap() - 0.015).abs() < 1e-15);
        assert!((parse_fortran_float("-3.0E0").unwrap() + 3.0).abs() < 1e-15);
    }

    /// Real-subprocess smoke test, gated on binary detection AND `#[ignore]`:
    /// runs an actual GFN2-xTB single point on H₂ only when xtb is installed.
    /// Skips cleanly (no failure) when the binary is absent, so the ignored
    /// tier is safe on machines without xtb.
    #[test]
    #[ignore = "requires the external xtb binary; run with --include-ignored"]
    fn real_xtb_single_point_h2() {
        if find_xtb().is_err() {
            eprintln!("xtb not found — skipping real-subprocess test");
            return;
        }
        let mol = Molecule::from_xyz("2\nh2\nH 0 0 0\nH 0 0 0.74\n").unwrap();
        let input = XtbInput::from_molecule(XtbMethod::Gfn2Xtb, &mol, None);
        let workdir = std::env::temp_dir().join("chemx_xtb_realtest");
        let res = run(&mol, &input, XtbRun::Gradient, &workdir).expect("xtb run");
        assert!(res.energy < 0.0, "H2 GFN2 energy should be negative");
        assert!(res.gradient.is_some());
    }

    #[test]
    fn find_xtb_absent_gives_named_error() {
        // With a bogus override and (almost certainly) no xtb on the test PATH,
        // the error is the named BinaryNotFound. If a machine *does* have xtb on
        // PATH this returns Ok — accept either, but assert the error shape when
        // it is an error.
        unsafe {
            std::env::set_var(XTB_PATH_ENV, "/nonexistent/xtb/binary/xyzzy");
        }
        match find_xtb() {
            Err(ExtError::BinaryNotFound {
                program, env_var, ..
            }) => {
                assert_eq!(program, "xtb");
                assert_eq!(env_var, XTB_PATH_ENV);
            }
            Err(other) => panic!("unexpected error: {other}"),
            Ok(_) => { /* xtb is on PATH on this machine */ }
        }
        unsafe {
            std::env::remove_var(XTB_PATH_ENV);
        }
    }
}