from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path
from types import SimpleNamespace
from typing import Any, Dict, List, Optional

import logging
import time

from ase import units
from ase.io import read, write, Trajectory
from ase.md import Langevin, VelocityVerlet
from ase.md.nvtberendsen import NVTBerendsen
from ase.md.nptberendsen import NPTBerendsen
from ase.md.velocitydistribution import MaxwellBoltzmannDistribution, Stationary, ZeroRotation

try:
    from ase.md.bussi import Bussi
except ModuleNotFoundError:  # pragma: no cover
    Bussi = None

try:
    from ase.md.andersen import Andersen
except ModuleNotFoundError:  # pragma: no cover
    Andersen = None

try:
    from ase.md.nose_hoover_chain import NoseHooverChainNVT, IsotropicMTKNPT, MTKNPT
except ModuleNotFoundError:  # pragma: no cover
    NoseHooverChainNVT = None
    IsotropicMTKNPT = None
    MTKNPT = None

try:
    from ase.md.langevinbaoab import LangevinBAOAB
except ModuleNotFoundError:  # pragma: no cover
    LangevinBAOAB = None

try:
    from ase.md.melchionna import MelchionnaNPT
except ModuleNotFoundError:  # pragma: no cover
    MelchionnaNPT = None

from ase.geometry import cellpar_to_cell

from .workflows import (
    log_key_value_block,
    log_section_header,
    log_timing,
    log_total_timing,
    log_timestamp_event,
    setup_calculated_atoms,
)

MD_ENGINE_LABELS = {
    "langevin": "Langevin (NVT)",
    "velocity_verlet": "VelocityVerlet (NVE)",
    "nvt_berendsen": "NVT Berendsen",
    "nose_hoover_chain": "Nose-Hoover Chain (NVT)",
    "bussi": "Bussi (NVT)",
    "andersen": "Andersen (NVT)",
    "langevin_baoab_nvt": "Langevin-Hoover BAOAB (NVT)",
    "npt_berendsen": "NPT Berendsen",
    "isotropic_mtk_npt": "Isotropic MTK (NPT)",
    "mtk_npt": "Full MTK (NPT)",
    "langevin_baoab_npt": "Langevin-Hoover BAOAB (NPT)",
    "melchionna_npt": "Melchionna NPT",
}


def _missing_engine_builder(engine_key: str, module_hint: str):
    def _builder(*_):
        raise RuntimeError(
            f"The '{engine_key}' MD engine requires {module_hint}, but it is not available in the installed ASE. "
            "Update ASE or choose a different engine."
        )

    return _builder


def _build_md_builder_params(md_options: Dict[str, Any]) -> SimpleNamespace:
    temperature = md_options.get("temperature", 300.0)
    timestep = md_options.get("timestep_fs", 0.5) * units.fs
    friction_coeff = max(0.0, md_options.get("friction", 0.0)) / units.fs
    relax_fs = max(1e-12, md_options.get("relax_fs", 100.0))
    relax_time = relax_fs * units.fs
    pressure_bar = md_options.get("pressure_bar", 1.0)
    compressibility_bar = md_options.get("compressibility_bar", 0.0)
    compressibility_au = compressibility_bar / units.bar if compressibility_bar and compressibility_bar > 0 else None
    barostat_relax = max(1e-12, md_options.get("barostat_relax_fs", 1000.0)) * units.fs
    nose_tdamp = max(1e-12, md_options.get("nose_tdamp_fs", relax_fs)) * units.fs
    bussi_taut = max(1e-12, md_options.get("bussi_taut_fs", relax_fs)) * units.fs
    baoab_ttau = max(1e-12, md_options.get("baoab_ttau_fs", 50.0)) * units.fs
    baoab_ptau = max(1e-12, md_options.get("baoab_ptau_fs", 1000.0)) * units.fs
    melchionna_ttime = max(0.0, md_options.get("melchionna_ttime_fs", 25.0)) * units.fs
    melchionna_ptime = max(0.0, md_options.get("melchionna_ptime_fs", 75.0)) * units.fs
    return SimpleNamespace(
        temperature=temperature,
        timestep=timestep,
        friction=friction_coeff,
        relax=relax_time,
        pressure_bar=pressure_bar,
        pressure_au=pressure_bar * units.bar,
        compressibility_au=compressibility_au,
        barostat_relax=barostat_relax,
        nose_tdamp=nose_tdamp,
        nose_tchain=int(md_options.get("nose_tchain", 3)),
        nose_tloop=int(md_options.get("nose_tloop", 1)),
        bussi_taut=bussi_taut,
        andersen_prob=float(md_options.get("andersen_prob", 0.0)),
        mtk_tdamp=max(1e-12, md_options.get("mtk_tdamp_fs", 100.0)) * units.fs,
        mtk_pdamp=max(1e-12, md_options.get("mtk_pdamp_fs", 1000.0)) * units.fs,
        mtk_tchain=int(md_options.get("mtk_tchain", 3)),
        mtk_pchain=int(md_options.get("mtk_pchain", 3)),
        mtk_tloop=int(md_options.get("mtk_tloop", 1)),
        mtk_ploop=int(md_options.get("mtk_ploop", 1)),
        baoab_ttau=baoab_ttau,
        baoab_ptau=baoab_ptau,
        baoab_pmass=md_options.get("baoab_pmass"),
        baoab_pmass_factor=md_options.get("baoab_pmass_factor", 1.0),
        baoab_hydrostatic=bool(md_options.get("baoab_hydrostatic")),
        melchionna_ttime=melchionna_ttime if melchionna_ttime > 0 else None,
        melchionna_ptime=melchionna_ptime if melchionna_ptime > 0 else None,
        melchionna_bulk=md_options.get("melchionna_bulk_gpa", 0.0) * units.GPa,
        melchionna_mask=md_options.get("melchionna_mask"),
    )


MD_ENGINE_BUILDERS = {
    "langevin": lambda atoms, p: Langevin(
        atoms,
        p.timestep,
        temperature_K=p.temperature,
        friction=p.friction,
    ),
    "velocity_verlet": lambda atoms, p: VelocityVerlet(atoms, p.timestep),
    "nvt_berendsen": lambda atoms, p: NVTBerendsen(
        atoms,
        p.timestep,
        temperature_K=p.temperature,
        taut=p.relax,
    ),
}

MD_ENGINE_BUILDERS["bussi"] = (
    (lambda atoms, p: Bussi(atoms, p.timestep, temperature_K=p.temperature, taut=p.bussi_taut))
    if Bussi is not None
    else _missing_engine_builder("bussi", "ase.md.bussi.Bussi")
)

MD_ENGINE_BUILDERS["andersen"] = (
    (lambda atoms, p: Andersen(atoms, p.timestep, temperature_K=p.temperature, andersen_prob=p.andersen_prob))
    if Andersen is not None
    else _missing_engine_builder("andersen", "ase.md.andersen.Andersen")
)

MD_ENGINE_BUILDERS["npt_berendsen"] = (
    (
        lambda atoms, p: NPTBerendsen(
            atoms,
            p.timestep,
            temperature_K=p.temperature,
            pressure_au=p.pressure_au,
            taut=p.relax,
            taup=p.barostat_relax,
            compressibility_au=p.compressibility_au,
        )
    )
    if NPTBerendsen is not None
    else _missing_engine_builder("npt_berendsen", "ase.md.nptberendsen.NPTBerendsen")
)

MD_ENGINE_BUILDERS["nose_hoover_chain"] = (
    (
        lambda atoms, p: NoseHooverChainNVT(
            atoms,
            p.timestep,
            temperature_K=p.temperature,
            tdamp=p.nose_tdamp,
            tchain=p.nose_tchain,
            tloop=p.nose_tloop,
        )
    )
    if NoseHooverChainNVT is not None
    else _missing_engine_builder("nose_hoover_chain", "ase.md.nose_hoover_chain.NoseHooverChainNVT")
)

MD_ENGINE_BUILDERS["isotropic_mtk_npt"] = (
    (
        lambda atoms, p: IsotropicMTKNPT(
            atoms,
            p.timestep,
            temperature_K=p.temperature,
            pressure_au=p.pressure_au,
            tdamp=p.mtk_tdamp,
            pdamp=p.mtk_pdamp,
            tchain=p.mtk_tchain,
            pchain=p.mtk_pchain,
            tloop=p.mtk_tloop,
            ploop=p.mtk_ploop,
        )
    )
    if IsotropicMTKNPT is not None
    else _missing_engine_builder("isotropic_mtk_npt", "ase.md.nose_hoover_chain.IsotropicMTKNPT")
)

MD_ENGINE_BUILDERS["mtk_npt"] = (
    (
        lambda atoms, p: MTKNPT(
            atoms,
            p.timestep,
            temperature_K=p.temperature,
            pressure_au=p.pressure_au,
            tdamp=p.mtk_tdamp,
            pdamp=p.mtk_pdamp,
            tchain=p.mtk_tchain,
            pchain=p.mtk_pchain,
            tloop=p.mtk_tloop,
            ploop=p.mtk_ploop,
        )
    )
    if MTKNPT is not None
    else _missing_engine_builder("mtk_npt", "ase.md.nose_hoover_chain.MTKNPT")
)

MD_ENGINE_BUILDERS["langevin_baoab_nvt"] = (
    (
        lambda atoms, p: LangevinBAOAB(
            atoms,
            p.timestep,
            temperature_K=p.temperature,
            T_tau=p.baoab_ttau,
        )
    )
    if LangevinBAOAB is not None
    else _missing_engine_builder("langevin_baoab_nvt", "ase.md.langevinbaoab.LangevinBAOAB")
)

MD_ENGINE_BUILDERS["langevin_baoab_npt"] = (
    (
        lambda atoms, p: LangevinBAOAB(
            atoms,
            p.timestep,
            temperature_K=p.temperature,
            externalstress=-p.pressure_au,
            hydrostatic=p.baoab_hydrostatic,
            T_tau=p.baoab_ttau,
            P_tau=p.baoab_ptau,
            P_mass=p.baoab_pmass if p.baoab_pmass else None,
            P_mass_factor=p.baoab_pmass_factor,
        )
    )
    if LangevinBAOAB is not None
    else _missing_engine_builder("langevin_baoab_npt", "ase.md.langevinbaoab.LangevinBAOAB")
)

MD_ENGINE_BUILDERS["melchionna_npt"] = (
    (
        lambda atoms, p: MelchionnaNPT(
            atoms,
            p.timestep,
            temperature_K=p.temperature,
            externalstress=-p.pressure_au,
            ttime=p.melchionna_ttime,
            pfactor=(p.melchionna_ptime**2 * (p.melchionna_bulk or 0.0)) if p.melchionna_ptime else None,
            mask=p.melchionna_mask,
        )
    )
    if MelchionnaNPT is not None
    else _missing_engine_builder("melchionna_npt", "ase.md.melchionna.MelchionnaNPT")
)


def create_md_dynamics(engine: str, atoms, *, md_options: Dict[str, Any]):
    params = _build_md_builder_params(md_options)
    builder = MD_ENGINE_BUILDERS.get(engine)
    if builder is None:
        raise ValueError(
            "Unsupported MD engine "
            f"'{engine}'. Available: {', '.join(sorted(MD_ENGINE_LABELS))}"
        )
    return builder(atoms, params)


def initialise_md_velocities(atoms, temperature: float, logger: logging.Logger) -> None:
    if temperature <= 0:
        logger.info("Temperature %.2f K <= 0. Skipping velocity initialisation.", temperature)
        return
    MaxwellBoltzmannDistribution(atoms, temperature_K=temperature)
    Stationary(atoms)
    ZeroRotation(atoms)
    logger.info("Initial velocities sampled at %.2f K.", temperature)


def build_md_metadata(
    *,
    md_options: Dict[str, Any],
    engine: str,
    device: str,
    charge: int,
    spin: int,
    use_pbc: bool,
    cell_parameters: Optional[List[float]],
) -> Dict[str, Any]:
    def add_value(metadata: Dict[str, Any], key: str, value: Any) -> None:
        if value is None:
            return
        metadata[key] = value

    checkpoint = md_options.get("mlff_checkpoint")
    task = md_options.get("mlff_task")

    metadata: Dict[str, Any] = {
        "run_type": "md",
        "engine": engine,
        "steps": md_options.get("steps"),
        "timestep_fs": md_options.get("timestep_fs"),
        "traj_interval": md_options.get("traj_interval"),
        "log_interval": md_options.get("log_interval"),
        "checkpoint": checkpoint,
        "method": "MLIP-UMA",
        "task": task,
        "device": device,
        "charge": charge,
        "spin": spin,
        "use_pbc": use_pbc,
    }
    use_input_velocities = bool(md_options.get("use_input_velocities"))
    if engine != "velocity_verlet" or not use_input_velocities:
        add_value(metadata, "temperature", md_options.get("temperature"))
    if use_input_velocities:
        add_value(metadata, "restart", True)

    if use_pbc and cell_parameters and len(cell_parameters) == 6:
        add_value(metadata, "cell_a", cell_parameters[0])
        add_value(metadata, "cell_b", cell_parameters[1])
        add_value(metadata, "cell_c", cell_parameters[2])
        add_value(metadata, "cell_alpha", cell_parameters[3])
        add_value(metadata, "cell_beta", cell_parameters[4])
        add_value(metadata, "cell_gamma", cell_parameters[5])

    if engine == "langevin":
        add_value(metadata, "friction", md_options.get("friction"))

    if engine in {"nvt_berendsen", "npt_berendsen"}:
        add_value(metadata, "relax_fs", md_options.get("relax_fs"))

    if engine in {"npt_berendsen", "isotropic_mtk_npt", "mtk_npt", "langevin_baoab_npt", "melchionna_npt"}:
        add_value(metadata, "pressure_bar", md_options.get("pressure_bar"))

    if engine == "npt_berendsen":
        add_value(metadata, "barostat_relax_fs", md_options.get("barostat_relax_fs"))
        add_value(metadata, "compressibility_bar", md_options.get("compressibility_bar"))

    if engine == "nose_hoover_chain":
        add_value(metadata, "nose_tdamp_fs", md_options.get("nose_tdamp_fs"))
        add_value(metadata, "nose_tchain", md_options.get("nose_tchain"))
        add_value(metadata, "nose_tloop", md_options.get("nose_tloop"))

    if engine == "bussi":
        add_value(metadata, "bussi_taut_fs", md_options.get("bussi_taut_fs"))

    if engine == "andersen":
        add_value(metadata, "andersen_prob", md_options.get("andersen_prob"))

    if engine in {"isotropic_mtk_npt", "mtk_npt"}:
        add_value(metadata, "mtk_tdamp_fs", md_options.get("mtk_tdamp_fs"))
        add_value(metadata, "mtk_pdamp_fs", md_options.get("mtk_pdamp_fs"))
        add_value(metadata, "mtk_tchain", md_options.get("mtk_tchain"))
        add_value(metadata, "mtk_pchain", md_options.get("mtk_pchain"))
        add_value(metadata, "mtk_tloop", md_options.get("mtk_tloop"))
        add_value(metadata, "mtk_ploop", md_options.get("mtk_ploop"))

    if engine in {"langevin_baoab_nvt", "langevin_baoab_npt"}:
        add_value(metadata, "baoab_ttau_fs", md_options.get("baoab_ttau_fs"))
        if engine == "langevin_baoab_npt":
            add_value(metadata, "baoab_ptau_fs", md_options.get("baoab_ptau_fs"))
            add_value(metadata, "baoab_pmass", md_options.get("baoab_pmass"))
            add_value(metadata, "baoab_pmass_factor", md_options.get("baoab_pmass_factor"))
            add_value(metadata, "baoab_hydrostatic", md_options.get("baoab_hydrostatic"))

    if engine == "melchionna_npt":
        add_value(metadata, "melchionna_ttime_fs", md_options.get("melchionna_ttime_fs"))
        add_value(metadata, "melchionna_ptime_fs", md_options.get("melchionna_ptime_fs"))
        add_value(metadata, "melchionna_bulk_gpa", md_options.get("melchionna_bulk_gpa"))
        add_value(metadata, "melchionna_mask", md_options.get("melchionna_mask"))

    return metadata


def log_md_metadata(logger: logging.Logger, metadata: Dict[str, Any]) -> None:
    ordered_keys = [
        "run_type",
        "engine",
        "steps",
        "timestep_fs",
        "temperature",
        "traj_interval",
        "log_interval",
        "checkpoint",
        "method",
        "task",
        "charge",
        "spin",
        "device",
        "use_pbc",
        "restart",
        "cell_a",
        "cell_b",
        "cell_c",
        "cell_alpha",
        "cell_beta",
        "cell_gamma",
        "pressure_bar",
        "compressibility_bar",
        "barostat_relax_fs",
        "relax_fs",
        "friction",
        "nose_tdamp_fs",
        "nose_tchain",
        "nose_tloop",
        "bussi_taut_fs",
        "andersen_prob",
        "mtk_tdamp_fs",
        "mtk_pdamp_fs",
        "mtk_tchain",
        "mtk_pchain",
        "mtk_tloop",
        "mtk_ploop",
        "baoab_ttau_fs",
        "baoab_ptau_fs",
        "baoab_pmass",
        "baoab_pmass_factor",
        "baoab_hydrostatic",
        "melchionna_ttime_fs",
        "melchionna_ptime_fs",
        "melchionna_bulk_gpa",
        "melchionna_mask",
    ]
    rows: list[tuple[str, Any]] = []
    for key in ordered_keys:
        if key in metadata:
            rows.append((key, metadata[key]))
    for key in sorted(metadata):
        if key in ordered_keys:
            continue
        rows.append((key, metadata[key]))
    log_key_value_block(logger, "Molecular Dynamics Metadata", rows, char="-")


def load_restart_velocities(traj_path: Path, logger: Optional[logging.Logger] = None) -> List[List[float]]:
    if logger is not None:
        with log_timing(logger, "ASE restart trajectory load"):
            atoms = read(str(traj_path), index=-1)
    else:
        atoms = read(str(traj_path), index=-1)
    velocities = atoms.get_velocities()
    if velocities is None:
        raise ValueError("No velocities found in trajectory.")
    return velocities.tolist()


@dataclass
class MdPaths:
    log: Path
    traj: Path
    final_snapshot: Path
    multi_xyz: Path


def build_md_paths(input_path: Path) -> MdPaths:
    parent = input_path.parent
    stem = input_path.stem
    return MdPaths(
        log=parent / f"{stem}-MD.log",
        traj=parent / f"{stem}-MD.traj",
        final_snapshot=parent / f"{stem}-MD-final.xyz",
        multi_xyz=parent / f"{stem}-MD.xyz",
    )


def build_md_options_from_args(args, logger: Optional[logging.Logger] = None) -> Dict[str, Any]:
    melchionna_mask: Optional[List[int]] = None
    melchionna_mask_raw = (getattr(args, "md_melchionna_mask", "") or "").strip()
    if melchionna_mask_raw:
        mask_values: List[int] = []
        for entry in melchionna_mask_raw.split(","):
            entry_clean = entry.strip()
            if entry_clean == "":
                continue
            value = 1 if int(round(float(entry_clean))) > 0 else 0
            mask_values.append(value)
        if mask_values:
            mask_values = (mask_values + [1, 1, 1])[:3]
            melchionna_mask = mask_values

    use_pbc = bool(getattr(args, "md_use_pbc", False))
    cell_parameters: Optional[List[float]] = None
    if use_pbc:
        cell_a = getattr(args, "md_cell_a", None)
        cell_b = getattr(args, "md_cell_b", None)
        cell_c = getattr(args, "md_cell_c", None)
        cell_alpha = getattr(args, "md_cell_alpha", None)
        cell_beta = getattr(args, "md_cell_beta", None)
        cell_gamma = getattr(args, "md_cell_gamma", None)
        if None in (cell_a, cell_b, cell_c, cell_alpha, cell_beta, cell_gamma):
            raise ValueError("All cell parameters are required when --md-use-pbc is set.")
        cell_parameters = [
            float(cell_a),
            float(cell_b),
            float(cell_c),
            float(cell_alpha),
            float(cell_beta),
            float(cell_gamma),
        ]

    use_input_velocities = bool(getattr(args, "md_use_input_velocities", False))
    input_velocities = getattr(args, "md_velocities", None)
    restart_traj = getattr(args, "md_restart_traj", None)
    if input_velocities:
        if not isinstance(input_velocities, list):
            raise ValueError("MD velocities must be provided as a list.")
        use_input_velocities = True
    elif restart_traj:
        input_velocities = load_restart_velocities(Path(restart_traj), logger=logger)
        use_input_velocities = True
    if use_input_velocities and input_velocities is None:
        raise ValueError("MD velocities required when md_use_input_velocities is set.")

    andersen_prob = float(getattr(args, "md_andersen_prob", 0.0))
    if not 0.0 <= andersen_prob <= 1.0:
        raise ValueError("Andersen collision probability must be between 0 and 1.")

    return {
        "engine": getattr(args, "md_engine", "langevin"),
        "steps": getattr(args, "md_steps", 0),
        "timestep_fs": getattr(args, "md_timestep_fs", 0.5),
        "temperature": getattr(args, "md_temperature", 300.0),
        "friction": getattr(args, "md_friction", 0.0),
        "traj_interval": getattr(args, "md_traj_interval", 10),
        "log_interval": getattr(args, "md_log_interval", 10),
        "relax_fs": getattr(args, "md_relax_fs", 100.0),
        "pressure_bar": getattr(args, "md_pressure_bar", 1.0),
        "compressibility_bar": getattr(args, "md_compressibility_bar", 0.0),
        "barostat_relax_fs": getattr(args, "md_barostat_relax_fs", 1000.0),
        "nose_tdamp_fs": getattr(args, "md_nose_tdamp_fs", 100.0),
        "nose_tchain": getattr(args, "md_nose_tchain", 3),
        "nose_tloop": getattr(args, "md_nose_tloop", 1),
        "bussi_taut_fs": getattr(args, "md_bussi_taut_fs", 100.0),
        "andersen_prob": andersen_prob,
        "mtk_tdamp_fs": getattr(args, "md_mtk_tdamp_fs", 100.0),
        "mtk_pdamp_fs": getattr(args, "md_mtk_pdamp_fs", 1000.0),
        "mtk_tchain": getattr(args, "md_mtk_tchain", 3),
        "mtk_pchain": getattr(args, "md_mtk_pchain", 3),
        "mtk_tloop": getattr(args, "md_mtk_tloop", 1),
        "mtk_ploop": getattr(args, "md_mtk_ploop", 1),
        "baoab_ttau_fs": getattr(args, "md_baoab_ttau_fs", 50.0),
        "baoab_ptau_fs": getattr(args, "md_baoab_ptau_fs", 1000.0),
        "baoab_pmass": getattr(args, "md_baoab_pmass", None),
        "baoab_pmass_factor": getattr(args, "md_baoab_pmass_factor", 1.0),
        "baoab_hydrostatic": bool(getattr(args, "md_baoab_hydrostatic", False)),
        "melchionna_ttime_fs": getattr(args, "md_melchionna_ttime_fs", 25.0),
        "melchionna_ptime_fs": getattr(args, "md_melchionna_ptime_fs", 75.0),
        "melchionna_bulk_gpa": getattr(args, "md_melchionna_bulk_gpa", 100.0),
        "melchionna_mask": melchionna_mask,
        "use_pbc": use_pbc,
        "cell_parameters": cell_parameters,
        "mlff_checkpoint": getattr(args, "mlff_chk", None),
        "mlff_task": getattr(args, "mlff_task", None),
        "use_input_velocities": use_input_velocities,
        "input_velocities": input_velocities,
    }


def run_md_simulation(
    args,
    logger: logging.Logger,
    *,
    md_options: Dict[str, Any],
    paths: MdPaths,
) -> int:
    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 or 1

        md_steps = int(md_options.get("steps", 0))
        traj_interval = max(1, int(md_options.get("traj_interval", 10)))
        log_interval = max(1, int(md_options.get("log_interval", 10)))
        engine = md_options.get("engine", "langevin")
        label = MD_ENGINE_LABELS.get(engine, engine)
        temperature = md_options.get("temperature", 300.0)
        timestep_fs = md_options.get("timestep_fs", 0.5)
        friction = md_options.get("friction", 0.0)
        relax_fs = md_options.get("relax_fs", 100.0)
        use_pbc = bool(md_options.get("use_pbc"))
        cell_parameters = md_options.get("cell_parameters")
        use_input_velocities = bool(md_options.get("use_input_velocities"))
        input_velocities = md_options.get("input_velocities")

        metadata = build_md_metadata(
            md_options=md_options,
            engine=engine,
            device=context.device,
            charge=int(args.chg),
            spin=int(args.spin),
            use_pbc=use_pbc,
            cell_parameters=cell_parameters,
        )
        log_md_metadata(logger, metadata)

        log_section_header(logger, "Molecular Dynamics Run")
        show_temperature = engine != "velocity_verlet" or not use_input_velocities
        md_summary: list[tuple[str, Any]] = [
            ("Engine", label),
            ("Steps", md_steps),
            ("Timestep (fs)", f"{timestep_fs:.4f}"),
        ]
        if show_temperature:
            md_summary.append(("Temperature (K)", f"{temperature:.2f}"))
        if engine == "langevin":
            md_summary.append(("Friction (1/fs)", f"{friction:.4f}"))
        md_summary.extend([
            ("Traj interval", traj_interval),
            ("Log interval", log_interval),
            ("Checkpoint", md_options.get("mlff_checkpoint")),
            ("Task", md_options.get("mlff_task")),
        ])
        if use_input_velocities:
            md_summary.append(("Velocities", "restart (input)"))
        md_summary.append(("Periodic BC", "enabled" if use_pbc else "disabled"))
        if use_pbc and cell_parameters:
            md_summary.append((
                "Cell (Å | °)",
                "a=%.3f b=%.3f c=%.3f | α=%.2f β=%.2f γ=%.2f"
                % (
                    cell_parameters[0],
                    cell_parameters[1],
                    cell_parameters[2],
                    cell_parameters[3],
                    cell_parameters[4],
                    cell_parameters[5],
                ),
            ))
        if engine in {"npt_berendsen", "isotropic_mtk_npt", "mtk_npt", "langevin_baoab_npt", "melchionna_npt"}:
            md_summary.append(("Pressure (bar)", f"{md_options.get('pressure_bar', 0.0):.6f}"))
        if engine in {"nvt_berendsen", "npt_berendsen"}:
            md_summary.append(("Thermostat tau", f"{relax_fs:.4f} fs"))
        if engine == "npt_berendsen":
            md_summary.extend([
                ("Barostat tau", f"{md_options.get('barostat_relax_fs', 1000.0):.4f} fs"),
                ("Compressibility", f"{md_options.get('compressibility_bar', 0.0):.6e} 1/bar"),
            ])
        if engine == "nose_hoover_chain":
            md_summary.append((
                "Nose-Hoover",
                "tdamp=%.4f fs | chain=%d | tloop=%d"
                % (
                    md_options.get("nose_tdamp_fs", relax_fs),
                    md_options.get("nose_tchain", 3),
                    md_options.get("nose_tloop", 1),
                ),
            ))
        if engine == "bussi":
            md_summary.append(("Bussi taut", f"{md_options.get('bussi_taut_fs', relax_fs):.4f} fs"))
        if engine == "andersen":
            md_summary.append(("Andersen prob", f"{md_options.get('andersen_prob', 0.0):.6f}"))
        if engine in {"isotropic_mtk_npt", "mtk_npt"}:
            md_summary.append((
                "MTK params",
                "tdamp=%.4f fs | pdamp=%.4f fs | tchain=%d | pchain=%d"
                % (
                    md_options.get("mtk_tdamp_fs", 100.0),
                    md_options.get("mtk_pdamp_fs", 1000.0),
                    md_options.get("mtk_tchain", 3),
                    md_options.get("mtk_pchain", 3),
                ),
            ))
        if engine == "langevin_baoab_nvt":
            md_summary.append(("BAOAB params", f"T_tau={md_options.get('baoab_ttau_fs', 50.0):.4f} fs"))
        if engine == "langevin_baoab_npt":
            md_summary.append((
                "BAOAB params",
                "T_tau=%.4f fs | P_tau=%.4f fs | hydrostatic=%s"
                % (
                    md_options.get("baoab_ttau_fs", 50.0),
                    md_options.get("baoab_ptau_fs", 1000.0),
                    md_options.get("baoab_hydrostatic", False),
                ),
            ))
        if engine == "melchionna_npt":
            md_summary.append((
                "Melchionna",
                "ttime=%.4f fs | ptime=%.4f fs | bulk=%.3f GPa"
                % (
                    md_options.get("melchionna_ttime_fs", 25.0),
                    md_options.get("melchionna_ptime_fs", 75.0),
                    md_options.get("melchionna_bulk_gpa", 100.0),
                ),
            ))
        log_key_value_block(logger, "Molecular Dynamics Parameters", md_summary, char="-")
        log_timestamp_event(logger, f"Starting {label} dynamics on device {context.device}")

        atoms = context.atoms
        if use_pbc:
            if cell_parameters:
                cell_matrix = cellpar_to_cell(cell_parameters)
                atoms.set_cell(cell_matrix, scale_atoms=True)
            atoms.set_pbc(True)
            cell_vectors = atoms.get_cell()
            log_section_header(logger, "Cell matrix (Å)", char="-")
            axis_labels = ("a", "b", "c")
            for axis_label, vector in zip(axis_labels, cell_vectors):
                logger.info("  %s = [%.6f %.6f %.6f]", axis_label, vector[0], vector[1], vector[2])
            logger.info("")
        else:
            atoms.set_pbc(False)

        velocities_loaded = False
        if use_input_velocities and input_velocities:
            if len(input_velocities) != len(atoms):
                raise ValueError("Velocity count does not match atoms count.")
            atoms.set_velocities(input_velocities)
            velocities_loaded = True
            logger.info("Input velocities loaded for MD restart.")
        if not velocities_loaded:
            initialise_md_velocities(atoms, temperature, logger)
        else:
            logger.info("MD restart: using velocities from previous trajectory.")
            logger.info("Velocity initialisation skipped (restart).")

        dynamics = create_md_dynamics(engine, atoms, md_options=md_options)
        trajectory = Trajectory(str(paths.traj), "w", atoms)

        def log_step():
            step = dynamics.nsteps
            potential = atoms.get_potential_energy()
            kinetic = atoms.get_kinetic_energy()
            total = potential + kinetic
            inst_temp = atoms.get_temperature()
            logger.info(
                "Step %d/%d | E_pot=%.6f eV | E_kin=%.6f eV | E_tot=%.6f eV | T=%.2f K",
                step,
                md_steps,
                potential,
                kinetic,
                total,
                inst_temp,
            )

        dynamics.attach(log_step, interval=log_interval)
        dynamics.attach(trajectory.write, interval=traj_interval)

        def record_xyz_frame():
            step = dynamics.nsteps
            try:
                potential = atoms.get_potential_energy()
            except Exception:
                potential = float("nan")
            time_fs = step * timestep_fs
            comment = f"step={step} time_fs={time_fs:.6f} energy_ev={potential:.8f}"
            write(
                paths.multi_xyz,
                atoms,
                format="xyz",
                append=True,
                comment=comment,
            )

        dynamics.attach(record_xyz_frame, interval=traj_interval)

        logger.info(
            "Starting %s dynamics for %d steps (dt=%.4f fs).",
            label,
            md_steps,
            timestep_fs,
        )
        with log_timing(logger, "MD integration"):
            chunk = max(1, min(100, traj_interval, log_interval))
            completed = 0
            while completed < md_steps:
                steps_this_round = min(chunk, md_steps - completed)
                dynamics.run(steps_this_round)
                completed += steps_this_round
        log_timestamp_event(logger, "Molecular dynamics finished successfully.")

        write(
            str(paths.final_snapshot),
            atoms,
            format="xyz",
            comment=f"MD final snapshot | charge={int(args.chg)} spin={int(args.spin)}",
        )
        logger.info("Final files saved to %s", paths.final_snapshot)
        log_timestamp_event(logger, "End of Molecular Dynamics Run")

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