#!/usr/bin/env python3
"""xTB implicit-solvent delta correction helpers.

The correction adds only solvent contributions by subtracting xTB vacuum terms
from xTB implicit-solvent terms on the same geometry:
    dE = E(solv) - E(vac)
    dF = F(solv) - F(vac)
    dH = H(solv) - H(vac)
"""

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

import numpy as np

EV_PER_HARTREE = 27.211386245988
ANG_PER_BOHR = 0.529177210903
FORCE_EV_ANG_PER_HA_BOHR = EV_PER_HARTREE / ANG_PER_BOHR
HESS_EV_ANG2_PER_HA_BOHR2 = EV_PER_HARTREE / (ANG_PER_BOHR * ANG_PER_BOHR)

_TOTAL_ENERGY_RE = re.compile(
    r"TOTAL ENERGY\s+([-+]?(?:\d+\.\d*|\d*\.\d+|\d+)(?:[DdEe][-+]?\d+)?)"
)
_INT_TOKEN_RE = re.compile(r"^[+-]?\d+$")
_SOLVENT_MODEL_TO_XTB_FLAG = {
    "alpb": "--alpb",
    "cpcmx": "--cpcmx",
}


class XTBError(RuntimeError):
    """Raised for xTB implicit-solvent correction failures."""


def normalize_solvent_name(solvent):
    text = "" if solvent is None else str(solvent).strip().lower()
    if (not text) or (text in ("none", "vac", "vacuum")):
        return "none"
    return text


def normalize_solvent_model(model):
    text = "" if model is None else str(model).strip().lower()
    if not text:
        return "alpb"
    if text not in _SOLVENT_MODEL_TO_XTB_FLAG:
        raise XTBError(
            "Unsupported --solvent-model '{}'. Choose from: {}.".format(
                model,
                ", ".join(sorted(_SOLVENT_MODEL_TO_XTB_FLAG.keys())),
            )
        )
    return text


def solvent_correction_enabled(solvent):
    return normalize_solvent_name(solvent) != "none"


def resolve_xtb_ncores(explicit_ncores=None, fallback=4):
    candidates = []
    if explicit_ncores is not None:
        candidates.append(explicit_ncores)
    candidates.extend(
        [
            os.getenv("OMP_NUM_THREADS"),
            os.getenv("SLURM_CPUS_PER_TASK"),
            os.getenv("PBS_NP"),
            os.getenv("NSLOTS"),
        ]
    )
    for item in candidates:
        try:
            val = int(item)
            if val > 0:
                return val
        except Exception:
            continue
    return max(1, int(fallback))


def convert_units_xtb_to_mlip(energy_ha=None, gradient_ha_bohr=None, hessian_ha_bohr2=None):
    """Convert xTB atomic units to MLIP units (eV, eV/Ang, eV/Ang^2)."""
    energy_ev = None
    forces_ev_ang = None
    hessian_ev_ang2 = None

    if energy_ha is not None:
        energy_ev = float(energy_ha) * EV_PER_HARTREE
    if gradient_ha_bohr is not None:
        grad = np.asarray(gradient_ha_bohr, dtype=np.float64).reshape(-1, 3)
        forces_ev_ang = -grad * FORCE_EV_ANG_PER_HA_BOHR
    if hessian_ha_bohr2 is not None:
        hess = np.asarray(hessian_ha_bohr2, dtype=np.float64)
        hessian_ev_ang2 = hess * HESS_EV_ANG2_PER_HA_BOHR2

    return energy_ev, forces_ev_ang, hessian_ev_ang2


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


def _all_int_tokens(parts):
    return bool(parts) and all(_INT_TOKEN_RE.match(str(tok)) for tok in parts)


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-alpb-", 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_alpb_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 _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,
    solvent,
    solvent_model,
    xtb_acc,
    mode,
):
    cmd = _xtb_cmd_tokens(xtb_cmd)
    cmd.extend(
        [
            str(xyz_filename),
            "--acc",
            str(float(xtb_acc)),
            "--norestart",
        ]
    )
    m = str(mode).strip().lower()
    if m == "grad":
        cmd.append("--grad")
    elif m == "hess":
        cmd.append("--hess")
    elif m == "sp":
        pass
    else:
        raise XTBError("Unknown xTB mode '{}'".format(mode))

    solvent_name = normalize_solvent_name(solvent)
    if solvent_name != "none":
        model_name = normalize_solvent_model(solvent_model)
        cmd.extend([_SOLVENT_MODEL_TO_XTB_FLAG[model_name], solvent_name])

    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,
    solvent,
    solvent_model,
    xtb_cmd,
    xtb_acc,
    mode,
    ncores,
):
    cmd = _build_xtb_cmd(
        xtb_cmd=xtb_cmd,
        xyz_filename=xyz_filename,
        charge=charge,
        multiplicity=multiplicity,
        solvent=solvent,
        solvent_model=solvent_model,
        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 (mode={}, solvent_model={}, solvent={}) rc={}.\nCMD: {}\nSTDOUT:\n{}\nSTDERR:\n{}".format(
                mode,
                normalize_solvent_model(solvent_model),
                normalize_solvent_name(solvent),
                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_xtb_hessian(path, natoms):
    """Parse xTB `hessian` file.

    Supported formats:
    1) Dense matrix format (xTB >= 6.7 commonly writes this):
         $hessian
         <H11> <H12> ...
         ...
    2) xTB blocked format with dimension/header rows.
    """
    with open(path, "r") as handle:
        lines = [ln.rstrip("\n") for ln in handle]

    ndim_expected = int(natoms) * 3
    start = None
    for i, line in enumerate(lines):
        if line.strip().lower() == "$hessian":
            start = i + 1
            break
    if start is None:
        raise XTBError("Could not locate '$hessian' block in {}".format(path))

    data_lines = []
    i = start
    while i < len(lines):
        text = lines[i].strip()
        if text.startswith("$"):
            break
        if text:
            data_lines.append(text)
        i += 1

    if not data_lines:
        raise XTBError("No Hessian data found after '$hessian' in {}".format(path))

    # Format A: dense matrix values (no explicit dimension line).
    first_parts = data_lines[0].split()
    if not (len(first_parts) == 1 and _INT_TOKEN_RE.match(first_parts[0])):
        values = []
        for row in data_lines:
            for tok in row.split():
                values.append(_parse_float_token(tok))
        expected = ndim_expected * ndim_expected
        if len(values) < expected:
            raise XTBError(
                "Dense Hessian parse failed in {}: expected at least {} values, got {}.".format(
                    path,
                    expected,
                    len(values),
                )
            )
        hess = np.asarray(values[:expected], dtype=np.float64).reshape(
            ndim_expected, ndim_expected
        )
        hess = 0.5 * (hess + hess.T)
        return hess

    # Format B: blocked matrix with explicit dimension line.
    try:
        ndim = int(first_parts[0])
    except Exception as exc:
        raise XTBError("Invalid Hessian dimension line '{}'".format(data_lines[0])) from exc

    if ndim != ndim_expected:
        raise XTBError(
            "xTB Hessian dimension mismatch: expected {}, got {} in {}".format(
                ndim_expected,
                ndim,
                path,
            )
        )

    hess = np.zeros((ndim, ndim), dtype=np.float64)
    i = 1
    saw_block = False

    while i < len(data_lines):
        line = data_lines[i].strip()
        if not line:
            i += 1
            continue

        header = line.split()
        if not _all_int_tokens(header):
            i += 1
            continue

        raw_cols = [int(tok) for tok in header]
        if min(raw_cols) == 0:
            cols = raw_cols
        elif min(raw_cols) == 1:
            cols = [c - 1 for c in raw_cols]
        else:
            raise XTBError("Unexpected Hessian column indices in '{}'".format(line))
        if min(cols) < 0 or max(cols) >= ndim:
            raise XTBError("Hessian column index out of range in '{}'".format(line))

        i += 1
        irow = 0
        while i < len(data_lines) and irow < ndim:
            row_line = data_lines[i].strip()
            if not row_line:
                i += 1
                continue
            row_parts = row_line.split()
            if _all_int_tokens(row_parts):
                break

            if len(row_parts) == (len(cols) + 1) and _INT_TOKEN_RE.match(row_parts[0]):
                data_tokens = row_parts[1:]
            elif len(row_parts) == len(cols):
                data_tokens = row_parts
            else:
                raise XTBError(
                    "Malformed Hessian row in {}: '{}'".format(path, row_line)
                )

            vals = [_parse_float_token(tok) for tok in data_tokens]
            hess[np.asarray(cols, dtype=int), irow] = np.asarray(vals, dtype=np.float64)
            irow += 1
            i += 1

        if irow != ndim:
            raise XTBError(
                "Incomplete Hessian block in {} (read {} rows, expected {}).".format(
                    path,
                    irow,
                    ndim,
                )
            )
        saw_block = True

    if not saw_block:
        raise XTBError("No Hessian matrix blocks parsed from {}".format(path))

    # Numerical noise can break strict symmetry slightly.
    hess = 0.5 * (hess + hess.T)
    return hess


def xtb_energy(
    symbols,
    coords_ang,
    charge,
    multiplicity,
    solvent=None,
    solvent_model="alpb",
    xtb_cmd="xtb",
    xtb_acc=0.2,
    xtb_workdir="tmp",
    xtb_keep_files=False,
    ncores=1,
):
    run_dir = _make_run_dir(xtb_workdir)
    xyz_name = "xtb_input.xyz"
    try:
        _write_xyz(os.path.join(run_dir, xyz_name), symbols, coords_ang)
        stdout = _run_xtb(
            run_dir=run_dir,
            xyz_filename=xyz_name,
            charge=charge,
            multiplicity=multiplicity,
            solvent=solvent,
            solvent_model=solvent_model,
            xtb_cmd=xtb_cmd,
            xtb_acc=xtb_acc,
            mode="sp",
            ncores=ncores,
        )
        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)
    finally:
        _cleanup_run_dir(run_dir, xtb_keep_files)


def xtb_engrad(
    symbols,
    coords_ang,
    charge,
    multiplicity,
    solvent=None,
    solvent_model="alpb",
    xtb_cmd="xtb",
    xtb_acc=0.2,
    xtb_workdir="tmp",
    xtb_keep_files=False,
    ncores=1,
):
    run_dir = _make_run_dir(xtb_workdir)
    xyz_name = "xtb_input.xyz"
    nat = len(symbols)
    try:
        _write_xyz(os.path.join(run_dir, xyz_name), symbols, coords_ang)
        _run_xtb(
            run_dir=run_dir,
            xyz_filename=xyz_name,
            charge=charge,
            multiplicity=multiplicity,
            solvent=solvent,
            solvent_model=solvent_model,
            xtb_cmd=xtb_cmd,
            xtb_acc=xtb_acc,
            mode="grad",
            ncores=ncores,
        )
        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_ha_bohr = _parse_engrad(engrad_path, nat)
        energy_ev, forces_ev_ang, _ = convert_units_xtb_to_mlip(
            energy_ha=energy_ha,
            gradient_ha_bohr=grad_ha_bohr,
        )
        return float(energy_ev), np.asarray(forces_ev_ang, dtype=np.float64)
    finally:
        _cleanup_run_dir(run_dir, xtb_keep_files)


def xtb_hessian(
    symbols,
    coords_ang,
    charge,
    multiplicity,
    solvent=None,
    solvent_model="alpb",
    xtb_cmd="xtb",
    xtb_acc=0.2,
    xtb_workdir="tmp",
    xtb_keep_files=False,
    ncores=1,
):
    run_dir = _make_run_dir(xtb_workdir)
    xyz_name = "xtb_input.xyz"
    try:
        _write_xyz(os.path.join(run_dir, xyz_name), symbols, coords_ang)
        _run_xtb(
            run_dir=run_dir,
            xyz_filename=xyz_name,
            charge=charge,
            multiplicity=multiplicity,
            solvent=solvent,
            solvent_model=solvent_model,
            xtb_cmd=xtb_cmd,
            xtb_acc=xtb_acc,
            mode="hess",
            ncores=ncores,
        )
        hess_path = os.path.join(run_dir, "hessian")
        if not os.path.isfile(hess_path):
            raise XTBError("xTB Hessian file not found: {}".format(hess_path))

        hess_ha_bohr2 = _parse_xtb_hessian(hess_path, len(symbols))
        _, _, hess_ev_ang2 = convert_units_xtb_to_mlip(hessian_ha_bohr2=hess_ha_bohr2)
        return np.asarray(hess_ev_ang2, dtype=np.float64)
    finally:
        _cleanup_run_dir(run_dir, xtb_keep_files)


def delta_alpb_minus_vac(
    symbols,
    coords_ang,
    charge,
    multiplicity,
    solvent,
    need_forces,
    need_hessian,
    solvent_model="alpb",
    xtb_cmd="xtb",
    xtb_acc=0.2,
    xtb_workdir="tmp",
    xtb_keep_files=False,
    ncores=4,
):
    """Return (dE, dF, dH) where dX = X(solv) - X(vacuum) in MLIP units."""
    solvent_name = normalize_solvent_name(solvent)
    nat = len(symbols)
    if solvent_name == "none":
        zero_forces = np.zeros((nat, 3), dtype=np.float64) if need_forces else None
        zero_hessian = np.zeros((3 * nat, 3 * nat), dtype=np.float64) if need_hessian else None
        return 0.0, zero_forces, zero_hessian

    model_name = normalize_solvent_model(solvent_model)

    ncores = resolve_xtb_ncores(ncores)
    common_kwargs = dict(
        symbols=symbols,
        coords_ang=coords_ang,
        charge=charge,
        multiplicity=multiplicity,
        solvent_model=model_name,
        xtb_cmd=xtb_cmd,
        xtb_acc=xtb_acc,
        xtb_workdir=xtb_workdir,
        xtb_keep_files=xtb_keep_files,
        ncores=ncores,
    )

    if need_forces or need_hessian:
        with concurrent.futures.ThreadPoolExecutor(max_workers=2) as pool:
            fut_vac = pool.submit(xtb_engrad, solvent="none", **common_kwargs)
            fut_sol = pool.submit(xtb_engrad, solvent=solvent_name, **common_kwargs)
            vac_e, vac_f = fut_vac.result()
            sol_e, sol_f = fut_sol.result()
    else:
        with concurrent.futures.ThreadPoolExecutor(max_workers=2) as pool:
            fut_vac = pool.submit(xtb_energy, solvent="none", **common_kwargs)
            fut_sol = pool.submit(xtb_energy, solvent=solvent_name, **common_kwargs)
            vac_e = fut_vac.result()
            sol_e = fut_sol.result()
        vac_f = None
        sol_f = None

    vac_h = None
    sol_h = None
    if need_hessian:
        with concurrent.futures.ThreadPoolExecutor(max_workers=2) as pool:
            fut_vac = pool.submit(xtb_hessian, solvent="none", **common_kwargs)
            fut_sol = pool.submit(xtb_hessian, solvent=solvent_name, **common_kwargs)
            vac_h = fut_vac.result()
            sol_h = fut_sol.result()

    de_ev = float(sol_e - vac_e)

    df_ev_ang = None
    if need_forces:
        if vac_f is None or sol_f is None:
            raise XTBError("xTB force delta requested but state force is missing.")
        vac_f = np.asarray(vac_f, dtype=np.float64).reshape(nat, 3)
        sol_f = np.asarray(sol_f, dtype=np.float64).reshape(nat, 3)
        df_ev_ang = sol_f - vac_f

    dh_ev_ang2 = None
    if need_hessian:
        if vac_h is None or sol_h is None:
            raise XTBError("xTB Hessian delta requested but state Hessian is missing.")
        vac_h = np.asarray(vac_h, dtype=np.float64).reshape(3 * nat, 3 * nat)
        sol_h = np.asarray(sol_h, dtype=np.float64).reshape(3 * nat, 3 * nat)
        dh_ev_ang2 = sol_h - vac_h

    return de_ev, df_ev_ang, dh_ev_ang2
