#!/usr/bin/env python3
"""xTB point-charge embedding correction helpers.

The correction term is defined as:
    dE = E(embed) - E(no-embed)
    dF_Q = F_Q(embed) - F_Q(no-embed)
    dF_M = F_M(embed) - 0

where Q are QM atoms passed to MLIP and M are external point charges.
"""

from __future__ import absolute_import, division, print_function

import concurrent.futures
import os
import re
import shlex
import shutil
import subprocess
import tempfile

import numpy as np

try:
    if __package__ in (None, ""):
        from xtb_alpb_correction import (
            XTBError,
            convert_units_xtb_to_mlip,
            resolve_xtb_ncores,
        )
    else:
        from .xtb_alpb_correction import (
            XTBError,
            convert_units_xtb_to_mlip,
            resolve_xtb_ncores,
        )
except Exception as exc:
    raise RuntimeError("Failed to import xTB helpers for embedcharge correction: {}".format(exc))


BOHR_PER_ANG = 1.8897261254578281
_TOTAL_ENERGY_RE = re.compile(
    r"TOTAL ENERGY\s+([-+]?(?:\d+\.\d*|\d*\.\d+|\d+)(?:[DdEe][-+]?\d+)?)"
)


def _parse_float_token(text):
    return float(str(text).replace("D", "E").replace("d", "e"))


def _make_run_dir(xtb_workdir):
    base = None
    mode = str(xtb_workdir or "tmp").strip()
    if mode and mode.lower() != "tmp":
        base = os.path.abspath(mode)
        os.makedirs(base, exist_ok=True)
    return tempfile.mkdtemp(prefix="xtb-embed-", dir=base)


def _cleanup_run_dir(path, keep_files):
    if keep_files:
        return
    try:
        shutil.rmtree(path)
    except Exception:
        pass


def _write_xyz(path, symbols, coords_ang):
    coords = np.asarray(coords_ang, dtype=np.float64).reshape(-1, 3)
    if len(symbols) != coords.shape[0]:
        raise XTBError(
            "XYZ write failed: symbol count {} != coord count {}".format(
                len(symbols),
                coords.shape[0],
            )
        )
    with open(path, "w") as handle:
        handle.write("{}\n".format(coords.shape[0]))
        handle.write("generated by xtb_embedcharge_correction\n")
        for sym, xyz in zip(symbols, coords):
            handle.write(
                "{:<2s} {: .12f} {: .12f} {: .12f}\n".format(
                    str(sym),
                    float(xyz[0]),
                    float(xyz[1]),
                    float(xyz[2]),
                )
            )


def _write_pcharge(path, mm_coords_ang, mm_charges):
    coords = np.asarray(mm_coords_ang, dtype=np.float64).reshape(-1, 3)
    charges = np.asarray(mm_charges, dtype=np.float64).reshape(-1)
    if coords.shape[0] != charges.shape[0]:
        raise XTBError(
            "Point-charge write failed: coord count {} != charge count {}".format(
                coords.shape[0],
                charges.shape[0],
            )
        )
    coords_bohr = coords * BOHR_PER_ANG
    with open(path, "w") as handle:
        handle.write("{}\n".format(coords_bohr.shape[0]))
        for q, xyz in zip(charges, coords_bohr):
            handle.write(
                "{: .12f} {: .12f} {: .12f} {: .12f}\n".format(
                    float(q),
                    float(xyz[0]),
                    float(xyz[1]),
                    float(xyz[2]),
                )
            )


def _xtb_cmd_tokens(xtb_cmd):
    tokens = shlex.split(str(xtb_cmd or "xtb"))
    if not tokens:
        raise XTBError("Empty --xtb-cmd value.")
    return tokens


def _build_xtb_cmd(
    xtb_cmd,
    xyz_filename,
    charge,
    multiplicity,
    xtb_acc,
    mode,
    use_restart=False,
):
    cmd = _xtb_cmd_tokens(xtb_cmd)
    cmd.extend(
        [
            str(xyz_filename),
            "--acc",
            str(float(xtb_acc)),
        ]
    )
    if not use_restart:
        cmd.append("--norestart")
    m = str(mode).strip().lower()
    if m == "grad":
        cmd.append("--grad")
    elif m == "sp":
        pass
    else:
        raise XTBError("Unknown xTB mode '{}'".format(mode))

    if charge is not None:
        cmd.extend(["--chrg", str(int(charge))])
    if multiplicity is not None:
        cmd.extend(["--uhf", str(max(0, int(multiplicity) - 1))])

    return cmd


def _run_xtb(
    run_dir,
    xyz_filename,
    charge,
    multiplicity,
    xtb_cmd,
    xtb_acc,
    mode,
    ncores,
):
    cmd = _build_xtb_cmd(
        xtb_cmd=xtb_cmd,
        xyz_filename=xyz_filename,
        charge=charge,
        multiplicity=multiplicity,
        xtb_acc=xtb_acc,
        mode=mode,
    )
    env = os.environ.copy()
    env["OMP_NUM_THREADS"] = str(resolve_xtb_ncores(ncores))

    try:
        proc = subprocess.run(
            cmd,
            cwd=run_dir,
            env=env,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True,
            check=False,
        )
    except FileNotFoundError as exc:
        raise XTBError(
            "xTB command not found: '{}'. Set --xtb-cmd correctly.".format(xtb_cmd)
        ) from exc
    except Exception as exc:
        raise XTBError("Failed to run xTB command '{}': {}".format(" ".join(cmd), exc))

    if proc.returncode != 0:
        out_tail = "\n".join((proc.stdout or "").splitlines()[-20:])
        err_tail = "\n".join((proc.stderr or "").splitlines()[-20:])
        raise XTBError(
            "xTB failed (embedcharge mode={}, rc={}).\nCMD: {}\nSTDOUT:\n{}\nSTDERR:\n{}".format(
                mode,
                proc.returncode,
                " ".join(cmd),
                out_tail,
                err_tail,
            )
        )

    return proc.stdout or ""


def _parse_energy_from_stdout(stdout_text):
    matches = _TOTAL_ENERGY_RE.findall(str(stdout_text or ""))
    if not matches:
        return None
    return _parse_float_token(matches[-1])


def _parse_engrad(path, natoms):
    with open(path, "r") as handle:
        lines = handle.readlines()

    n = len(lines)
    energy_ha = None
    grad_vals = []

    i = 0
    while i < n:
        if "the current total energy in eh" in lines[i].strip().lower():
            i += 1
            while i < n:
                text = lines[i].strip()
                if text and (not text.startswith("#")):
                    energy_ha = _parse_float_token(text.split()[0])
                    break
                i += 1
            break
        i += 1

    grad_start = None
    for idx, line in enumerate(lines):
        if "the current gradient in eh/bohr" in line.strip().lower():
            grad_start = idx + 1
            break

    if grad_start is not None:
        j = grad_start
        while j < n and len(grad_vals) < (3 * int(natoms)):
            text = lines[j].strip()
            if text and (not text.startswith("#")):
                grad_vals.append(_parse_float_token(text.split()[0]))
            j += 1

    if energy_ha is None:
        raise XTBError("Could not parse energy from engrad file: {}".format(path))
    if len(grad_vals) != 3 * int(natoms):
        raise XTBError(
            "Could not parse full gradient from engrad (expected {}, got {}) in {}".format(
                3 * int(natoms),
                len(grad_vals),
                path,
            )
        )

    grad = np.asarray(grad_vals, dtype=np.float64).reshape(int(natoms), 3)
    return float(energy_ha), grad


def _parse_pcgrad(path, ncharges):
    if int(ncharges) <= 0:
        return np.zeros((0, 3), dtype=np.float64)
    if not os.path.isfile(path):
        raise XTBError("xTB point-charge gradient file not found: {}".format(path))

    rows = []
    with open(path, "r") as handle:
        for line in handle:
            text = line.strip()
            if not text:
                continue
            parts = text.split()
            if len(parts) < 3:
                raise XTBError("Malformed pcgrad row in {}: '{}'".format(path, text))
            rows.append([_parse_float_token(parts[0]), _parse_float_token(parts[1]), _parse_float_token(parts[2])])

    if len(rows) != int(ncharges):
        raise XTBError(
            "pcgrad row count mismatch: expected {}, got {} in {}".format(
                int(ncharges),
                len(rows),
                path,
            )
        )

    return np.asarray(rows, dtype=np.float64).reshape(int(ncharges), 3)


# ---------------------------------------------------------------------------
# Persistent xTB worker (inspired by tblite Calculator pattern)
# ---------------------------------------------------------------------------


class XTBWorker(object):
    """Persistent xTB CLI worker that reuses a working directory and restart files.

    Mirrors tblite's Calculator pattern: init once, update inputs, evaluate.
    The xTB restart file from the previous step is used as the SCF initial
    guess, significantly reducing iteration count for MD trajectories.
    """

    def __init__(self, role, xtb_cmd, xtb_acc, ncores,
                 xtb_workdir="tmp", xtb_keep_files=False):
        self._role = str(role)
        self._xtb_cmd = str(xtb_cmd or "xtb")
        self._xtb_acc = float(xtb_acc)
        self._ncores = ncores
        self._xtb_workdir = xtb_workdir
        self._xtb_keep_files = bool(xtb_keep_files)

        self._run_dir = None
        self._xyz_path = None
        self._pcharge_path = None
        self._env = None
        self._initialized = False
        self._step_count = 0
        self._last_charge = None
        self._last_multiplicity = None
        self._last_mode = None
        self._cmd_cache = None

    def _ensure_initialized(self):
        if self._initialized:
            return
        base_dir = str(self._xtb_workdir or "tmp").strip()
        if base_dir.lower() == "tmp":
            base_dir = tempfile.gettempdir()
        else:
            base_dir = os.path.abspath(base_dir)
        os.makedirs(base_dir, exist_ok=True)

        self._run_dir = os.path.join(base_dir, "xtb-persistent-{}".format(self._role))
        os.makedirs(self._run_dir, exist_ok=True)

        self._xyz_path = os.path.join(self._run_dir, "xtb_input.xyz")
        self._pcharge_path = os.path.join(self._run_dir, "pcharge")

        self._env = os.environ.copy()
        self._env["OMP_NUM_THREADS"] = str(resolve_xtb_ncores(self._ncores))

        self._initialized = True

    def _get_cmd(self, charge, multiplicity, mode):
        """Return cached command list, rebuilding only when parameters change."""
        if (self._cmd_cache is not None
                and charge == self._last_charge
                and multiplicity == self._last_multiplicity
                and mode == self._last_mode):
            return self._cmd_cache

        # Parameters changed — rebuild command and invalidate restart file.
        if (self._last_charge is not None
                and (charge != self._last_charge
                     or multiplicity != self._last_multiplicity)):
            self.reset()

        self._cmd_cache = _build_xtb_cmd(
            xtb_cmd=self._xtb_cmd,
            xyz_filename="xtb_input.xyz",
            charge=charge,
            multiplicity=multiplicity,
            xtb_acc=self._xtb_acc,
            mode=mode,
            use_restart=True,
        )
        self._last_charge = charge
        self._last_multiplicity = multiplicity
        self._last_mode = mode
        return self._cmd_cache

    def _run_subprocess(self, cmd):
        try:
            proc = subprocess.run(
                cmd,
                cwd=self._run_dir,
                env=self._env,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                text=True,
                check=False,
            )
        except FileNotFoundError as exc:
            raise XTBError(
                "xTB command not found: '{}'. Set --xtb-cmd correctly.".format(
                    self._xtb_cmd)
            ) from exc
        except Exception as exc:
            raise XTBError(
                "Failed to run xTB command '{}': {}".format(" ".join(cmd), exc))

        if proc.returncode != 0:
            out_tail = "\n".join((proc.stdout or "").splitlines()[-20:])
            err_tail = "\n".join((proc.stderr or "").splitlines()[-20:])
            raise XTBError(
                "xTB failed (worker={}, rc={}).\nCMD: {}\nSTDOUT:\n{}\nSTDERR:\n{}".format(
                    self._role, proc.returncode, " ".join(cmd), out_tail, err_tail,
                )
            )
        return proc.stdout or ""

    def evaluate(self, symbols, coords_q_ang, mode, with_embedding,
                 mm_coords_ang, mm_charges, charge, multiplicity):
        """Write inputs, run xTB, parse outputs.  Returns (energy_ev, fq, fm)."""
        self._ensure_initialized()

        nm = int(np.asarray(mm_charges, dtype=np.float64).reshape(-1).size)
        cmd = self._get_cmd(charge, multiplicity, mode)

        # Overwrite input files in persistent directory.
        _write_xyz(self._xyz_path, symbols, coords_q_ang)
        if with_embedding and nm > 0:
            _write_pcharge(self._pcharge_path, mm_coords_ang, mm_charges)
        else:
            # Remove stale pcharge file so xTB does not pick it up.
            if os.path.isfile(self._pcharge_path):
                os.remove(self._pcharge_path)

        # Run xTB; retry once with restart reset on failure.
        try:
            stdout = self._run_subprocess(cmd)
        except XTBError:
            if self._step_count > 0:
                self.reset()
                stdout = self._run_subprocess(cmd)
            else:
                raise

        # Parse output.
        m = str(mode).strip().lower()
        if m == "sp":
            energy_ha = _parse_energy_from_stdout(stdout)
            if energy_ha is None:
                raise XTBError("Could not parse xTB energy from stdout.")
            energy_ev, _, _ = convert_units_xtb_to_mlip(energy_ha=energy_ha)
            self._step_count += 1
            return float(energy_ev), None, None

        xyz_stem = os.path.splitext("xtb_input.xyz")[0]
        engrad_path = os.path.join(self._run_dir, "{}.engrad".format(xyz_stem))
        if not os.path.isfile(engrad_path):
            raise XTBError("xTB engrad file not found: {}".format(engrad_path))

        energy_ha, grad_q_ha_bohr = _parse_engrad(engrad_path, len(symbols))
        energy_ev, forces_q_ev_ang, _ = convert_units_xtb_to_mlip(
            energy_ha=energy_ha, gradient_ha_bohr=grad_q_ha_bohr,
        )

        forces_m_ev_ang = np.zeros((nm, 3), dtype=np.float64)
        if with_embedding and nm > 0:
            pcgrad_ha_bohr = _parse_pcgrad(
                os.path.join(self._run_dir, "pcgrad"), nm)
            _, forces_m_ev_ang, _ = convert_units_xtb_to_mlip(
                gradient_ha_bohr=pcgrad_ha_bohr)

        self._step_count += 1
        return (
            float(energy_ev),
            np.asarray(forces_q_ev_ang, dtype=np.float64).reshape(len(symbols), 3),
            np.asarray(forces_m_ev_ang, dtype=np.float64).reshape(nm, 3),
        )

    def reset(self):
        """Delete restart file to force fresh SCF on next evaluate()."""
        if self._run_dir is None:
            return
        restart = os.path.join(self._run_dir, "xtbrestart")
        if os.path.isfile(restart):
            try:
                os.remove(restart)
            except OSError:
                pass

    def shutdown(self):
        """Clean up the persistent working directory."""
        if self._run_dir is not None and not self._xtb_keep_files:
            try:
                shutil.rmtree(self._run_dir)
            except Exception:
                pass
        self._run_dir = None
        self._initialized = False


class EmbedchargeWorkerPool(object):
    """Manages persistent embed + vacuum xTB workers with a shared thread pool.

    Created once at server startup, reused for the entire simulation.
    """

    def __init__(self, xtb_cmd="xtb", xtb_acc=0.2, xtb_workdir="tmp",
                 xtb_keep_files=False, ncores=4):
        self._embed_worker = XTBWorker(
            role="embed", xtb_cmd=xtb_cmd, xtb_acc=xtb_acc,
            ncores=ncores, xtb_workdir=xtb_workdir,
            xtb_keep_files=xtb_keep_files,
        )
        self._vacuum_worker = XTBWorker(
            role="vacuum", xtb_cmd=xtb_cmd, xtb_acc=xtb_acc,
            ncores=ncores, xtb_workdir=xtb_workdir,
            xtb_keep_files=xtb_keep_files,
        )
        self._executor = concurrent.futures.ThreadPoolExecutor(max_workers=2)

    def delta_embed_minus_vac(self, symbols, coords_q_ang, mm_coords_ang,
                              mm_charges, charge, multiplicity, need_forces,
                              ncores):
        """Compute dE = E(embed) - E(vacuum) using persistent workers."""
        nm = int(np.asarray(mm_charges, dtype=np.float64).reshape(-1).size)
        if nm <= 0:
            zq = np.zeros((len(symbols), 3), dtype=np.float64) if need_forces else None
            zm = np.zeros((0, 3), dtype=np.float64) if need_forces else None
            return 0.0, zq, zm

        mode = "grad" if need_forces else "sp"

        fut_embed = self._executor.submit(
            self._embed_worker.evaluate,
            symbols=symbols, coords_q_ang=coords_q_ang, mode=mode,
            with_embedding=True, mm_coords_ang=mm_coords_ang,
            mm_charges=mm_charges, charge=charge, multiplicity=multiplicity,
        )
        fut_vac = self._executor.submit(
            self._vacuum_worker.evaluate,
            symbols=symbols, coords_q_ang=coords_q_ang, mode=mode,
            with_embedding=False,
            mm_coords_ang=np.zeros((0, 3), dtype=np.float64),
            mm_charges=np.zeros((0,), dtype=np.float64),
            charge=charge, multiplicity=multiplicity,
        )

        embed_e, embed_fq, embed_fm = fut_embed.result()
        vac_e, vac_fq, _ = fut_vac.result()

        de_ev = float(embed_e - vac_e)
        if not need_forces:
            return de_ev, None, None

        df_q = np.asarray(embed_fq, dtype=np.float64).reshape(len(symbols), 3) - \
            np.asarray(vac_fq, dtype=np.float64).reshape(len(symbols), 3)
        df_m = np.asarray(embed_fm, dtype=np.float64).reshape(nm, 3)
        return de_ev, df_q, df_m

    def shutdown(self):
        """Shut down workers and thread pool."""
        if self._executor is not None:
            self._executor.shutdown(wait=True)
            self._executor = None
        for worker in (self._embed_worker, self._vacuum_worker):
            if worker is not None:
                try:
                    worker.shutdown()
                except Exception:
                    pass


def _run_xtb_state(
    symbols,
    coords_q_ang,
    charge,
    multiplicity,
    mode,
    with_embedding,
    mm_coords_ang,
    mm_charges,
    xtb_cmd,
    xtb_acc,
    xtb_workdir,
    xtb_keep_files,
    ncores,
):
    run_dir = _make_run_dir(xtb_workdir)
    xyz_name = "xtb_input.xyz"
    nm = int(np.asarray(mm_charges, dtype=np.float64).reshape(-1).size)
    try:
        _write_xyz(os.path.join(run_dir, xyz_name), symbols, coords_q_ang)
        if with_embedding and nm > 0:
            _write_pcharge(
                os.path.join(run_dir, "pcharge"),
                mm_coords_ang=mm_coords_ang,
                mm_charges=mm_charges,
            )

        stdout = _run_xtb(
            run_dir=run_dir,
            xyz_filename=xyz_name,
            charge=charge,
            multiplicity=multiplicity,
            xtb_cmd=xtb_cmd,
            xtb_acc=xtb_acc,
            mode=mode,
            ncores=ncores,
        )

        if str(mode).strip().lower() == "sp":
            energy_ha = _parse_energy_from_stdout(stdout)
            if energy_ha is None:
                raise XTBError("Could not parse xTB energy from stdout.")
            energy_ev, _, _ = convert_units_xtb_to_mlip(energy_ha=energy_ha)
            return float(energy_ev), None, None

        engrad_path = os.path.join(run_dir, "{}.engrad".format(os.path.splitext(xyz_name)[0]))
        if not os.path.isfile(engrad_path):
            raise XTBError("xTB engrad file not found: {}".format(engrad_path))
        energy_ha, grad_q_ha_bohr = _parse_engrad(engrad_path, len(symbols))
        energy_ev, forces_q_ev_ang, _ = convert_units_xtb_to_mlip(
            energy_ha=energy_ha,
            gradient_ha_bohr=grad_q_ha_bohr,
        )

        forces_m_ev_ang = np.zeros((nm, 3), dtype=np.float64)
        if with_embedding and nm > 0:
            pcgrad_ha_bohr = _parse_pcgrad(os.path.join(run_dir, "pcgrad"), nm)
            _, forces_m_ev_ang, _ = convert_units_xtb_to_mlip(
                gradient_ha_bohr=pcgrad_ha_bohr
            )

        return (
            float(energy_ev),
            np.asarray(forces_q_ev_ang, dtype=np.float64).reshape(len(symbols), 3),
            np.asarray(forces_m_ev_ang, dtype=np.float64).reshape(nm, 3),
        )
    finally:
        _cleanup_run_dir(run_dir, xtb_keep_files)


def _delta_embed_minus_vac(
    symbols,
    coords_q_ang,
    mm_coords_ang,
    mm_charges,
    charge,
    multiplicity,
    need_forces,
    xtb_cmd,
    xtb_acc,
    xtb_workdir,
    xtb_keep_files,
    ncores,
):
    nm = int(np.asarray(mm_charges, dtype=np.float64).reshape(-1).size)
    if nm <= 0:
        zq = np.zeros((len(symbols), 3), dtype=np.float64) if need_forces else None
        zm = np.zeros((0, 3), dtype=np.float64) if need_forces else None
        return 0.0, zq, zm

    mode = "grad" if need_forces else "sp"

    common_kw = dict(
        symbols=symbols,
        coords_q_ang=coords_q_ang,
        charge=charge,
        multiplicity=multiplicity,
        mode=mode,
        xtb_cmd=xtb_cmd,
        xtb_acc=xtb_acc,
        xtb_workdir=xtb_workdir,
        xtb_keep_files=xtb_keep_files,
        ncores=ncores,
    )

    # Run embed and vacuum in parallel (they are independent).
    with concurrent.futures.ThreadPoolExecutor(max_workers=2) as pool:
        fut_embed = pool.submit(
            _run_xtb_state,
            with_embedding=True,
            mm_coords_ang=mm_coords_ang,
            mm_charges=mm_charges,
            **common_kw
        )
        fut_vac = pool.submit(
            _run_xtb_state,
            with_embedding=False,
            mm_coords_ang=np.zeros((0, 3), dtype=np.float64),
            mm_charges=np.zeros((0,), dtype=np.float64),
            **common_kw
        )
        embed_e, embed_fq, embed_fm = fut_embed.result()
        vac_e, vac_fq, _ = fut_vac.result()

    de_ev = float(embed_e - vac_e)
    if not need_forces:
        return de_ev, None, None

    df_q_ev_ang = np.asarray(embed_fq, dtype=np.float64).reshape(len(symbols), 3) - np.asarray(
        vac_fq, dtype=np.float64
    ).reshape(len(symbols), 3)
    df_m_ev_ang = np.asarray(embed_fm, dtype=np.float64).reshape(nm, 3)
    return de_ev, df_q_ev_ang, df_m_ev_ang


def _assemble_full_force(df_q_ev_ang, df_m_ev_ang):
    fq = np.asarray(df_q_ev_ang, dtype=np.float64).reshape(-1, 3)
    fm = np.asarray(df_m_ev_ang, dtype=np.float64).reshape(-1, 3)
    if fm.shape[0] == 0:
        return fq
    return np.vstack([fq, fm])


def _numerical_hessian_from_forces(
    eval_full_forces,
    coords_q_ang,
    mm_coords_ang,
    step_ang,
    force0_ev_ang=None,
):
    q0 = np.asarray(coords_q_ang, dtype=np.float64).reshape(-1, 3)
    m0 = np.asarray(mm_coords_ang, dtype=np.float64).reshape(-1, 3)
    nq = q0.shape[0]
    nm = m0.shape[0]
    nall = nq + nm
    dof = nall * 3
    step = float(step_ang)
    if step <= 0.0:
        raise XTBError("Embedcharge Hessian step must be positive, got {}".format(step_ang))

    if force0_ev_ang is None:
        f0 = np.asarray(eval_full_forces(q0, m0), dtype=np.float64).reshape(nall, 3)
    else:
        f0 = np.asarray(force0_ev_ang, dtype=np.float64).reshape(nall, 3)

    hess = np.zeros((dof, dof), dtype=np.float64)
    nqdof = 3 * nq
    for k in range(dof):
        qp = q0.copy()
        qm = q0.copy()
        mp = m0.copy()
        mm = m0.copy()

        if k < nqdof:
            a = k // 3
            c = k % 3
            qp[a, c] += step
            qm[a, c] -= step
        else:
            kk = k - nqdof
            a = kk // 3
            c = kk % 3
            mp[a, c] += step
            mm[a, c] -= step

        f_plus = np.asarray(eval_full_forces(qp, mp), dtype=np.float64).reshape(-1)
        f_minus = np.asarray(eval_full_forces(qm, mm), dtype=np.float64).reshape(-1)
        g_plus = -f_plus
        g_minus = -f_minus
        hess[:, k] = (g_plus - g_minus) / (2.0 * step)

    hess = 0.5 * (hess + hess.T)
    return f0, hess


def delta_embedcharge_minus_noembed(
    symbols,
    coords_q_ang,
    mm_coords_ang,
    mm_charges,
    charge,
    multiplicity,
    need_forces,
    need_hessian,
    xtb_cmd="xtb",
    xtb_acc=0.2,
    xtb_workdir="tmp",
    xtb_keep_files=False,
    ncores=4,
    hessian_step=1.0e-3,
):
    """Return xTB point-charge embedding correction in MLIP units.

    Returns
    -------
    tuple
        ``(dE_ev, dF_full_ev_ang, dH_full_ev_ang2)``, where full index order is
        concatenated as ``[Q atoms, MM point charges]``.
    """
    q_symbols = list(symbols or [])
    q_coords = np.asarray(coords_q_ang, dtype=np.float64).reshape(-1, 3)
    m_coords = np.asarray(mm_coords_ang, dtype=np.float64).reshape(-1, 3)
    m_charges = np.asarray(mm_charges, dtype=np.float64).reshape(-1)

    if len(q_symbols) != q_coords.shape[0]:
        raise XTBError(
            "Embedcharge correction got symbol/coord mismatch: {} vs {}".format(
                len(q_symbols),
                q_coords.shape[0],
            )
        )
    if m_coords.shape[0] != m_charges.shape[0]:
        raise XTBError(
            "Embedcharge correction got MM coord/charge mismatch: {} vs {}".format(
                m_coords.shape[0],
                m_charges.shape[0],
            )
        )

    nq = q_coords.shape[0]
    nm = m_coords.shape[0]
    nall = nq + nm
    if nm <= 0:
        zero_force = np.zeros((nall, 3), dtype=np.float64) if (need_forces or need_hessian) else None
        zero_hess = np.zeros((3 * nall, 3 * nall), dtype=np.float64) if need_hessian else None
        return 0.0, zero_force, zero_hess

    # xTB warns that point-charge Hessian runs can fail in parallel due to I/O.
    run_ncores = resolve_xtb_ncores(ncores)
    if need_hessian:
        run_ncores = 1

    req_force = bool(need_forces) or bool(need_hessian)
    de_ev, df_q, df_m = _delta_embed_minus_vac(
        symbols=q_symbols,
        coords_q_ang=q_coords,
        mm_coords_ang=m_coords,
        mm_charges=m_charges,
        charge=charge,
        multiplicity=multiplicity,
        need_forces=req_force,
        xtb_cmd=xtb_cmd,
        xtb_acc=xtb_acc,
        xtb_workdir=xtb_workdir,
        xtb_keep_files=xtb_keep_files,
        ncores=run_ncores,
    )

    df_full = None
    if req_force:
        df_full = _assemble_full_force(df_q, df_m)

    dh_full = None
    if need_hessian:
        step = float(hessian_step)

        def _eval_full_forces(cur_q, cur_m):
            _de, _fq, _fm = _delta_embed_minus_vac(
                symbols=q_symbols,
                coords_q_ang=cur_q,
                mm_coords_ang=cur_m,
                mm_charges=m_charges,
                charge=charge,
                multiplicity=multiplicity,
                need_forces=True,
                xtb_cmd=xtb_cmd,
                xtb_acc=xtb_acc,
                xtb_workdir=xtb_workdir,
                xtb_keep_files=xtb_keep_files,
                ncores=run_ncores,
            )
            return _assemble_full_force(_fq, _fm)

        _f0, dh_full = _numerical_hessian_from_forces(
            eval_full_forces=_eval_full_forces,
            coords_q_ang=q_coords,
            mm_coords_ang=m_coords,
            step_ang=step,
            force0_ev_ang=df_full,
        )

    return float(de_ev), df_full, dh_full

