from __future__ import annotations

from argparse import Namespace
from collections import Counter
from contextlib import contextmanager, suppress
from dataclasses import dataclass
from datetime import datetime
from pathlib import Path
from typing import Any, Dict, Iterable, Optional, Tuple
import logging
import sys
import shutil
import inspect
import io
import zipfile
import uuid
import time

import numpy as np
from ase.io import read, write
from ase.optimize import BFGS, BFGSLineSearch, FIRE, LBFGS, MDMin
from ase.mep import DimerControl, MinModeAtoms, MinModeTranslate
from ase.vibrations import Vibrations
from ase.vibrations.data import VibrationsData
from ase.thermochemistry import IdealGasThermo
from ase.visualize import view
from ase import Atoms, units

from .utils import (
    Sum_of_atomic_energies,
    ensure_fairchem_loaded,
    extract_xyz_metadata,
    get_fairchem_import_error,
    get_ground_state_atomic_multiplicity_source,
    get_predict_unit_with_local_fallback,
    rmsd_simple,
)
from . import __version__


@dataclass(frozen=True)
class WorkflowPaths:
    """Collection of standard output artifacts emitted by a run."""

    trajectory: Path
    log: Path
    final_geometry: Path
    freq_archive: Optional[Path] = None


@dataclass
class AtomContext:
    """Container bundling an ASE atoms object with metadata for downstream steps."""

    atoms: Atoms
    counts: Counter
    device: str
    predictor: Any


LOG_SECTION_WIDTH = 77


def format_log_timestamp(moment: Optional[datetime] = None) -> str:
    """Return a compact timestamp string for log messages."""
    return (moment or datetime.now()).strftime("%Y-%m-%d %H:%M:%S")


def timestamp_token(moment: Optional[datetime] = None) -> str:
    """Return a bracketed timestamp token for log output."""
    return f"[{format_log_timestamp(moment)}]"


def _log_context_break(logger: logging.Logger, lines: int = 2) -> None:
    """Add visible vertical spacing before a new log context begins."""
    for _ in range(lines):
        logger.info("")


def _format_banner_line(text: str, marker: str, width: int = LOG_SECTION_WIDTH) -> str:
    """Return a centered banner line framed with repeated marker characters."""
    clean = f" {text.strip()} "
    if len(clean) >= width:
        return clean
    padding = width - len(clean)
    left = padding // 2
    right = padding - left
    return f"{marker * left}{clean}{marker * right}"


def log_timestamp_event(
    logger: logging.Logger,
    message: str,
    *,
    moment: Optional[datetime] = None,
    level: int = logging.INFO,
    marker: str = "-",
) -> None:
    """Emit a timestamped event line with visible separators."""
    _log_context_break(logger)
    banner = f"{timestamp_token(moment)} {message}"
    is_timing_banner = marker == "=" and message.startswith("Timing [")
    if is_timing_banner:
        log_rule(logger, char="=")
    logger.log(level, "%s", _format_banner_line(banner, marker))
    if is_timing_banner:
        log_rule(logger, char="=")


def log_timestamp_message(
    logger: logging.Logger,
    message: str,
    *,
    moment: Optional[datetime] = None,
    level: int = logging.INFO,
    prefix: str = ">>>",
) -> None:
    """Emit a compact timestamped progress message."""
    logger.log(level, "%s %s %s", prefix, timestamp_token(moment), message)


def log_rule(logger: logging.Logger, char: str = "=") -> None:
    """Emit a full-width separator rule."""
    logger.info(char * LOG_SECTION_WIDTH)


def log_main_banner(logger: logging.Logger) -> None:
    """Emit the UMA-ASE banner shown at the top of CLI log files."""
    _log_context_break(logger)
    log_rule(logger, char="=")
    logger.info("U M A - A S E   v%s", __version__)
    logger.info("Universal Model for Atoms - Atomistic Simulation Environment")
    log_rule(logger, char="=")
    logger.info("")


def log_section_header(logger: logging.Logger, title: str, *, char: str = "=") -> None:
    """Emit a titled full-width section header."""
    _log_context_break(logger)
    log_rule(logger, char=char)
    logger.info("%s", _format_banner_line(title, char))
    log_rule(logger, char=char)


def log_key_value_block(
    logger: logging.Logger,
    title: str,
    items: Iterable[tuple[str, Any]],
    *,
    char: str = "-",
) -> None:
    """Emit a titled block of aligned key/value pairs."""
    log_section_header(logger, title, char=char)
    for key, value in items:
        logger.info("  %-20s: %s", key, value)
    logger.info("")


def log_atomic_reference_method(logger: logging.Logger, *, title: str = "Atomic Reference Energies") -> None:
    """Describe how atomic reference energies are computed for bonding-energy estimates."""
    source = get_ground_state_atomic_multiplicity_source()
    log_section_header(logger, title, char="-")
    logger.info("  Reference model             : single isolated atom per element")
    logger.info("  Reference cell              : periodic cubic box (a = 5.0 Å)")
    logger.info("  UMA cutoff condition        : box kept inside the default 6.0 Å cutoff")
    if source.get("primary"):
        logger.info("  Multiplicity source         : %s", source["primary"])
    if source.get("fallback"):
        logger.info("  Multiplicity fallback       : %s", source["fallback"])
    if source.get("rule"):
        logger.info("  Multiplicity rule           : %s", source["rule"])
    logger.info("")


def capture_stdout(func, *args, **kwargs):
    """Capture stdout emitted by *func* and replay it to the original stream."""
    buffer = io.StringIO()
    original_stdout = sys.stdout
    try:
        sig = inspect.signature(func)
        if "log" in sig.parameters and "log" not in kwargs:
            kwargs["log"] = buffer
            result = func(*args, **kwargs)
        else:
            sys.stdout = buffer
            result = func(*args, **kwargs)
    finally:
        sys.stdout = original_stdout
    output = buffer.getvalue()
    if output:
        original_stdout.write(output)
        original_stdout.flush()
    return result, output


class LoggerWriter:
    """Adapter that lets ASE optimizers log through a standard logger."""

    def __init__(self, logger: logging.Logger):
        self._logger = logger

    def write(self, message: str) -> None:
        message = message.strip()
        if message:
            self._logger.info(message)

    def flush(self) -> None:  # pragma: no cover - interface requirement
        pass

    def close(self) -> None:  # pragma: no cover - interface requirement
        pass


def _orthonormalize_mode_basis(vectors: Iterable[np.ndarray], tol: float = 1e-10) -> np.ndarray:
    """Return an orthonormal basis assembled with modified Gram-Schmidt."""

    basis: list[np.ndarray] = []
    for vector in vectors:
        candidate = np.asarray(vector, dtype=float).copy()
        for item in basis:
            candidate -= np.dot(item, candidate) * item
        norm = float(np.linalg.norm(candidate))
        if np.isfinite(norm) and norm > tol:
            basis.append(candidate / norm)
    if not basis:
        return np.empty((0, 0), dtype=float)
    return np.column_stack(basis)


def _rigid_body_basis(atoms: Atoms, indices: np.ndarray) -> np.ndarray:
    """Build mass-weighted rigid-body translation/rotation vectors."""

    active_atoms = atoms[indices]
    n_atoms = len(active_atoms)
    if n_atoms == 0:
        return np.empty((0, 0), dtype=float)

    masses = active_atoms.get_masses()
    sqrt_masses = np.sqrt(masses)
    vectors: list[np.ndarray] = []

    for axis in range(3):
        translation = np.zeros((n_atoms, 3), dtype=float)
        translation[:, axis] = sqrt_masses
        vectors.append(translation.ravel())

    if not np.any(active_atoms.get_pbc()):
        positions = active_atoms.get_positions()
        center_of_mass = np.average(positions, axis=0, weights=masses)
        relative_positions = positions - center_of_mass
        for axis_vector in np.eye(3):
            rotation = np.cross(np.broadcast_to(axis_vector, relative_positions.shape), relative_positions)
            rotation *= sqrt_masses[:, None]
            vectors.append(rotation.ravel())

    return _orthonormalize_mode_basis(vectors)


def _project_rigid_body_modes(H: np.ndarray, atoms: Atoms, indices: np.ndarray) -> np.ndarray:
    """Project rigid translations/rotations from a Cartesian Hessian."""

    active_masses = atoms[indices].get_masses()
    inv_sqrt_masses = np.repeat(active_masses ** -0.5, 3)
    mw_hessian = inv_sqrt_masses[:, None] * H * inv_sqrt_masses

    basis = _rigid_body_basis(atoms, indices)
    if basis.size:
        projector = np.eye(mw_hessian.shape[0]) - basis @ basis.T
        mw_hessian = projector @ mw_hessian @ projector
        mw_hessian = 0.5 * (mw_hessian + mw_hessian.T)

    sqrt_masses = np.repeat(np.sqrt(active_masses), 3)
    projected = sqrt_masses[:, None] * mw_hessian * sqrt_masses
    return 0.5 * (projected + projected.T)


class SymmetrizedVibrations(Vibrations):
    """ASE Vibrations variant with explicit symmetrization and rigid-mode projection."""

    def read(self, method: str = "standard", direction: str = "central") -> None:
        self.method = method.lower()
        self.direction = direction.lower()
        assert self.method in ["standard", "frederiksen"]
        assert self.direction in ["central", "forward", "backward"]

        n = 3 * len(self.indices)
        H = np.empty((n, n))
        r = 0

        eq_disp = self._eq_disp()
        if self.direction != "central":
            feq = eq_disp.forces()

        for a, i in self._iter_ai():
            disp_minus = self._disp(a, i, -1)
            disp_plus = self._disp(a, i, 1)

            fminus = disp_minus.forces()
            fplus = disp_plus.forces()
            if self.method == "frederiksen":
                fminus[a] -= fminus.sum(0)
                fplus[a] -= fplus.sum(0)

            if self.nfree == 4:
                fminusminus = self._disp(a, i, -2).forces()
                fplusplus = self._disp(a, i, 2).forces()
                if self.method == "frederiksen":
                    fminusminus[a] -= fminusminus.sum(0)
                    fplusplus[a] -= fplusplus.sum(0)

            if self.direction == "central":
                if self.nfree == 2:
                    row = (fminus - fplus)[self.indices].ravel() / (2 * self.delta)
                else:
                    assert self.nfree == 4
                    row = (
                        -fminusminus + 8 * fminus - 8 * fplus + fplusplus
                    )[self.indices].ravel() / (12 * self.delta)
            elif self.direction == "forward":
                row = (feq - fplus)[self.indices].ravel() / self.delta
            else:
                assert self.direction == "backward"
                row = (fminus - feq)[self.indices].ravel() / self.delta

            H[r] = row
            r += 1

        # Explicitly average mirrored matrix elements: H_ij <- (H_ij + H_ji) / 2.
        H = 0.5 * (H + H.T)
        H = _project_rigid_body_modes(H, self.atoms, self.indices)
        self.H = H

        masses = self.atoms.get_masses()
        if any(masses[self.indices] == 0):
            raise RuntimeError(
                "Zero mass encountered in one or more of the vibrated atoms. "
                "Use Atoms.set_masses() to set all masses to non-zero values."
            )

        self.im = np.repeat(masses[self.indices] ** -0.5, 3)
        self._vibrations = VibrationsData.from_2d(self.atoms, self.H, indices=self.indices)

        omega2, modes = np.linalg.eigh(self.im[:, None] * H * self.im)
        self.modes = modes.T.copy()

        scale = units._hbar * 1e10 / np.sqrt(units._e * units._amu)
        self.hnu = scale * omega2.astype(complex) ** 0.5


class TorchUnavailable(RuntimeError):
    """Lightweight exception indicating torch is not present."""


OPTIMIZER_CLASSES = {
    "BFGS": BFGS,
    "LBFGS": LBFGS,
    "BFGS_LINESEARCH": BFGSLineSearch,
    "FIRE": FIRE,
    "MDMIN": MDMin,
}

RUN_LABELS = {
    "sp": "SP",
    "geoopt": "OPT",
    "freqs": "FREQS",
    "ts": "TS",
}

RUN_TITLES = {
    "sp": "Single Point",
    "geoopt": "Geometry Optimization",
    "freqs": "Vibrational Analysis",
    "ts": "Transition State Search",
    "md": "Molecular Dynamics",
}

MODE_VISUAL_MAX_DISPLACEMENT_ANG = 0.30
MODE_VISUAL_FRAMES = 30


def build_output_paths(
    input_path: Path,
    run_sequence: Optional[Iterable[str]] = None,
    explicit_geoopt: bool = True,
) -> WorkflowPaths:
    """Generate canonical output filenames derived from the input geometry."""
    parent = input_path.parent
    stem = input_path.stem
    sequence = list(run_sequence) if run_sequence else ["geoopt"]
    normalized = [
        item.lower() if isinstance(item, str) else str(item).lower()
        for item in sequence
    ]
    if set(normalized) == {"freqs"}:
        labels = ["FREQS"]
    else:
        labels = []
        for item in normalized:
            label = RUN_LABELS.get(item, item.upper())
            if label not in labels:
                labels.append(label)
    log_suffix = "-".join(labels)
    log_name = f"{stem}-{log_suffix}.log"
    freq_archive = None
    if any(item in {"freqs", "ts"} for item in normalized):
        freq_archive = parent / f"{stem}-FREQS-modes.zip"
    return WorkflowPaths(
        trajectory=parent / f"{stem}-{log_suffix}.traj",
        log=parent / log_name,
        final_geometry=parent / f"{stem}-geoopt-{log_suffix}.xyz",
        freq_archive=freq_archive,
    )


def remove_if_exists(paths: Iterable[Path]) -> None:
    """Delete pre-existing files or directories produced by earlier runs."""
    for path in paths:
        if path.is_dir():
            shutil.rmtree(path, ignore_errors=True)
        elif path.exists():
            path.unlink()


def _convert_traj_to_xyz(traj_file: Path, destination: Path, logger: logging.Logger) -> bool:
    """Convert an ASE trajectory to multi-frame XYZ, deleting the original."""

    try:
        frames = read(traj_file, ":")
    except Exception as exc:
        logger.warning("Unable to read vibrational trajectory %s: %s", traj_file, exc)
        return False
    try:
        destination.parent.mkdir(parents=True, exist_ok=True)
        write(destination, frames, format="xyz")
    except OSError as exc:
        logger.warning("Unable to write %s: %s", destination, exc)
        return False
    try:
        traj_file.unlink()
    except OSError:
        pass
    return True


def _cleanup_vibration_cache(cache_dir: Path, target_dir: Path, logger: logging.Logger) -> None:
    """Move trajectory files to target dir and remove cache artifacts."""
    if not cache_dir.exists():
        return

    moved = 0
    for traj_file in cache_dir.glob("*.traj"):
        destination = target_dir / f"{traj_file.stem}.xyz"
        if _convert_traj_to_xyz(traj_file, destination, logger):
            moved += 1

    removed = 0
    for cache_file in cache_dir.glob("cache*.json"):
        try:
            cache_file.unlink()
            removed += 1
        except OSError as exc:
            logger.warning("Unable to delete %s: %s", cache_file, exc)

    try:
        shutil.rmtree(cache_dir)
    except OSError as exc:
        logger.warning("Unable to remove cache directory %s: %s", cache_dir, exc)

    if moved:
        logger.info("Moved %d normal mode components to %s", moved, target_dir)
    if removed:
        logger.info("Removed %d temporary cache files", removed)


def _package_vibrational_data(
    source_dir: Path,
    archive_path: Path,
    logger: logging.Logger,
) -> Optional[Path]:
    """Create a ZIP archive containing the vibrational trajectories."""

    converted = 0
    for traj_file in source_dir.rglob("*.traj"):
        xyz_path = traj_file.with_suffix(".xyz")
        if _convert_traj_to_xyz(traj_file, xyz_path, logger):
            converted += 1
    if converted:
        logger.info("Converted %d normal mode trajectories to XYZ", converted)

    files = [path for path in source_dir.rglob("*") if path.is_file()]
    if not files:
        logger.warning("No vibrational files found to package in %s", source_dir)
        return None
    try:
        archive_path.parent.mkdir(parents=True, exist_ok=True)
        if archive_path.exists():
            archive_path.unlink()
        with zipfile.ZipFile(archive_path, "w", zipfile.ZIP_DEFLATED) as archive:
            for file_path in files:
                try:
                    relative = file_path.relative_to(source_dir.parent)
                except ValueError:
                    relative = file_path.name
                archive.write(file_path, str(relative))
    except OSError as exc:
        logger.warning("Unable to package vibrational data into %s: %s", archive_path, exc)
        return None
    logger.info("Packaged vibrational data into %s", archive_path)
    return archive_path


def _scale_mode_for_visualization(
    mode: np.ndarray,
    max_displacement: float = MODE_VISUAL_MAX_DISPLACEMENT_ANG,
) -> np.ndarray:
    """Scale a Cartesian normal mode to a fixed maximum atomic displacement."""

    cartesian_mode = np.real_if_close(np.asarray(mode, dtype=np.complex128)).real
    atom_norms = np.linalg.norm(cartesian_mode, axis=1)
    largest_displacement = float(np.max(atom_norms)) if atom_norms.size else 0.0
    if not np.isfinite(largest_displacement) or largest_displacement <= 0.0:
        return np.zeros_like(cartesian_mode, dtype=float)
    return cartesian_mode * (max_displacement / largest_displacement)


def _write_all_vibrational_modes(
    vib_data: VibrationsData,
    target_dir: Path,
    base_name: str,
    logger: logging.Logger,
    *,
    frames: int = MODE_VISUAL_FRAMES,
    max_displacement: float = MODE_VISUAL_MAX_DISPLACEMENT_ANG,
) -> None:
    """Write one XYZ animation per vibrational mode using a fixed display amplitude."""

    target_dir.mkdir(parents=True, exist_ok=True)
    base_atoms = vib_data.get_atoms()
    modes = np.asarray(vib_data.get_modes(all_atoms=True))
    written = 0

    for mode_index, mode in enumerate(modes):
        try:
            scaled_mode = _scale_mode_for_visualization(mode, max_displacement=max_displacement)
            frames_data = []
            for phase in np.linspace(0.0, 2.0 * np.pi, frames, endpoint=False):
                image = base_atoms.copy()
                image.positions = base_atoms.positions + np.sin(phase) * scaled_mode
                frames_data.append(image)
            destination = target_dir / f"{base_name}.{mode_index}.xyz"
            write(destination, frames_data, format="xyz")
            written += 1
        except Exception as exc:  # pragma: no cover - ASE behavior varies
            logger.warning("Unable to write vibrational mode %d: %s", mode_index, exc)

    logger.info(
        "Wrote %d/%d vibrational mode trajectories using normalized eigenvectors "
        "with a maximum atomic displacement of %.2f Å.",
        written,
        len(modes),
        max_displacement,
    )


def _log_normal_mode_matrices(
    vib_data: VibrationsData,
    atoms: Atoms,
    logger: logging.Logger,
) -> None:
    """Emit formatted displacement matrices for every normal mode."""
    try:
        mode_matrices = np.asarray(vib_data.get_modes(all_atoms=True)).real
    except Exception as exc:  # pragma: no cover - ASE internals may vary
        logger.warning("Unable to extract normal-mode displacement matrices: %s", exc)
        return

    n_modes, n_atoms, _ = mode_matrices.shape
    symbols = atoms.get_chemical_symbols()
    component_labels = ("dx", "dy", "dz")
    col_width = 13
    chunk_size = 6

    logger.info(
        "Normal mode displacement matrix (columns = modes, rows = atom component dx/dy/dz)."
    )
    for start in range(0, n_modes, chunk_size):
        end = min(start + chunk_size, n_modes)
        header = "Atom Sym Comp" + "".join(
            f"{f'Mode {mode_index}':>{col_width}}"
            for mode_index in range(start, end)
        )
        logger.info(header)

        for atom_index in range(n_atoms):
            symbol = symbols[atom_index] if atom_index < len(symbols) else "?"
            for comp_index, label in enumerate(component_labels):
                row_label = f"{atom_index + 1:4d} {symbol:>3} {label:>4}"
                values = "".join(
                    f"{mode_matrices[mode_index, atom_index, comp_index]:{col_width}.6f}"
                    for mode_index in range(start, end)
                )
                logger.info("%s%s", row_label, values)
        logger.info("")


@contextmanager
def configure_logging(log_path: Path):
    """Context manager that emits log messages to stdout and a logfile."""
    logger = logging.Logger(f"uma_workflow.{uuid.uuid4().hex}")
    logger.setLevel(logging.INFO)
    logger.propagate = False

    formatter = logging.Formatter("%(message)s")
    stream_handler = logging.StreamHandler(sys.stdout)
    stream_handler.setFormatter(formatter)
    file_handler = logging.FileHandler(log_path, mode="w")
    file_handler.setFormatter(formatter)

    logger.addHandler(stream_handler)
    logger.addHandler(file_handler)

    try:
        yield logger
    finally:
        for handler in logger.handlers[:]:
            handler.close()
            logger.removeHandler(handler)


def select_device(force_cpu: bool = False) -> str:
    """Return the preferred execution device based on torch availability."""
    if force_cpu:
        return "cpu"
    try:
        import torch
    except ModuleNotFoundError as exc:  # pragma: no cover - environment dependent
        raise TorchUnavailable("PyTorch is required to determine execution device.") from exc
    try:
        return "cuda" if torch.cuda.is_available() else "cpu"
    except Exception as exc:
        raise TorchUnavailable("Failed to query torch CUDA availability.") from exc


def log_arguments(args: Namespace, logger: logging.Logger, device: str) -> None:
    """Echo the parsed CLI arguments to aid reproducibility."""
    run_types = _normalize_run_types(args.run_type)
    includes_geoopt = "geoopt" in run_types
    includes_freqs = "freqs" in run_types
    includes_ts = "ts" in run_types
    includes_md = "md" in run_types
    items: list[tuple[str, Any]] = [
        ("input", args.input),
        ("run_type", args.run_type),
        ("chg", args.chg),
        ("spin", args.spin),
        ("checkpoint", args.mlff_chk),
        ("task", args.mlff_task),
    ]
    if includes_geoopt or includes_ts:
        items.extend([
            ("optimizer", getattr(args, "optimizer", "-")),
            ("grad", getattr(args, "grad", "-")),
            ("iter", getattr(args, "iter", "-")),
        ])
    if includes_freqs or includes_ts:
        items.extend([
            ("temperature", f"{getattr(args, 'temp', 0.0):.4f} K"),
            ("pressure", f"{getattr(args, 'press', 0.0):.2f} Pa"),
        ])
    if includes_ts:
        items.extend([
            ("ts_displacement", getattr(args, "ts_displacement", "-")),
            ("ts_dimer_separation", getattr(args, "ts_dimer_separation", "-")),
            ("ts_trial_step", getattr(args, "ts_trial_step", "-")),
            ("ts_max_step", getattr(args, "ts_max_step", "-")),
            ("ts_max_rot", getattr(args, "ts_max_rot", "-")),
        ])
    if includes_md:
        md_engine = getattr(args, "md_engine", "langevin")
        md_restart_traj = getattr(args, "md_restart_traj", None)
        items.extend([
            ("md_engine", md_engine),
            ("md_steps", getattr(args, "md_steps", "-")),
            ("md_timestep_fs", getattr(args, "md_timestep_fs", "-")),
        ])
        if md_engine != "velocity_verlet" or not md_restart_traj:
            items.append(("md_temperature", getattr(args, "md_temperature", "-")))
        items.extend([
            ("md_traj_interval", getattr(args, "md_traj_interval", "-")),
            ("md_log_interval", getattr(args, "md_log_interval", "-")),
        ])
        if md_engine == "langevin":
            items.append(("md_friction", getattr(args, "md_friction", "-")))
        if md_engine in {"nvt_berendsen", "npt_berendsen"}:
            items.append(("md_relax_fs", getattr(args, "md_relax_fs", "-")))
        if md_engine in {
            "npt_berendsen",
            "isotropic_mtk_npt",
            "mtk_npt",
            "langevin_baoab_npt",
            "melchionna_npt",
        }:
            items.append(("md_pressure_bar", getattr(args, "md_pressure_bar", "-")))
        if md_engine == "npt_berendsen":
            items.extend([
                ("md_barostat_relax_fs", getattr(args, "md_barostat_relax_fs", "-")),
                ("md_compressibility_bar", getattr(args, "md_compressibility_bar", "-")),
            ])
        if md_engine == "nose_hoover_chain":
            items.extend([
                ("md_nose_tdamp_fs", getattr(args, "md_nose_tdamp_fs", "-")),
                ("md_nose_tchain", getattr(args, "md_nose_tchain", "-")),
                ("md_nose_tloop", getattr(args, "md_nose_tloop", "-")),
            ])
        if md_engine == "bussi":
            items.append(("md_bussi_taut_fs", getattr(args, "md_bussi_taut_fs", "-")))
        if md_engine == "andersen":
            items.append(("md_andersen_prob", getattr(args, "md_andersen_prob", "-")))
        if md_engine in {"isotropic_mtk_npt", "mtk_npt"}:
            items.extend([
                ("md_mtk_tdamp_fs", getattr(args, "md_mtk_tdamp_fs", "-")),
                ("md_mtk_pdamp_fs", getattr(args, "md_mtk_pdamp_fs", "-")),
                ("md_mtk_tchain", getattr(args, "md_mtk_tchain", "-")),
                ("md_mtk_pchain", getattr(args, "md_mtk_pchain", "-")),
                ("md_mtk_tloop", getattr(args, "md_mtk_tloop", "-")),
                ("md_mtk_ploop", getattr(args, "md_mtk_ploop", "-")),
            ])
        if md_engine in {"langevin_baoab_nvt", "langevin_baoab_npt"}:
            items.append(("md_baoab_ttau_fs", getattr(args, "md_baoab_ttau_fs", "-")))
            if md_engine == "langevin_baoab_npt":
                items.extend([
                    ("md_baoab_ptau_fs", getattr(args, "md_baoab_ptau_fs", "-")),
                    ("md_baoab_pmass", getattr(args, "md_baoab_pmass", "-")),
                    ("md_baoab_pmass_factor", getattr(args, "md_baoab_pmass_factor", "-")),
                    ("md_baoab_hydrostatic", getattr(args, "md_baoab_hydrostatic", "-")),
                ])
        if md_engine == "melchionna_npt":
            items.extend([
                ("md_melchionna_ttime_fs", getattr(args, "md_melchionna_ttime_fs", "-")),
                ("md_melchionna_ptime_fs", getattr(args, "md_melchionna_ptime_fs", "-")),
                ("md_melchionna_bulk_gpa", getattr(args, "md_melchionna_bulk_gpa", "-")),
                ("md_melchionna_mask", getattr(args, "md_melchionna_mask", "-")),
            ])
        if getattr(args, "md_use_pbc", False):
            items.extend([
                ("md_use_pbc", getattr(args, "md_use_pbc", "-")),
                ("md_cell_a", getattr(args, "md_cell_a", "-")),
                ("md_cell_b", getattr(args, "md_cell_b", "-")),
                ("md_cell_c", getattr(args, "md_cell_c", "-")),
                ("md_cell_alpha", getattr(args, "md_cell_alpha", "-")),
                ("md_cell_beta", getattr(args, "md_cell_beta", "-")),
                ("md_cell_gamma", getattr(args, "md_cell_gamma", "-")),
            ])
        if md_restart_traj:
            items.append(("md_restart_traj", md_restart_traj))
    items.extend([
        ("cpu_only", getattr(args, "cpu", False)),
        ("device", device),
    ])
    log_main_banner(logger)
    log_key_value_block(logger, "CLI Arguments", items, char="-")


def summarise_structure(
    args: Namespace,
    atoms: Atoms,
    logger: logging.Logger,
    *,
    comment: Optional[str] = None,
) -> Counter:
    """Log structural information about the current Atoms object."""
    counts = Counter(atoms.get_chemical_symbols())
    items: list[tuple[str, Any]] = [
        ("input", args.input),
        ("number_of_atoms", len(atoms)),
        ("formula", atoms.get_chemical_formula()),
        ("element_counts", dict(counts)),
        ("charge", args.chg),
        ("spin_multiplicity", args.spin),
    ]
    if comment:
        items.append(("xyz_comment", comment))
    log_key_value_block(logger, "Atomic System Summary", items, char="-")
    return counts


def log_geometry(logger: logging.Logger, atoms: Atoms, title: str) -> None:
    """Log atomic coordinates of *atoms* with a descriptive *title*."""
    log_section_header(logger, title, char="-")
    symbols = atoms.get_chemical_symbols()
    positions = atoms.get_positions()
    for symbol, (x_coord, y_coord, z_coord) in zip(symbols, positions):
        logger.info("  %3s %12.6f %12.6f %12.6f", symbol, x_coord, y_coord, z_coord)
    logger.info("")


@contextmanager
def log_timing(logger: logging.Logger, label: str):
    """Log wall timing for a named phase."""
    wall_start = time.perf_counter()
    try:
        yield
    finally:
        wall_elapsed = time.perf_counter() - wall_start
        log_timestamp_event(
            logger,
            f"Timing [{label}]: wall={wall_elapsed:.3f}s",
            marker="=",
        )


def log_total_timing(
    logger: logging.Logger,
    label: str,
    wall_start: float,
    cpu_start: float,
) -> None:
    """Log total CPU and wall timing for a run."""
    wall_elapsed = time.perf_counter() - wall_start
    cpu_elapsed = time.process_time() - cpu_start
    log_timestamp_event(
        logger,
        f"Timing [Run total - {label}]: wall={wall_elapsed:.3f}s cpu={cpu_elapsed:.3f}s",
        marker="=",
    )


def log_run_metadata(
    args: Namespace,
    logger: logging.Logger,
    run_type: str,
    device: str,
) -> None:
    """Emit run metadata for traceability."""
    run_title = RUN_TITLES.get(run_type, run_type.upper())
    metadata: Dict[str, Any] = {
        "run_type_requested": getattr(args, "run_type", run_type),
        "run_type": run_type,
        "checkpoint": getattr(args, "mlff_chk", None),
        "method": "MLIP-UMA",
        "task": getattr(args, "mlff_task", None),
        "chg": getattr(args, "chg", None),
        "spin": getattr(args, "spin", None),
        "device": device,
    }
    if run_type in {"geoopt", "ts"}:
        metadata["optimizer"] = getattr(args, "optimizer", None)
        metadata["grad"] = getattr(args, "grad", None)
        metadata["iter"] = getattr(args, "iter", None)
    if run_type in {"freqs", "ts"}:
        metadata["temperature"] = getattr(args, "temp", None)
        metadata["pressure"] = getattr(args, "press", None)
    if run_type == "ts":
        metadata["ts_displacement"] = getattr(args, "ts_displacement", None)
        metadata["ts_dimer_separation"] = getattr(args, "ts_dimer_separation", None)
        metadata["ts_trial_step"] = getattr(args, "ts_trial_step", None)
        metadata["ts_max_step"] = getattr(args, "ts_max_step", None)
        metadata["ts_max_rot"] = getattr(args, "ts_max_rot", None)
    log_section_header(logger, f"{run_title} Metadata", char="-")
    if run_type in {"freqs", "ts"}:
        logger.info("Thermochemistry inputs are required for vibrational analysis.")
    ordered_keys = [
        "run_type_requested",
        "run_type",
        "checkpoint",
        "method",
        "task",
        "chg",
        "spin",
        "device",
        "optimizer",
        "grad",
        "iter",
        "temperature",
        "pressure",
        "ts_displacement",
        "ts_dimer_separation",
        "ts_trial_step",
        "ts_max_step",
        "ts_max_rot",
    ]
    rows: list[tuple[str, Any]] = []
    for key in ordered_keys:
        if key in metadata and metadata[key] is not None:
            rows.append((key, metadata[key]))
    for key in sorted(metadata):
        if key in ordered_keys or metadata[key] is None:
            continue
        rows.append((key, metadata[key]))
    for key, value in rows:
        logger.info("  %-20s: %s", key, value)
    logger.info("")


def _normalize_run_types(raw: Any) -> set[str]:
    if raw is None:
        return set()
    if isinstance(raw, str):
        items = [raw]
    else:
        items = list(raw)
    return {str(item).strip().lower() for item in items if str(item).strip()}


def setup_calculated_atoms(
    args: Namespace,
    logger: logging.Logger,
) -> Tuple[int, Optional[AtomContext]]:
    """Create an ASE atoms object wired with the requested UMA calculator."""
    _, fairchem_calculator_cls, _ = ensure_fairchem_loaded()
    fairchem_import_error = get_fairchem_import_error()
    if fairchem_import_error is not None or fairchem_calculator_cls is None:
        logger.error("UMA FairChem libraries unavailable: %s", fairchem_import_error)
        return 1, None

    try:
        device = select_device(force_cpu=getattr(args, "cpu", False))
    except TorchUnavailable as exc:
        logger.error(str(exc))
        return 1, None

    if device == "cpu" and not getattr(args, "cpu", False):
        logger.info("CUDA device unavailable; running on CPU.")

    input_path = Path(args.input)
    metadata = extract_xyz_metadata(input_path)

    chg_explicit = bool(getattr(args, "_chg_explicit", False))
    if metadata.charge is not None and not chg_explicit:
        args.chg = metadata.charge
        logger.info("Charge inferred from XYZ metadata: %s", args.chg)
    elif metadata.charge is not None and chg_explicit and metadata.charge != args.chg:
        logger.debug(
            "XYZ metadata charge %s ignored because CLI provided charge %s.",
            metadata.charge,
            args.chg,
        )

    spin_explicit = bool(getattr(args, "_spin_explicit", False))
    if metadata.spin is not None and not spin_explicit:
        args.spin = metadata.spin
        logger.info("Spin multiplicity inferred from XYZ metadata: %s", args.spin)
    elif metadata.spin is not None and spin_explicit and metadata.spin != args.spin:
        logger.debug(
            "XYZ metadata spin %s ignored because CLI provided spin %s.",
            metadata.spin,
            args.spin,
        )

    log_arguments(args, logger, device)

    log_timestamp_event(logger, f"Preparing ASE objects for {input_path.name}")
    with log_timing(logger, "ASE geometry load"):
        atoms = read(args.input)
    atoms.info["charge"] = args.chg
    atoms.info["spin"] = args.spin

    xyz_comment = metadata.comment
    if xyz_comment:
        atoms.info.setdefault("uma_comment", xyz_comment)
    if metadata.url:
        atoms.info.setdefault("uma_comment_url", metadata.url)

    counts = summarise_structure(args, atoms, logger, comment=xyz_comment)

    try:
        with log_timing(logger, "UMA checkpoint load"):
            predictor = get_predict_unit_with_local_fallback(
                args.mlff_chk,
                device=device,
                logger=logger,
            )
    except FileNotFoundError as exc:
        logger.error(str(exc))
        return 1, None
    except KeyError:
        logger.error(
            "UMA checkpoint '%s' not found. Please provide a valid identifier.",
            args.mlff_chk,
        )
        return 1, None
    except Exception as exc:
        logger.error("Failed to load UMA checkpoint '%s': %s", args.mlff_chk, exc)
        return 1, None

    log_key_value_block(
        logger,
        "UMA Calculator Setup",
        [
            ("uma_checkpoint", args.mlff_chk),
            ("uma_model", args.mlff_task),
        ],
        char="-",
    )
    try:
        with log_timing(logger, "ASE calculator setup"):
            atoms.calc = fairchem_calculator_cls(predictor, task_name=args.mlff_task)
    except Exception as exc:
        logger.error("Failed to initialise UMA calculator for task '%s': %s", args.mlff_task, exc)
        return 1, None
    log_timestamp_event(logger, f"ASE objects ready on device {device}")

    return 0, AtomContext(atoms=atoms, counts=counts, device=device, predictor=predictor)


def prepare_atoms_for_vibrations(
    args: Namespace,
    logger: logging.Logger,
) -> Tuple[int, Optional[AtomContext]]:
    """Set up atoms for vibrational analysis without running an optimisation."""
    status, context = setup_calculated_atoms(args, logger)
    if status != 0 or context is None:
        return status, None
    return 0, context


def _optimizer_converged(optimizer, run_result) -> bool:
    """Determine whether an ASE optimizer reported convergence."""
    converged: Optional[bool] = None
    if run_result is not None:
        with suppress(Exception):
            converged = bool(run_result)

    if converged is None:
        converged_attr = getattr(optimizer, "converged", None)
        if callable(converged_attr):
            gradient = None
            optimizable = getattr(optimizer, "optimizable", None)
            if optimizable is not None and hasattr(optimizable, "get_gradient"):
                with suppress(Exception):
                    gradient = optimizable.get_gradient()
            if gradient is None:
                atoms_obj = getattr(optimizer, "atoms", None)
                if atoms_obj is not None and hasattr(atoms_obj, "get_forces"):
                    with suppress(Exception):
                        gradient = atoms_obj.get_forces().ravel()
            if gradient is not None:
                with suppress(TypeError, AttributeError, AssertionError):
                    converged = converged_attr(gradient)
            else:
                with suppress(TypeError, AttributeError):
                    converged = converged_attr()
        elif isinstance(converged_attr, bool):
            converged = converged_attr

    if converged is None:
        converged = False
    return converged

def _resolve_optimizer(name: str):
    """Map a user-supplied optimiser name to an ASE optimiser class."""
    key = name.replace("-", "_").upper()
    return OPTIMIZER_CLASSES.get(key)


def run_geometry_optimization(
    args: Namespace,
    paths: WorkflowPaths,
    logger: logging.Logger,
) -> Tuple[int, Optional[AtomContext]]:
    """Execute a geometry optimisation and capture derived metrics."""
    run_wall_start = time.perf_counter()
    run_cpu_start = time.process_time()
    try:
        status, context = setup_calculated_atoms(args, logger)
        if status != 0 or context is None:
            return status, None

        log_run_metadata(args, logger, "geoopt", context.device)

        optimizer_cls = _resolve_optimizer(args.optimizer)
        if optimizer_cls is None:
            logger.error(
                "Unknown optimizer '%s'. Available options: %s",
                args.optimizer,
                ", ".join(sorted(OPTIMIZER_CLASSES)),
            )
            return 1, None

        optimizer = optimizer_cls(
            context.atoms,
            trajectory=str(paths.trajectory),
            logfile=LoggerWriter(logger),
        )

        start_time = datetime.now()
        log_section_header(logger, "Geometry Optimization Run")
        log_timestamp_event(
            logger,
            f"Running Geometry Optimization on device {context.device}",
            moment=start_time,
        )
        logger.info("Target fmax %s eV/A. Max number of iterations %s", args.grad, args.iter)

        log_geometry(logger, context.atoms, "Initial geometry for GeoOpt (Å):")

        try:
            with log_timing(logger, "GeoOpt optimization"):
                run_result = optimizer.run(fmax=args.grad, steps=args.iter)
        except Exception as exc:
            logger.error("Optimizer failed: %s", exc)
            return 1, None

        finish_time = datetime.now()
        converged = _optimizer_converged(optimizer, run_result)

        if converged:
            log_timestamp_event(logger, "Optimization finished", moment=finish_time)
        else:
            log_timestamp_event(
                logger,
                "GeoOpt stopped: maximum iterations reached without convergence.",
                moment=finish_time,
                level=logging.ERROR,
                marker="!",
            )
            return 1, None
        logger.info("")
        log_section_header(logger, "Geometry Optimization Results")
        log_geometry(logger, context.atoms, "Final geometry (Å):")

        potential_energy = context.atoms.get_potential_energy()
        total_energy = context.atoms.get_total_energy()

        log_section_header(logger, "Energy Summary", char="-")
        logger.info("  %-28s: %.12f eV", "Potential energy", potential_energy)
        logger.info("  %-28s: %.12f eV", "Total energy", total_energy)

        log_atomic_reference_method(logger)
        try:
            atomic_sum = Sum_of_atomic_energies(
                context.counts,
                args.mlff_chk,
                args.mlff_task,
                context.device,
                logger=logger,
                predictor=context.predictor,
            )
        except Exception as exc:
            logger.error("Failed to compute sum of atomic energies: %s", exc)
            return 1, None

        bonding_energy = total_energy - atomic_sum
        logger.info("")
        logger.info("  %-28s: %.12f eV", "Sum of atomic energies", atomic_sum)
        logger.info("  %-28s: %.12f eV", "Bonding energy", bonding_energy)
        logger.info("  %-28s: %.12f kcal.mol-1", "Bonding energy", bonding_energy * 23.0609)

        trajectory = read(filename=str(paths.trajectory), index=":")
        rmsd_value = rmsd_simple(trajectory[0], trajectory[-1])

        log_section_header(logger, "Geometry Analysis", char="-")
        logger.info("  %-28s: %.6f Å", "RMSD (initial vs optimized)", rmsd_value)

        write(str(paths.final_geometry), context.atoms)
        log_key_value_block(
            logger,
            "Output Files",
            [
                ("Final geometry", str(paths.final_geometry)),
                ("Log file", str(paths.log)),
                ("Trajectory", str(paths.trajectory)),
            ],
            char="-",
        )

        if args.visualize:
            view(trajectory)

        log_timestamp_event(logger, "End of Geometry Optimization Run")

        return 0, context
    finally:
        log_total_timing(logger, "geoopt", run_wall_start, run_cpu_start)


def run_single_point(
    args: Namespace,
    paths: WorkflowPaths,
    logger: logging.Logger,
) -> Tuple[int, Optional[AtomContext]]:
    """Compute single-point energies and bonding energy."""
    run_wall_start = time.perf_counter()
    run_cpu_start = time.process_time()
    try:
        status, context = setup_calculated_atoms(args, logger)
        if status != 0 or context is None:
            return status, None

        atoms = context.atoms
        log_run_metadata(args, logger, "sp", context.device)
        log_section_header(logger, "Single Point Run")
        log_timestamp_event(logger, "Running Single Point")

        log_geometry(logger, atoms, "Initial geometry for Single Point (Å):")

        potential_energy = atoms.get_potential_energy()
        total_energy_wall_start = time.perf_counter()
        total_energy = atoms.get_total_energy()
        total_energy_wall_elapsed = time.perf_counter() - total_energy_wall_start

        logger.info("Potential Energy: %s eV", potential_energy)
        logger.info("Total Energy:     %s eV", total_energy)
        logger.info(" ")
        log_timestamp_event(
            logger,
            f"Timing [Single Point total energy]: wall={total_energy_wall_elapsed:.3f}s",
            marker="=",
        )

        log_atomic_reference_method(logger)
        try:
            with log_timing(logger, "Single Point atomic energies"):
                atomic_sum = Sum_of_atomic_energies(
                    context.counts,
                    args.mlff_chk,
                    args.mlff_task,
                    context.device,
                    logger=logger,
                    predictor=context.predictor,
                )
        except Exception as exc:
            logger.error("Failed to compute sum of atomic energies: %s", exc)
            return 1, None

        with log_timing(logger, "Single Point bonding energy"):
            bonding_energy = total_energy - atomic_sum
            logger.info("Sum of atomic energies: %s eV", atomic_sum)
            logger.info("Bonding Energy: %s eV", bonding_energy)
            logger.info("Bonding Energy: %s kcal.mol-1", bonding_energy*23.0609)
        log_timestamp_event(logger, "End of Single Point Run")

        return 0, context
    finally:
        log_total_timing(logger, "sp", run_wall_start, run_cpu_start)


def run_vibrational_analysis(
    context: AtomContext,
    logger: logging.Logger,
    base_path: Path,
    temperature: float,
    pressure: float,
    freq_archive: Optional[Path] = None,
    *,
    skip_thermochemistry: bool = False,
    args: Optional[Namespace] = None,
) -> Tuple[int, Optional[VibrationsData]]:
    """Compute vibrational frequencies, thermochemistry (optional), and archive normal modes."""

    run_wall_start = time.perf_counter()
    run_cpu_start = time.process_time()
    try:
        if args is not None:
            log_run_metadata(args, logger, "freqs", context.device)

        start_time = datetime.now()
        log_section_header(logger, "Vibrational Analysis Run")
 
        log_timestamp_event(
            logger,
            f"Starting vibrational analysis on device {context.device} "
            f"(T={temperature:.4f} K, P={pressure:.2f} Pa)",
            moment=start_time,
        )

        log_geometry(logger, context.atoms, "Geometry for Frequencies (Å):")
        freq_root = base_path.parent / "freqs"
        freq_root.mkdir(parents=True, exist_ok=True)
        base_name = base_path.stem
        target_dir = freq_root / base_name
        target_dir.mkdir(parents=True, exist_ok=True)
        cache_dir = target_dir / "cache"
        cache_dir.mkdir(parents=True, exist_ok=True)
        archive_destination = freq_archive or (base_path.parent / f"{base_name}-FREQS-modes.zip")
        freq_packaged = False

        def finalize_modes() -> None:
            nonlocal freq_packaged
            if freq_packaged:
                return
            with log_timing(logger, "Frequency packaging"):
                _cleanup_vibration_cache(cache_dir, target_dir, logger)
                _package_vibrational_data(target_dir, archive_destination, logger)
            freq_packaged = True

        vib_prefix = cache_dir / base_name
        vib = SymmetrizedVibrations(context.atoms, name=str(vib_prefix))
        vib.clean()
        vib_data: Optional[VibrationsData] = None

        original_displacements = vib.displacements

        def wrapped_displacements():
            raw_disp = list(original_displacements())
            total = len(raw_disp)
            start = datetime.now()
            for index, disp in enumerate(raw_disp, start=1):
                if disp.name != "eq":
                    log_timestamp_message(
                        logger,
                        f"Vibrational displacement {index}/{total} ({disp.name}) in progress.",
                    )
                yield disp
            log_timestamp_event(
                logger,
                f"Completed vibrational displacements ({total} entries) in {datetime.now() - start}.",
            )

        vib.displacements = wrapped_displacements  # type: ignore[assignment]

        try:
            with log_timing(logger, "Vibrational displacements"):
                vib.run()
        except Exception as exc:
            logger.error("Failed to compute vibrational modes: %s", exc)
            _cleanup_vibration_cache(cache_dir, target_dir, logger)
            return 1, None
        finally:
            vib.displacements = original_displacements

        skip_thermo_reason = None

        with log_timing(logger, "Vibrational analysis"):
            try:
                vib_data = vib.get_vibrations()
            except Exception as exc:  # pragma: no cover - ASE internals may vary
                logger.warning("Unable to cache VibrationsData for mode inspection: %s", exc)
                vib_data = None

            end_time = datetime.now()
            log_timestamp_event(
                logger,
                f"Vibrational analysis completed. Normal modes stored under {target_dir}",
                moment=end_time,
            )
            if vib_data is not None:
                _log_normal_mode_matrices(vib_data, context.atoms, logger)

            raw_energies: Optional[np.ndarray] = None
            try:
                if vib_data is not None:
                    raw_energies = vib_data.get_energies()
                else:
                    raw_energies = vib.get_energies()
            except ValueError as exc:
                logger.error("Unable to retrieve vibrational energies: %s", exc)
                finalize_modes()
                return 1, vib_data

            try:
                vib_energies = _sanitize_vibrational_energies(raw_energies, logger)
            except ValueError as exc:
                logger.warning("Imaginary vibrational energies are present: %s", exc)
                vib_energies = None
                skip_thermochemistry = True
                skip_thermo_reason = "imaginary vibrational modes detected"

            try:
                _, summary_output = capture_stdout(vib.summary)
                if summary_output:
                    logger.info("Vibrational summary:\n%s", summary_output.rstrip())
                else:
                    logger.info("Vibrational summary: <no output>")
            except Exception as exc:  # pragma: no cover - ASE behaviour varies
                logger.warning("Unable to print vibration summary: %s", exc)

            try:
                if vib_data is not None:
                    _write_all_vibrational_modes(vib_data, target_dir, base_name, logger)
                elif raw_energies is not None:
                    fallback_vib_data = vib.get_vibrations()
                    _write_all_vibrational_modes(fallback_vib_data, target_dir, base_name, logger)
            except Exception as exc:  # pragma: no cover - optional
                logger.warning("Unable to write vibrational modes: %s", exc)
            finally:
                finalize_modes()
        if skip_thermochemistry:
            if skip_thermo_reason:
                logger.info("Thermochemistry analysis skipped (%s).", skip_thermo_reason)
            else:
                logger.info("Thermochemistry analysis skipped for transition-state preparation.")
            log_timestamp_event(logger, "Timing [Thermochemistry]: wall=0.000s (skipped)", marker="=")
            log_timestamp_event(logger, "End of Vibrational Analysis Run")
            return 0, vib_data

        with log_timing(logger, "Thermochemistry"):
            potential_energy = context.atoms.get_potential_energy()
            try:
                thermo = IdealGasThermo(
                    vib_energies=vib_energies,
                    potentialenergy=potential_energy,
                    atoms=context.atoms,
                    geometry="nonlinear",
                    symmetrynumber=1,
                    spin=0,
                )
            except ValueError as exc:
                logger.warning(
                    "Thermochemistry skipped: %s",
                    exc or "Imaginary vibrational energies prevent further analysis.",
                )
                finish_timestamp = datetime.now()
                log_timestamp_event(logger, "End of Vibrational Analysis Run", moment=finish_timestamp)
                finalize_modes()
                return 0, vib_data
            if hasattr(thermo, "summary"):
                try:
                    _, thermo_output = capture_stdout(
                        thermo.summary,
                        temperature=temperature,
                        pressure=pressure,
                        verbose=True,
                    )
                    if thermo_output:
                        logger.info(
                            "Thermochemistry summary (IdealGasThermo):\n%s",
                            thermo_output.rstrip(),
                        )
                except Exception as exc:  # pragma: no cover - optional
                    logger.warning("IdealGasThermo.summary failed: %s", exc)
            else:
                property_specs = []
                recorded_lines = []
                for label, method_name, kwargs, unit in property_specs:
                    method = getattr(thermo, method_name, None)
                    if callable(method):
                        try:
                            value, prop_output = capture_stdout(method, **kwargs)
                        except Exception as exc:  # pragma: no cover - optional
                            logger.warning("Unable to evaluate %s: %s", label.lower(), exc)
                        else:
                            if prop_output:
                                logger.info(prop_output.rstrip())
                            recorded_lines.append(f"  {label}: {value} {unit}")
                if recorded_lines:
                    logger.info("Thermochemistry (IdealGasThermo):\n%s", "\n".join(recorded_lines))

            try:
                gibbs, gibbs_output = capture_stdout(
                    thermo.get_gibbs_energy,
                    temperature=temperature,
                    pressure=pressure,
                )
                if gibbs_output:
                    logger.info("IdealGasThermo main output:\n%s", gibbs_output.rstrip())
            except Exception as exc:
                logger.warning("Unable to compute Gibbs free energy: %s", exc)
                finalize_modes()
                return 1, vib_data
        finish_timestamp = datetime.now()
        log_timestamp_event(logger, "End of Vibrational Analysis Run", moment=finish_timestamp)

        finalize_modes()
        return 0, vib_data
    finally:
        log_total_timing(logger, "freqs", run_wall_start, run_cpu_start)


def _select_lowest_curvature_mode(
    vib_data: VibrationsData,
    logger: logging.Logger,
) -> Tuple[Optional[np.ndarray], Optional[int], Optional[float], Optional[float]]:
    """Return the normalized displacement vector of the softest vibrational mode."""
    try:
        energies, modes = vib_data.get_energies_and_modes()
    except Exception as exc:  # pragma: no cover - ASE internals
        logger.error("Unable to access vibrational modes: %s", exc)
        return None, None, None, None

    if energies is None or len(energies) == 0 or modes is None or len(modes) == 0:
        logger.error("Vibrational calculation did not produce any modes.")
        return None, None, None, None

    energy_array = np.asarray(energies)
    if energy_array.size == 0:
        logger.error("Unable to identify lowest curvature mode.")
        return None, None, None, None

    freq_complex = energy_array / units.invcm if units.invcm else energy_array
    imag_component = np.abs(np.imag(freq_complex))
    real_component = np.real(freq_complex)
    curvature_cm_values = np.where(imag_component > 1e-6, -imag_component, real_component)

    mode_index = int(np.argmin(curvature_cm_values))
    curvature_cm = float(curvature_cm_values[mode_index])
    curvature = curvature_cm * units.invcm if units.invcm else float(energy_array[mode_index].real)

    mode_candidate = np.array(modes)[mode_index]
    mode_candidate = np.array(mode_candidate, dtype=float)
    norm = float(np.linalg.norm(mode_candidate))
    if not np.isfinite(norm) or norm == 0.0:
        logger.error("Selected eigenmode %d has invalid norm.", mode_index)
        return None, None, None, None

    return mode_candidate / norm, mode_index, curvature, curvature_cm


def _sanitize_vibrational_energies(
    energies: Iterable[complex],
    logger: logging.Logger,
    threshold_cm: float = 100.0,
) -> np.ndarray:
    """Convert near-zero imaginary frequencies into positive reals for thermochemistry."""
    array = np.asarray(list(energies), dtype=complex)
    if array.size == 0:
        return np.array([], dtype=float)

    sanitized = np.empty_like(array, dtype=float)
    adjustments: list[tuple[int, float]] = []
    sizable_imag: list[tuple[int, float]] = []
    tol_cm = 1e-3

    for idx, value in enumerate(array):
        freq_cm = value / units.invcm if units.invcm else value
        freq_real = float(np.real(freq_cm))
        freq_imag = float(np.imag(freq_cm))

        if not np.isfinite(freq_imag) or not np.isfinite(freq_real):
            sizable_imag.append((idx, freq_imag))
            continue

        if abs(freq_imag) <= tol_cm:
            sanitized[idx] = abs(freq_real) * units.invcm
            continue

        if abs(freq_imag) <= threshold_cm:
            sanitized[idx] = abs(freq_imag) * units.invcm
            adjustments.append((idx, freq_imag))
        else:
            sanitized[idx] = freq_cm.real * units.invcm
            sizable_imag.append((idx, freq_imag))

    if len(sizable_imag) > 1:
        details = ", ".join(f"mode {idx} ({value:.2f} cm^-1)" for idx, value in sizable_imag)
        raise ValueError(f"Imaginary vibrational modes exceed tolerance: {details}")
    if len(sizable_imag) == 1:
        idx, value = sizable_imag[0]
        logger.info(
            "Single significant imaginary mode retained for TS validation: mode %d (%.2f cm^-1).",
            idx,
            value,
        )

    if adjustments:
        logger.info(
            "Converted %d near-zero imaginary modes (<%.1f cm^-1) to real frequencies: %s",
            len(adjustments),
            threshold_cm,
            ", ".join(f"mode {idx} ({value:.2f} cm^-1)" for idx, value in adjustments),
        )

    return sanitized


def run_transition_state_search(
    args: Namespace,
    paths: WorkflowPaths,
    logger: logging.Logger,
    base_path: Path,
    context: Optional[AtomContext],
) -> Tuple[int, Optional[AtomContext]]:
    """Locate a transition state using a freshly computed vibrational analysis plus dimer search."""
    run_wall_start = time.perf_counter()
    run_cpu_start = time.process_time()
    try:
        if context is None:
            status, context = setup_calculated_atoms(args, logger)
            if status != 0 or context is None:
                return status, None

        log_run_metadata(args, logger, "ts", context.device)

        log_section_header(logger, "Transition State Search Run")
        logger.info("* Plan: (1) preconditioning frequency run (thermochemistry skipped),")
        logger.info("*       (2) dimer-based TS search,")
        logger.info("*       (3) final frequency/thermochemistry validation on the saddle point.")

        log_section_header(logger, "Phase 1: Vibrational Seed Run", char="-")

        with log_timing(logger, "TS phase 1: preconditioning frequencies"):
            freq_status, vib_data = run_vibrational_analysis(
                context,
                logger,
                base_path,
                args.temp,
                args.press,
                paths.freq_archive,
                skip_thermochemistry=True,
                args=args,
            )
        if freq_status != 0:
            return freq_status, None

        if vib_data is None:
            logger.error("Vibrational data unavailable; cannot seed dimer search.")
            return 1, None

        log_section_header(logger, "Phase 2: Dimer Search Run", char="-")
        logger.info(
            "Running dimer TS search with fmax %.4f eV/Å and up to %d steps.",
            args.grad,
            args.iter,
        )

        with log_timing(logger, "TS phase 2: dimer search"):
            eigenmode, mode_index, curvature, curvature_cm = _select_lowest_curvature_mode(
                vib_data, logger
            )
            if eigenmode is None or mode_index is None or curvature is None:
                logger.error("Unable to select a vibrational eigenmode for the TS search.")
                return 1, None

            if curvature_cm is not None:
                logger.info(
                    "Using vibrational mode #%d with curvature %.2f cm^-1 (%.6f eV) to initialize the dimer search.",
                    mode_index,
                    curvature_cm,
                    curvature,
                )
            else:
                logger.info(
                    "Using vibrational mode #%d with curvature %.6f eV to initialize the dimer search.",
                    mode_index,
                    curvature,
                )

            control_log = LoggerWriter(logger)
            dimer_control = DimerControl(
                logfile=control_log,
                eigenmode_logfile=control_log,
                dimer_separation=args.ts_dimer_separation,
                trial_trans_step=args.ts_trial_step,
                maximum_translation=args.ts_max_step,
                max_num_rot=args.ts_max_rot,
            )

            dimer_atoms = MinModeAtoms(context.atoms, control=dimer_control)
            try:
                dimer_atoms.initialize_eigenmodes(eigenmodes=[eigenmode])
            except Exception as exc:
                logger.error("Failed to initialise dimer eigenmode: %s", exc)
                return 1, None

            displacement = getattr(args, "ts_displacement", 0.1)
            if displacement > 0:
                displacement_vector = (eigenmode * displacement).tolist()
                try:
                    full_mask = [True] * len(context.atoms)
                    dimer_atoms.displace(
                        displacement_vector=displacement_vector,
                        method="vector",
                        mask=full_mask,
                    )
                except Exception as exc:
                    logger.warning("Unable to displace along the selected eigenmode: %s", exc)

            optimizer = MinModeTranslate(
                dimer_atoms,
                logfile=LoggerWriter(logger),
                trajectory=str(paths.trajectory),
            )

            try:
                run_result = optimizer.run(fmax=args.grad, steps=args.iter)
            except Exception as exc:
                logger.error("Dimer optimizer failed: %s", exc)
                return 1, None

            finish_time = datetime.now()
            converged = _optimizer_converged(optimizer, run_result)
            final_curvature = float(dimer_atoms.get_curvature())
            if not converged or final_curvature >= 0.0:
                logger.error(
                    "%s TS search stopped without convergence (curvature %.6f eV).",
                    finish_time,
                    final_curvature,
                )
                return 1, None

            log_timestamp_event(logger, "Transition state search converged.", moment=finish_time)
            context.atoms = dimer_atoms.get_atoms()
            logger.info("Final dimer curvature: %.6f eV", final_curvature)
            log_geometry(logger, context.atoms, "Transition state geometry (Å):")

            potential_energy = context.atoms.get_potential_energy()
            total_energy = context.atoms.get_total_energy()
            logger.info("Transition state potential energy: %s eV", potential_energy)
            logger.info("Transition state total energy    : %s eV", total_energy)

            log_atomic_reference_method(logger)
            try:
                atomic_sum = Sum_of_atomic_energies(
                    context.counts,
                    args.mlff_chk,
                    args.mlff_task,
                    context.device,
                    logger=logger,
                    predictor=context.predictor,
                )
            except Exception as exc:
                logger.error("Failed to compute sum of atomic energies: %s", exc)
                return 1, None

            bonding_energy = total_energy - atomic_sum
            logger.info("Sum of atomic energies: %s eV", atomic_sum)
            logger.info("Relative bonding energy: %s eV", bonding_energy)
            logger.info("Relative bonding energy: %s kcal.mol-1", bonding_energy * 23.0609)

            write(str(paths.final_geometry), context.atoms)
            logger.info("Transition state geometry stored in %s", str(paths.final_geometry))
            logger.info("Dimer trajectory written to %s", str(paths.trajectory))

        log_section_header(logger, "Phase 3: Vibrational Validation Run", char="-")
        logger.info("Initiating post-TS vibrational analysis to validate the saddle point.")
        with log_timing(logger, "TS phase 3: final frequencies"):
            freq_dir = base_path.parent / "freqs" / base_path.stem
            cleanup_targets = [freq_dir]
            if paths.freq_archive is not None:
                cleanup_targets.append(paths.freq_archive)
            remove_if_exists(cleanup_targets)

            final_freq_status, _ = run_vibrational_analysis(
                context,
                logger,
                base_path,
                args.temp,
                args.press,
                paths.freq_archive,
                args=args,
            )
        if final_freq_status != 0:
            logger.error("Post-TS vibrational analysis failed.")
            return final_freq_status, None

        log_timestamp_event(logger, "End of Transition State Search Run")

        if args.visualize:
            try:
                frames = read(paths.trajectory, ":")
            except Exception:  # pragma: no cover - optional viewer
                frames = None
            if frames is not None:
                view(frames)

        return 0, context
    finally:
        log_total_timing(logger, "ts", run_wall_start, run_cpu_start)
