src/chemgraph/schemas/ase_input.py

import importlib.util
import json
import os
import shutil
from pydantic import BaseModel, Field, model_validator, field_validator
from typing import Union, Optional, Any, List, Type
from chemgraph.schemas.atomsdata import AtomsData

from chemgraph.schemas.calculators.emt_calc import EMTCalc
from chemgraph.schemas.calculators.nwchem_calc import NWChemCalc
from chemgraph.schemas.calculators.orca_calc import OrcaCalc
from chemgraph.schemas.calculators.fairchem_calc import FAIRChemCalc
from chemgraph.schemas.calculators.mace_calc import MaceCalc
from chemgraph.schemas.calculators.tblite_calc import TBLiteCalc
from chemgraph.schemas.calculators.aimnet2_calc import AIMNET2Calc

# Gate optional calculators on whether their engine package is installed.
# Schema classes are always importable (internal to ChemGraph), so we must
# probe the underlying engine with importlib.util.find_spec().
# find_spec() can raise ModuleNotFoundError for sub-packages when the parent
# package is missing, so we guard with try/except.

def _engine_available(module_name: str) -> bool:
    """Return whether a Python calculator engine module is importable.

    Parameters
    ----------
    module_name : str
        Module name passed to ``importlib.util.find_spec``.

    Returns
    -------
    bool
        ``True`` when the module can be found.
    """
    try:
        return importlib.util.find_spec(module_name) is not None
    except (ImportError, ModuleNotFoundError):
        return False


def _command_available(command_name: str, env_var_name: str) -> bool:
    """Return whether a calculator command is configured or on ``PATH``.

    Parameters
    ----------
    command_name : str
        Executable name to locate.
    env_var_name : str
        Environment variable that can provide the command.

    Returns
    -------
    bool
        ``True`` when the command is configured or discoverable.
    """
    return bool(os.environ.get(env_var_name)) or shutil.which(command_name) is not None


if not _engine_available("fairchem.core"):
    FAIRChemCalc = None

if not _engine_available("mace"):
    MaceCalc = None

if not _engine_available("tblite"):
    TBLiteCalc = None

if not _engine_available("aimnet2calc"):
    AIMNET2Calc = None

if not _command_available("nwchem", "ASE_NWCHEM_COMMAND"):
    NWChemCalc = None

if not _command_available("orca", "ASE_ORCA_COMMAND"):
    OrcaCalc = None


_CALCULATOR_ALIASES = {
    "xtb": "tbli",
    "gfn1xtb": "tbli",
    "gfn2xtb": "tbli",
}


def _calculator_key(name: str) -> str:
    """Return a normalized calculator lookup key.

    Parameters
    ----------
    name : str
        Calculator name or alias.

    Returns
    -------
    str
        Four-character normalized key used for calculator matching.
    """
    normalized = "".join(ch for ch in name.lower() if ch.isalnum())
    return _CALCULATOR_ALIASES.get(normalized, normalized[:4])


# Define all possible calculator classes
_all_calculator_classes: List[Optional[Type[BaseModel]]] = [
    FAIRChemCalc,
    MaceCalc,
    AIMNET2Calc,
    TBLiteCalc,
    EMTCalc,
    NWChemCalc,
    OrcaCalc,
]

# Filter out unavailable calculators
available_calculator_classes: List[Type[BaseModel]] = [
    calc for calc in _all_calculator_classes if calc
]

# Create a union for type hinting
CalculatorUnion = Union[tuple(available_calculator_classes)]

_calculator_name_items = [calc.__name__ for calc in available_calculator_classes]
_calculator_name_items = [
    "TBLiteCalc (aliases: xTB, GFN1-xTB, GFN2-xTB)"
    if name == "TBLiteCalc"
    else name
    for name in _calculator_name_items
]
_calculator_names = ", ".join(_calculator_name_items)

# Determine default calculator using only calculators detected as available.
default_calculator = available_calculator_classes[0]


def get_available_calculator_names() -> List[str]:
    """Return calculator class names detected as available in this environment."""
    return [calc.__name__ for calc in available_calculator_classes]


def get_default_calculator_name() -> str:
    """Return the default calculator class name selected for this environment."""
    return default_calculator.__name__


def get_calculator_selection_context() -> str:
    """Return prompt text describing available calculators and default choice."""
    return (
        "\n\nCalculator availability detected during ChemGraph initialization:\n"
        f"- Available ASE calculators: {_calculator_names}.\n"
        f"- Default calculator when the user does not specify one: "
        f"{default_calculator.__name__}.\n"
        "- When calling run_ase, choose only from the available calculators above. "
        "If the user requests an unavailable calculator, choose the default "
        "available calculator when that substitution is appropriate; otherwise "
        "ask for clarification or explain that the requested calculator is not "
        "available."
    )


class ASEInputSchema(BaseModel):
    """
    Schema for defining input parameters used in ASE-based molecular simulations.

    Attributes
    ----------
    input_structure_file : str
        Path to the input coordinate file (e.g., CIF, XYZ, POSCAR) containing
        the atomic structure for the simulation.
    output_results_file: str
        Path to a JSON file where simulation results will be saved.
    driver : str
        Specifies the type of simulation to perform. Options:
        - 'energy': Single-point electronic energy calculation.
        - 'opt': Geometry optimization.
        - 'vib': Vibrational frequency analysis.
        - 'ir': Infrared spectrum calculation.
        - 'thermo': Thermochemical property calculation (enthalpy, entropy, Gibbs free energy).
    optimizer : str
        Optimization algorithm for geometry optimization. Options:
        - 'bfgs', 'lbfgs', 'gpmin', 'fire', 'mdmin'.
    calculator : Union[FAIRChemCalc, MaceCalc, NWChemCalc, OrcaCalc, TBLiteCalc, EMTCalc, AIMNET2Calc]
        ASE-compatible calculator used for the simulation. Supported types are determined
        by installed packages and may include:
        - FAIRChem, MACE, NWChem, Orca, TBLite and EMT. The order determines the priority of the calculators.
        - Use MACE or FAIRChem if the calculator is not specified.
    fmax : float
        Force convergence criterion in eV/Å. Optimization stops when all force components fall below this threshold.
    steps : int
        Maximum number of steps for geometry optimization.
    temperature : Optional[float]
        Temperature in Kelvin, required for thermochemical calculations (e.g., when using 'thermo' as the driver).
    pressure : float
        Pressure in Pascal (Pa), used in thermochemistry calculations (default is 1 atm).
    """

    input_structure_file: str = Field(
        description="Path to the input coordinate file (e.g., CIF, XYZ, POSCAR) containing the atomic structure for the simulation."
    )
    output_results_file: str = Field(
        default="output.json",
        description="Path to a JSON file where simulation results will be saved.",
    )
    driver: str = Field(
        default=None,
        description="Specifies the type of simulation to run. Options: 'energy' for electronic energy calculations, 'dipole' for dipole moment calculation, 'opt' for geometry optimization, 'vib' for vibrational frequency analysis, 'ir' for calculating infrared spectrum, and 'thermo' for thermochemical properties (including enthalpy, entropy, and Gibbs free energy). Use 'thermo' when the query involves enthalpy, entropy, or Gibbs free energy calculations.",
    )
    optimizer: str = Field(
        default="bfgs",
        description="The optimization algorithm used for geometry optimization. Options are 'bfgs', 'lbfgs', 'gpmin', 'fire', 'mdmin'",
    )
    calculator: CalculatorUnion = Field(
        default_factory=default_calculator,
        description=f"The ASE calculator to be used. Support {_calculator_names}. Use {default_calculator.__name__} if not specified.",
    )
    fmax: float = Field(
        default=0.01,
        description="The convergence criterion for forces (in eV/Å). Optimization stops when all force components are smaller than this value.",
    )
    steps: int = Field(
        default=1000,
        description="Maximum number of optimization steps. Internally 'vib', 'thermo' and 'ir' run geometry optimization before performing their respective calculations.",
    )
    temperature: Optional[float] = Field(
        default=None,
        description="Temperature for thermochemistry calculations in Kelvin (K).",
    )
    pressure: float = Field(
        default=101325.0,
        description="Pressure for thermochemistry calculations in Pascal (Pa).",
    )

    @model_validator(mode="before")
    @classmethod
    def _validate_calculator_type(cls, data: Any):
        """Validate and coerce the calculator payload.

        Parameters
        ----------
        data : Any
            Raw ASE input payload before Pydantic validation.

        Returns
        -------
        Any
            Payload with calculator converted to an available calculator model.
        """
        if not isinstance(data, dict):
            return data

        calc = data.get("calculator")
        if calc is None:
            calc = default_calculator()
            data["calculator"] = calc

        available_calcs = {
            _calculator_key(c.__name__.removesuffix("Calc")): c
            for c in available_calculator_classes
        }
        available_calc_names = [c.__name__ for c in available_calculator_classes]

        if isinstance(calc, str):
            calc_key = _calculator_key(calc.removesuffix("Calc"))
            if calc_key not in available_calcs:
                raise ValueError(
                    f"Calculator {calc} is not an allowed or available calculator. "
                    f"Available calculators are: {available_calc_names}"
                )
            data["calculator"] = available_calcs[calc_key]()
            return data

        if isinstance(calc, dict):
            calc_name = calc.get("calculator_type")
            if not calc_name:
                raise ValueError("Calculator dictionary must have a 'calculator_type' key.")

            calc_key = _calculator_key(calc_name)
            if calc_key not in available_calcs:
                raise ValueError(
                    f"Calculator {calc_name} is not an allowed or available calculator. "
                    f"Available calculators are: {available_calc_names}"
                )

            init_args = calc.copy()
            init_args.pop("calculator_type", None)
            data["calculator"] = available_calcs[calc_key](**init_args)
            return data

        elif hasattr(calc, "__class__"):
            calc_type_name = calc.__class__.__name__
            calc_key = _calculator_key(calc_type_name.removesuffix("Calc"))
            if calc_key not in available_calcs:
                raise ValueError(
                    f"Calculator {calc_type_name} is not an allowed or available calculator. "
                    f"Available calculators are: {available_calc_names}"
                )
        return data


class ASEOutputSchema(BaseModel):
    """
    Schema for defining outputs from ASE-based molecular simulations.
    """

    input_structure_file: str = Field(
        description=(
            "Path to the input coordinate file (e.g., CIF, XYZ, POSCAR) "
            "containing the initial atomic structure for the simulation."
        )
    )
    converged: bool = Field(
        default=False,
        description="Indicates if the optimization successfully converged.",
    )
    final_structure: AtomsData = Field(description="Final structure.")
    simulation_input: ASEInputSchema = Field(
        description="Simulation input for Atomic Simulation Environment."
    )
    single_point_energy: float = Field(
        default=None, description="Single-point energy/Potential energy"
    )
    energy_unit: str = Field(default="eV", description="The unit of the energy reported.")
    dipole_value: List[Optional[float]] = Field(
        default=[None, None, None],
        description="The value of the dipole moment reported.",
    )
    dipole_unit: str = Field(
        default=" e * angstrom", description="The unit of the dipole moment reported."
    )
    vibrational_frequencies: dict = Field(
        default={},
        description="Vibrational frequencies (in cm-1) and energies (in eV).",
    )
    ir_data: dict = Field(
        default={},
        description="Infrared spectrum related data.",
    )
    thermochemistry: dict = Field(default={}, description="Thermochemistry data in eV.")
    success: bool = Field(
        default=False, description="Indicates if the simulation finished correctly."
    )
    error: str = Field(default="", description="Error captured during the simulation")
    wall_time: float = Field(
        default=None,
        description="Total wall time (in seconds) taken to complete the simulation.",
    )

    @field_validator("vibrational_frequencies", "ir_data", "thermochemistry", mode="before")
    @classmethod
    def _coerce_json_string_to_dict(cls, v: Any) -> dict:
        """Accept dict-like payloads serialized as JSON strings.

        Parameters
        ----------
        v : Any
            Raw field value.

        Returns
        -------
        dict
            Parsed dictionary or empty dictionary.
        """
        if v is None:
            return {}
        if isinstance(v, dict):
            return v
        if isinstance(v, str):
            text = v.strip()
            if not text:
                return {}
            try:
                parsed = json.loads(text)
                return parsed if isinstance(parsed, dict) else {}
            except json.JSONDecodeError:
                return {}
        return {}

    @field_validator("error", mode="before")
    @classmethod
    def _coerce_error_to_string(cls, v: Any) -> str:
        """Allow null/non-string error fields from intermediate tool payloads.

        Parameters
        ----------
        v : Any
            Raw error value.

        Returns
        -------
        str
            Normalized error string.
        """
        if v is None:
            return ""
        return v if isinstance(v, str) else str(v)