src/chemgraph/tools/ase_core.py

"""Core simulation functions — the single source of truth.

Every callable here is a plain Python function (no LangChain ``@tool``,
no MCP ``@mcp.tool``, no Parsl ``@python_app``).  Framework-specific
wrappers in ``ase_tools.py``, ``mcp_tools.py``, and ``parsl_tools.py``
simply delegate to these functions.
"""

from __future__ import annotations

import glob
import hashlib
import json
import os
import shutil
import tempfile
import time
from pathlib import Path
from typing import List, Optional

import numpy as np

from chemgraph.schemas.atomsdata import AtomsData
from chemgraph.schemas.ase_input import ASEInputSchema, ASEOutputSchema


# ---------------------------------------------------------------------------
# Path helpers
# ---------------------------------------------------------------------------

def _resolve_path(path: str) -> str:
    """Resolve a path relative to ``CHEMGRAPH_LOG_DIR`` when appropriate.

    Parameters
    ----------
    path : str
        Absolute or relative file path.

    Returns
    -------
    str
        Resolved path.
    """
    log_dir = os.environ.get("CHEMGRAPH_LOG_DIR")
    if log_dir and not os.path.isabs(path):
        os.makedirs(log_dir, exist_ok=True)
        return os.path.join(log_dir, path)
    return path


def _ase_cache_enabled() -> bool:
    """Return whether run_ase result caching is enabled."""
    value = os.environ.get("CHEMGRAPH_ASE_CACHE", "1").strip().lower()
    return value not in {"0", "false", "no", "off"}


def _ase_cache_dir() -> str:
    """Return the cache directory for ASE result payloads."""
    configured = os.environ.get("CHEMGRAPH_ASE_CACHE_DIR")
    if configured:
        return configured
    log_dir = os.environ.get("CHEMGRAPH_LOG_DIR") or os.getcwd()
    log_path = Path(log_dir)
    if log_path.name.startswith("query_"):
        log_path = log_path.parent
    return os.path.join(str(log_path), ".ase_cache")


def _file_sha256(path: str) -> str:
    """Return a SHA256 digest for a file."""
    digest = hashlib.sha256()
    with open(path, "rb") as rf:
        for chunk in iter(lambda: rf.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def _ase_cache_key(params: ASEInputSchema, input_structure_file: str) -> str:
    """Build a content-based cache key for an ASE calculation."""
    try:
        payload = params.model_dump(mode="json")
    except Exception:
        payload = params.model_dump()
    payload.pop("input_structure_file", None)
    payload.pop("output_results_file", None)
    payload["input_structure_sha256"] = _file_sha256(input_structure_file)
    encoded = json.dumps(payload, sort_keys=True, default=str).encode("utf-8")
    return hashlib.sha256(encoded).hexdigest()


def _ase_cache_paths(cache_key: str) -> tuple[str, str]:
    cache_dir = _ase_cache_dir()
    return (
        os.path.join(cache_dir, f"{cache_key}.return.json"),
        os.path.join(cache_dir, f"{cache_key}.output.json"),
    )


def _load_ase_cache(cache_key: str, output_results_file: str) -> Optional[dict]:
    """Load a cached ASE result and materialize the requested output JSON."""
    if not _ase_cache_enabled():
        return None
    return_path, output_path = _ase_cache_paths(cache_key)
    if not (os.path.isfile(return_path) and os.path.isfile(output_path)):
        return None
    os.makedirs(os.path.dirname(output_results_file) or ".", exist_ok=True)
    shutil.copy2(output_path, output_results_file)
    with open(return_path, "r", encoding="utf-8") as rf:
        cached = json.load(rf)
    cached = _compact_ase_return_payload(cached)
    cached["cache_hit"] = True
    if isinstance(cached.get("message"), str):
        cached["message"] = f"{cached['message']} Loaded from ChemGraph ASE cache."
    return cached


def _store_ase_cache(
    cache_key: str,
    output_results_file: str,
    result: dict,
) -> dict:
    """Persist a successful ASE result in the local cache."""
    if not _ase_cache_enabled() or result.get("status") != "success":
        return result
    result = _compact_ase_return_payload(result)
    if not os.path.isfile(output_results_file):
        return result
    return_path, output_path = _ase_cache_paths(cache_key)
    os.makedirs(os.path.dirname(return_path), exist_ok=True)
    result_to_store = dict(result)
    result_to_store.pop("cache_hit", None)
    with open(return_path, "w", encoding="utf-8") as wf:
        json.dump(result_to_store, wf, indent=2, default=str)
    shutil.copy2(output_results_file, output_path)
    return result


def _compact_ase_return_payload(result: dict) -> dict:
    """Keep tool-return payloads small while preserving full output files."""
    if not isinstance(result, dict):
        return result
    if not isinstance(result.get("result"), dict):
        return result
    compact = dict(result)
    compact_result = dict(result["result"])
    compact_result.pop("ir_data", None)
    compact["result"] = compact_result
    return compact


# ---------------------------------------------------------------------------
# AtomsData <-> ASE Atoms conversions
# ---------------------------------------------------------------------------

def atoms_to_atomsdata(atoms) -> AtomsData:
    """Convert an ASE ``Atoms`` object to :class:`AtomsData`.

    Parameters
    ----------
    atoms : ase.Atoms
        ASE Atoms object.

    Returns
    -------
    AtomsData
    """
    return AtomsData(
        numbers=atoms.numbers.tolist(),
        positions=atoms.positions.tolist(),
        cell=atoms.cell.tolist(),
        pbc=atoms.pbc.tolist(),
    )


def atomsdata_to_atoms(atomsdata: AtomsData):
    """Convert :class:`AtomsData` to an ASE ``Atoms`` object.

    Parameters
    ----------
    atomsdata : AtomsData

    Returns
    -------
    ase.Atoms
    """
    from ase import Atoms

    return Atoms(
        numbers=atomsdata.numbers,
        positions=atomsdata.positions,
        cell=atomsdata.cell,
        pbc=atomsdata.pbc,
    )


# ---------------------------------------------------------------------------
# Molecular property helpers
# ---------------------------------------------------------------------------

def is_linear_molecule(atomsdata: AtomsData, tol: float = 1e-3) -> bool:
    """Determine whether a molecule is linear.

    Parameters
    ----------
    atomsdata : AtomsData
        Molecular structure.
    tol : float, optional
        Tolerance for the second singular value ratio, by default 1e-3.

    Returns
    -------
    bool
        ``True`` if the molecule is linear.
    """
    coords = np.array(atomsdata.positions)
    centered = coords - np.mean(coords, axis=0)
    _, s, _ = np.linalg.svd(centered)
    if s[0] == 0:
        return False  # degenerate — all atoms at one point
    return (s[1] / s[0]) < tol


def get_symmetry_number(atomsdata: AtomsData) -> int:
    """Return the rotational symmetry number using Pymatgen.

    Parameters
    ----------
    atomsdata : AtomsData

    Returns
    -------
    int
    """
    from pymatgen.symmetry.analyzer import PointGroupAnalyzer
    from ase import Atoms
    from pymatgen.io.ase import AseAtomsAdaptor

    atoms = Atoms(
        numbers=atomsdata.numbers,
        positions=atomsdata.positions,
        cell=atomsdata.cell,
        pbc=atomsdata.pbc,
    )
    aaa = AseAtomsAdaptor()
    molecule = aaa.get_molecule(atoms)
    pga = PointGroupAnalyzer(molecule)
    return pga.get_rotational_symmetry_number()


# ---------------------------------------------------------------------------
# Calculator loading
# ---------------------------------------------------------------------------

def load_calculator(calculator: dict) -> tuple[object, dict, object]:
    """Instantiate an ASE calculator from a config dictionary.

    Parameters
    ----------
    calculator : dict
        Must contain a ``"calculator_type"`` key.

    Returns
    -------
    tuple[object, dict, object]
        ``(ase_calculator, extra_info, calc_schema_instance)``

    Raises
    ------
    ValueError
        If the calculator type is unsupported.
    """
    calc_type = calculator["calculator_type"].lower()

    if "emt" in calc_type:
        from chemgraph.schemas.calculators.emt_calc import EMTCalc
        calc = EMTCalc(**calculator)
    elif "tblite" in calc_type or "xtb" in calc_type:
        from chemgraph.schemas.calculators.tblite_calc import TBLiteCalc
        calc = TBLiteCalc(**calculator)
    elif "orca" in calc_type:
        from chemgraph.schemas.calculators.orca_calc import OrcaCalc
        calc = OrcaCalc(**calculator)
    elif "nwchem" in calc_type:
        from chemgraph.schemas.calculators.nwchem_calc import NWChemCalc
        calc = NWChemCalc(**calculator)
    elif "fairchem" in calc_type:
        from chemgraph.schemas.calculators.fairchem_calc import FAIRChemCalc
        calc = FAIRChemCalc(**calculator)
    elif "mace" in calc_type:
        from chemgraph.schemas.calculators.mace_calc import MaceCalc
        calc = MaceCalc(**calculator)
    elif "aimnet2" in calc_type:
        from chemgraph.schemas.calculators.aimnet2_calc import AIMNET2Calc
        calc = AIMNET2Calc(**calculator)
    else:
        raise ValueError(
            f"Unsupported calculator: {calculator}. "
            "Available calculators are EMT, TBLite (GFN2-xTB, GFN1-xTB), "
            "Orca, NWChem, FAIRChem, MACE, or AIMNET2."
        )

    extra_info: dict = {}
    if hasattr(calc, "get_atoms_properties"):
        extra_info = calc.get_atoms_properties()

    return calc.get_calculator(), extra_info, calc


# ---------------------------------------------------------------------------
# Misc helpers (kept for backward compat / UI)
# ---------------------------------------------------------------------------

def extract_ase_atoms_from_tool_result(tool_result: dict):
    """Extract ``(atomic_numbers, positions)`` from a tool-result dict.

    Returns ``(None, None)`` if extraction fails.

    Parameters
    ----------
    tool_result : dict
        Tool result that may contain atom numbers and positions.

    Returns
    -------
    tuple
        ``(atomic_numbers, positions)`` or ``(None, None)``.
    """
    for keyset in ({"numbers", "positions"}, {"atomic_numbers", "positions"}):
        if keyset.issubset(tool_result.keys()):
            return tool_result[keyset.pop()], tool_result["positions"]

    if "atoms" in tool_result:
        atoms_data = tool_result["atoms"]
        if {"numbers", "positions"}.issubset(atoms_data):
            return atoms_data["numbers"], atoms_data["positions"]

    return None, None


def create_ase_atoms(atomic_numbers, positions):
    """Create an ASE ``Atoms`` object from atomic numbers and positions.

    Parameters
    ----------
    atomic_numbers : sequence
        Atomic numbers for each atom.
    positions : sequence
        Cartesian coordinates for each atom.

    Returns
    -------
    ase.Atoms or None
        Constructed atoms object, or ``None`` if construction fails.
    """
    from ase import Atoms

    try:
        return Atoms(numbers=atomic_numbers, positions=positions)
    except Exception as e:
        print(f"Error creating ASE Atoms object: {e}")
        return None


def create_xyz_string(atomic_numbers, positions) -> Optional[str]:
    """Create an XYZ-format string from atomic numbers and positions.

    Parameters
    ----------
    atomic_numbers : sequence
        Atomic numbers for each atom.
    positions : sequence
        Cartesian coordinates for each atom.

    Returns
    -------
    str or None
        XYZ-format structure text, or ``None`` if conversion fails.
    """
    from ase import Atoms

    try:
        atoms = Atoms(numbers=atomic_numbers, positions=positions)
        xyz_lines = [str(len(atoms)), "Generated by ChemGraph"]
        for symbol, pos in zip(atoms.get_chemical_symbols(), atoms.positions):
            xyz_lines.append(
                f"{symbol:2s} {pos[0]:12.6f} {pos[1]:12.6f} {pos[2]:12.6f}"
            )
        return "\n".join(xyz_lines)
    except Exception as e:
        print(f"Error creating XYZ string: {e}")
        return None


# ---------------------------------------------------------------------------
# Unified ASE simulation core
# ---------------------------------------------------------------------------

def run_ase_core(params: ASEInputSchema) -> dict:
    """Run an ASE simulation — the single implementation for all call methods.

    This function implements energy, dipole, optimization, vibrational,
    thermochemistry, and IR calculations.  Framework-specific wrappers
    (LangChain ``@tool``, MCP ``@mcp.tool``, Parsl) delegate here.

    Parameters
    ----------
    params : ASEInputSchema
        Fully validated simulation input.

    Returns
    -------
    dict
        Minimal result payload (status, message, key numbers).
    """
    from ase.io import read
    from ase.optimize import BFGS, LBFGS, GPMin, FIRE, MDMin

    # ---- unpack params ----
    try:
        calculator = params.calculator.model_dump()
    except Exception as e:
        return {
            "status": "failure",
            "error_type": "ValidationError",
            "message": f"Missing calculator parameter for the simulation. Raised exception: {e}",
        }

    start_time = time.time()

    input_structure_file = _resolve_path(params.input_structure_file)
    output_results_file = _resolve_path(params.output_results_file)
    optimizer = params.optimizer
    fmax = params.fmax
    steps = params.steps
    driver = params.driver
    temperature = params.temperature
    pressure = params.pressure
    if driver == "thermo" and temperature is None:
        temperature = 298.15
        params = params.model_copy(update={"temperature": temperature})

    # ---- input validation ----
    if not os.path.isfile(input_structure_file):
        return {
            "status": "failure",
            "error_type": "FileNotFoundError",
            "message": f"Input structure file {input_structure_file} does not exist.",
        }

    if not output_results_file.endswith(".json"):
        return {
            "status": "failure",
            "error_type": "ValueError",
            "message": f"Output results file must end with '.json', got: {params.output_results_file}",
        }

    cache_key = _ase_cache_key(params, input_structure_file)
    cached_result = _load_ase_cache(cache_key, output_results_file)
    if cached_result is not None:
        return cached_result

    calc, system_info, calc_model = load_calculator(calculator)

    if calc is None:
        return {
            "status": "failure",
            "error_type": "ValueError",
            "message": (
                f"Unsupported calculator: {calculator}. Available calculators are "
                "MACE (mace_mp, mace_off, mace_anicc), EMT, TBLite (GFN2-xTB, GFN1-xTB), NWChem and Orca"
            ),
        }

    try:
        atoms = read(input_structure_file)
    except Exception as e:
        return {
            "status": "failure",
            "error_type": type(e).__name__,
            "message": f"Cannot read {input_structure_file} using ASE. Exception from ASE: {e}",
        }

    atoms.info.update(system_info)
    atoms.calc = calc

    # ------------------------------------------------------------------
    # Driver: energy / dipole  (single-point, no optimization)
    # ------------------------------------------------------------------
    if driver in ("energy", "dipole"):
        energy = atoms.get_potential_energy()
        final_structure = atoms_to_atomsdata(atoms)

        dipole: List[Optional[float]] = [None, None, None]
        if driver == "dipole":
            try:
                dipole = [round(x, 4) for x in atoms.get_dipole_moment()]
            except Exception:
                pass

        end_time = time.time()
        wall_time = end_time - start_time

        simulation_output = ASEOutputSchema(
            input_structure_file=input_structure_file,
            converged=True,
            final_structure=final_structure,
            simulation_input=params,
            success=True,
            dipole_value=dipole,
            single_point_energy=energy,
            wall_time=wall_time,
        )
        with open(output_results_file, "w", encoding="utf-8") as wf:
            wf.write(simulation_output.model_dump_json(indent=4))

        if driver == "energy":
            return _store_ase_cache(cache_key, output_results_file, {
                "status": "success",
                "message": f"Simulation completed. Results saved to {os.path.abspath(output_results_file)}",
                "single_point_energy": energy,
                "unit": "eV",
            })
        else:  # dipole
            return _store_ase_cache(cache_key, output_results_file, {
                "status": "success",
                "message": f"Simulation completed. Results saved to {os.path.abspath(output_results_file)}",
                "dipole_moment": dipole,
                "dipole_unit": "e * Angstrom",
            })

    # ------------------------------------------------------------------
    # Drivers that require optimization: opt / vib / thermo / ir
    # ------------------------------------------------------------------
    OPTIMIZERS = {
        "bfgs": BFGS,
        "lbfgs": LBFGS,
        "gpmin": GPMin,
        "fire": FIRE,
        "mdmin": MDMin,
    }
    try:
        optimizer_class = OPTIMIZERS.get(optimizer.lower())
        if optimizer_class is None:
            raise ValueError(f"Unsupported optimizer: {optimizer}")

        if len(atoms) > 1:
            dyn = optimizer_class(atoms)
            converged = dyn.run(fmax=fmax, steps=steps)
        else:
            converged = True

        single_point_energy = float(atoms.get_potential_energy())
        final_structure = AtomsData(
            numbers=atoms.numbers,
            positions=atoms.positions,
            cell=atoms.cell,
            pbc=atoms.pbc,
        )
        thermo_data: dict = {}
        vib_data: dict = {}
        ir_data: dict = {}

        # --------------------------------------------------------------
        # Vibrational / thermo / IR analysis
        # --------------------------------------------------------------
        if driver in {"vib", "thermo", "ir"}:
            from ase.vibrations import Vibrations
            from ase import units

            ir_plot_path: Optional[str] = None
            mol_stem = (
                Path(input_structure_file).stem if input_structure_file else "mol"
            )

            with tempfile.TemporaryDirectory(
                prefix=f"chemgraph_vib_{mol_stem}_"
            ) as tmpdir:
                vib_name = os.path.join(tmpdir, "vib")
                vib = Vibrations(atoms, name=vib_name)
                vib.clean()
                vib.run()

                vib_data = {
                    "energies": [],
                    "energy_unit": "meV",
                    "frequencies": [],
                    "frequency_unit": "cm-1",
                }

                energies = vib.get_energies()

                for _idx, e in enumerate(energies):
                    is_imag = abs(e.imag) > 1e-8
                    e_val = e.imag if is_imag else e.real
                    energy_meV = 1e3 * e_val
                    freq_cm1 = e_val / units.invcm
                    suffix = "i" if is_imag else ""
                    vib_data["energies"].append(f"{energy_meV}{suffix}")
                    vib_data["frequencies"].append(f"{freq_cm1}{suffix}")

                # Write frequencies CSV
                freq_file_path = _resolve_path(f"frequencies_{mol_stem}.csv")
                freq_file = Path(freq_file_path)
                if freq_file.exists():
                    freq_file.unlink()
                with freq_file.open("w", encoding="utf-8") as f:
                    for i, freq in enumerate(vib_data["frequencies"], start=0):
                        f.write(f"{mol_stem}_vib.{i}.traj,{freq}\n")

                # Write normal-mode .traj files, then copy out of tmpdir
                for i in range(len(energies)):
                    vib.write_mode(n=i, kT=units.kB * 300, nimages=30)

                traj_dest_dir = _resolve_path("")
                if traj_dest_dir:
                    os.makedirs(traj_dest_dir, exist_ok=True)
                for traj_file in glob.glob(os.path.join(tmpdir, "vib.*.traj")):
                    dest_name = f"{mol_stem}_{Path(traj_file).name}"
                    dest_path = (
                        os.path.join(traj_dest_dir, dest_name)
                        if traj_dest_dir
                        else dest_name
                    )
                    shutil.copy2(traj_file, dest_path)

                # ---- IR ----
                if driver == "ir":
                    from ase.vibrations import Infrared
                    import matplotlib

                    matplotlib.use("Agg")
                    import matplotlib.pyplot as plt

                    ir_data["spectrum_frequencies"] = []
                    ir_data["spectrum_frequencies_units"] = "cm-1"
                    ir_data["spectrum_intensities"] = []
                    ir_data["spectrum_intensities_units"] = "D/Å^2 amu^-1"

                    ir_name = os.path.join(tmpdir, "ir")
                    ir = Infrared(atoms, name=ir_name)
                    ir.clean()
                    ir.run()

                    IR_SPECTRUM_START = 500
                    IR_SPECTRUM_END = 4000
                    freq_intensity = ir.get_spectrum(
                        start=IR_SPECTRUM_START, end=IR_SPECTRUM_END
                    )
                    mode_frequencies = [float(value) for value in ir.get_frequencies().real]
                    mode_intensities = [float(value) for value in ir.intensities]
                    ir_data["mode_frequencies"] = mode_frequencies
                    ir_data["mode_frequencies_units"] = "cm-1"
                    ir_data["mode_intensities"] = mode_intensities
                    ir_data["mode_intensities_units"] = "D/Å^2 amu^-1"
                    ir_data["spectrum_frequencies"] = [
                        float(value) for value in freq_intensity[0]
                    ]
                    ir_data["spectrum_intensities"] = [
                        float(value) for value in freq_intensity[1]
                    ]
                    fig, ax = plt.subplots()
                    ax.plot(freq_intensity[0], freq_intensity[1])
                    ax.set_xlabel("Frequency (cm⁻¹)")
                    ax.set_ylabel("Intensity (a.u.)")
                    ax.set_title("Infrared Spectrum")
                    ax.grid(True)
                    ir_plot_path = _resolve_path(f"ir_spectrum_{mol_stem}.png")
                    fig.savefig(ir_plot_path, format="png", dpi=300)
                    plt.close(fig)

                    ir_data["IR Plot"] = f"Saved to {os.path.abspath(ir_plot_path)}"
                    ir_spectrum = {
                        "frequency_cm1": [str(value) for value in mode_frequencies],
                        "intensity": [str(value) for value in mode_intensities],
                        "plot": os.path.abspath(ir_plot_path),
                    }
                    ir_data["Normal mode data"] = (
                        f"Normal modes saved as individual .traj files with prefix {mol_stem}_"
                    )

                # ---- Thermochemistry ----
                if driver == "thermo":
                    if len(atoms) == 1:
                        thermo_data = {
                            "enthalpy": single_point_energy,
                            "entropy": 0.0,
                            "gibbs_free_energy": single_point_energy,
                            "unit": "eV",
                        }
                    else:
                        from ase.thermochemistry import IdealGasThermo

                        linear = is_linear_molecule(final_structure)
                        geometry = "linear" if linear else "nonlinear"
                        symmetrynumber = get_symmetry_number(final_structure)

                        # IdealGasThermo expects total spin S; calculators expose
                        # multiplicity (2S+1) via get_multiplicity() when supported.
                        multiplicity = (
                            getattr(calc_model, "get_multiplicity", lambda: None)()
                            or 1
                        )
                        spin_S = (multiplicity - 1) / 2.0

                        thermo = IdealGasThermo(
                            vib_energies=energies,
                            potentialenergy=single_point_energy,
                            atoms=atoms,
                            geometry=geometry,
                            symmetrynumber=symmetrynumber,
                            spin=spin_S,
                        )
                        thermo_data = {
                            "enthalpy": float(
                                thermo.get_enthalpy(temperature=temperature)
                            ),
                            "entropy": float(
                                thermo.get_entropy(
                                    temperature=temperature, pressure=pressure
                                )
                            ),
                            "gibbs_free_energy": float(
                                thermo.get_gibbs_energy(
                                    temperature=temperature, pressure=pressure
                                )
                            ),
                            "unit": "eV",
                        }

        # ---- serialise full output ----
        end_time = time.time()
        wall_time = end_time - start_time

        simulation_output = ASEOutputSchema(
            input_structure_file=input_structure_file,
            converged=converged,
            final_structure=final_structure,
            simulation_input=params,
            vibrational_frequencies=vib_data,
            thermochemistry=thermo_data,
            success=True,
            ir_data=ir_data,
            single_point_energy=single_point_energy,
            wall_time=wall_time,
        )
        with open(output_results_file, "w", encoding="utf-8") as wf:
            wf.write(simulation_output.model_dump_json(indent=4))

        # ---- minimal return payload ----
        abs_output = os.path.abspath(output_results_file)
        if driver == "opt":
            return _store_ase_cache(cache_key, output_results_file, {
                "status": "success",
                "message": f"Simulation completed. Results saved to {abs_output}",
                "single_point_energy": single_point_energy,
                "unit": "eV",
            })
        elif driver == "vib":
            return _store_ase_cache(cache_key, output_results_file, {
                "status": "success",
                "result": {"vibrational_frequencies": vib_data},
                "message": (
                    "Vibrational analysis completed; frequencies returned. "
                    f"Full results (structure, vibrations and metadata) saved to {abs_output}."
                ),
            })
        elif driver == "thermo":
            return _store_ase_cache(cache_key, output_results_file, {
                "status": "success",
                "result": {"thermochemistry": thermo_data},
                "message": (
                    "Thermochemistry computed and returned. "
                    f"Full results (structure, vibrations, thermochemistry and metadata) saved to {abs_output}"
                ),
            })
        elif driver == "ir":
            return _store_ase_cache(cache_key, output_results_file, {
                "status": "success",
                "result": {
                    "vibrational_frequencies": vib_data,
                    "ir_spectrum": ir_spectrum,
                },
                "message": (
                    "Infrared computed and returned. "
                    f"Full results (structure, vibrations, IR arrays and metadata) saved to {abs_output}. "
                    f"IR plot saved to {os.path.abspath(ir_plot_path) if ir_plot_path else 'N/A'}. "
                    "Normal modes saved as individual .traj files"
                ),
            })

    except Exception as e:
        return {
            "status": "failure",
            "error_type": type(e).__name__,
            "message": str(e),
        }


# ---------------------------------------------------------------------------
# JSON result loader
# ---------------------------------------------------------------------------


def extract_output_json_core(json_file: str) -> dict:
    """Load simulation results from a JSON file produced by ``run_ase_core``.

    Parameters
    ----------
    json_file : str
        Path to the JSON file containing ASE simulation results.

    Returns
    -------
    dict
        Parsed results from the JSON file as a Python dictionary.

    Raises
    ------
    FileNotFoundError
        If the specified file does not exist.
    json.JSONDecodeError
        If the file is not valid JSON.
    """
    with open(json_file, "r", encoding="utf-8") as f:
        data = json.load(f)
    return data