#!/usr/bin/env python3
"""Non-MPI qchem shim for AMBER EXTERN (&qc) QM/MM.

AMBER calls this executable as:
    qchem <inpfile> <logfile> <savfile>

The shim reads backend/ml_keywords from environment variables:
- AMBER_MLIPS_BACKEND     : uma|orb|mace|aimnet2
- AMBER_MLIPS_ML_KEYWORDS : backend options, same syntax as ml_keywords
- AMBER_MLIPS_DEBUG       : 1 enables verbose stderr logs

It parses Q-Chem-style input generated by AMBER (&qc), evaluates MLIP,
and writes outputs expected by qm2_extern_qc_module.F90:
- <logfile>: must contain line beginning with " SCF   energy"
- efield.dat:
  - first nclatoms rows: MM electric field vectors (au)
  - next  nqmatoms rows: QM gradients (Eh/Bohr)
"""

from __future__ import absolute_import, division, print_function

import argparse
import os
import shlex
import sys

import numpy as np

from .mlip_backends import (
    AIMNet2Evaluator,
    BackendError,
    MACEEvaluator,
    OrbMolEvaluator,
    UMAEvaluator,
    ev_to_ha,
    forces_ev_ang_to_gradient_ha_bohr,
)
from .mlip_server import ServerError, client_evaluate
from .xtb_alpb_correction import XTBError
from .xtb_embedcharge_correction import delta_embedcharge_minus_noembed


class QCShimError(RuntimeError):
    """Raised for qchem shim runtime failures."""


class _BuiltinEvaluator(object):
    """Small dependency-free evaluator used for smoke/integration tests.

    Model spec:
    - builtin:zero
    - builtin:harmonic
    - builtin:harmonic:<k>
    """

    def __init__(self, model_spec):
        spec = str(model_spec or "builtin:zero").strip().lower()
        self.kind = "zero"
        self.k = 0.05

        if spec.startswith("builtin:harmonic"):
            self.kind = "harmonic"
            parts = spec.split(":")
            if len(parts) >= 3:
                try:
                    self.k = float(parts[2])
                except Exception:
                    self.k = 0.05
            if self.k <= 0.0:
                self.k = 0.05
        elif spec.startswith("builtin:zero"):
            self.kind = "zero"
        else:
            self.kind = "zero"

    def evaluate(
        self,
        symbols,
        coords_ang,
        charge,
        multiplicity,
        need_forces,
        need_hessian,
        hessian_mode,
        hessian_step,
    ):
        _ = (
            symbols,
            charge,
            multiplicity,
            need_forces,
            need_hessian,
            hessian_mode,
            hessian_step,
        )
        coords = np.asarray(coords_ang, dtype=np.float64).reshape(-1, 3)
        nat = coords.shape[0]
        if self.kind == "harmonic":
            # E = 1/2 k r^2, F = -k r
            energy_ev = 0.5 * self.k * float(np.sum(coords * coords))
            forces = -self.k * coords
        else:
            energy_ev = 0.0
            forces = np.zeros((nat, 3), dtype=np.float64)
        return float(energy_ev), np.asarray(forces, dtype=np.float64), None


def _normalize_symbol(sym_text):
    txt = str(sym_text or "").strip()
    if not txt:
        return "X"
    if len(txt) == 1:
        return txt.upper()
    return txt[0].upper() + txt[1:].lower()


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


def _read_sections(inpfile):
    sections = {}
    current = None

    with open(inpfile, "r") as handle:
        for raw in handle:
            line = raw.strip()
            if not line:
                continue

            low = line.lower()
            if low.startswith("$"):
                name = low[1:].split()[0]
                if name == "end":
                    current = None
                else:
                    current = name
                    sections.setdefault(current, [])
                continue

            if current is not None:
                sections[current].append(line)

    return sections


def _parse_qchem_input(inpfile):
    sec = _read_sections(inpfile)

    mol = sec.get("molecule")
    if not mol:
        raise QCShimError("Input does not contain a $molecule section: {}".format(inpfile))

    if len(mol) < 2:
        raise QCShimError("$molecule section is incomplete: {}".format(inpfile))

    header = mol[0].split()
    if len(header) < 2:
        raise QCShimError("Failed to parse charge/multiplicity from $molecule header.")

    charge = int(_f64(header[0]))
    spinmult = int(_f64(header[1]))

    symbols = []
    coords_q = []
    for line in mol[1:]:
        parts = line.split()
        if len(parts) < 4:
            continue
        symbols.append(_normalize_symbol(parts[0]))
        coords_q.append([_f64(parts[1]), _f64(parts[2]), _f64(parts[3])])

    if not symbols:
        raise QCShimError("No QM atoms parsed from $molecule section.")

    ext = sec.get("external_charges", [])
    coords_m = []
    charges_m = []
    for line in ext:
        parts = line.split()
        if len(parts) < 4:
            continue
        coords_m.append([_f64(parts[0]), _f64(parts[1]), _f64(parts[2])])
        charges_m.append(_f64(parts[3]))

    return (
        charge,
        spinmult,
        symbols,
        np.asarray(coords_q, dtype=np.float64),
        np.asarray(coords_m, dtype=np.float64).reshape(-1, 3),
        np.asarray(charges_m, dtype=np.float64).reshape(-1),
    )


def _build_keyword_parser():
    parser = argparse.ArgumentParser(prog="amber-mlips-qchem-keywords", add_help=False)
    parser.add_argument("--model", default=None, help="Model alias/path.")
    parser.add_argument("--device", default="auto", help="cpu|cuda|auto")
    parser.add_argument("--embedcharge", action="store_true", help="Enable xTB point-charge embedding correction.")
    parser.add_argument("--xtb-cmd", default="xtb", help="xTB executable for --embedcharge correction.")
    parser.add_argument("--xtb-acc", type=float, default=0.2, help="xTB --acc value for embedcharge correction.")
    parser.add_argument("--xtb-workdir", default="tmp", help="xTB scratch base dir for embedcharge correction.")
    parser.add_argument("--xtb-keep-files", action="store_true", help="Keep xTB scratch directories for debugging.")
    parser.add_argument("--xtb-ncores", type=int, default=4, help="xTB OMP threads for embedcharge correction.")

    # UMA
    parser.add_argument("--uma-task", default="omol")
    parser.add_argument("--uma-workers", type=int, default=1)

    # ORB
    parser.add_argument("--orb-precision", default="float32")
    parser.add_argument("--orb-compile", action="store_true")

    # MACE
    parser.add_argument("--mace-default-dtype", default="float32")

    parser.add_argument("--debug", action="store_true")
    return parser


def _parse_keywords(backend, ml_keywords):
    parser = _build_keyword_parser()
    tokens = shlex.split(str(ml_keywords or ""))
    args, unknown = parser.parse_known_args(tokens)
    if unknown:
        raise QCShimError("Unknown ml_keywords option(s): {}".format(" ".join(unknown)))
    args.backend = str(backend).strip().lower()
    return args


def _create_evaluator(args):
    backend = str(args.backend).lower()
    model = str(args.model or "").strip()

    if model.lower().startswith("builtin:"):
        return _BuiltinEvaluator(model)

    if backend == "uma":
        return UMAEvaluator(
            model=model or "uma-s-1p1",
            task=args.uma_task,
            device=args.device,
            workers=max(1, int(args.uma_workers)),
            workers_per_node=None,
            max_neigh=None,
            radius=None,
            r_edges=False,
            otf_graph=True,
        )

    if backend == "orb":
        return OrbMolEvaluator(
            model=model or "orb-v3-conservative-omol",
            device=args.device,
            precision=args.orb_precision,
            compile_model=bool(args.orb_compile),
            loader_kwargs=None,
            calc_kwargs=None,
        )

    if backend == "mace":
        return MACEEvaluator(
            model=model or "MACE-OMOL-0",
            device=args.device,
            default_dtype=args.mace_default_dtype,
            calc_kwargs=None,
        )

    if backend == "aimnet2":
        return AIMNet2Evaluator(
            model=model or "aimnet2",
            device=args.device,
            calc_kwargs=None,
        )

    raise QCShimError("Unsupported backend: {}".format(backend))


def _write_outputs(logfile, savfile, energy_ha, efield_mm_au, grad_qm_ha_bohr):
    logdir = os.path.dirname(os.path.abspath(logfile))
    if logdir and not os.path.isdir(logdir):
        os.makedirs(logdir)

    with open(logfile, "w") as handle:
        handle.write(" AMBER-MLIPS qchem shim\n")
        handle.write(" Charge-charge energy = {:17.10f}\n".format(0.0))
        handle.write(" SCF   energy in the final basis set = {: .16f}\n".format(float(energy_ha)))

    # AMBER qc reader expects efield.dat in current working directory.
    with open("efield.dat", "w") as handle:
        for row in np.asarray(efield_mm_au, dtype=np.float64).reshape(-1, 3):
            handle.write("{:22.16f}{:22.16f}{:22.16f}\n".format(float(row[0]), float(row[1]), float(row[2])))
        for row in np.asarray(grad_qm_ha_bohr, dtype=np.float64).reshape(-1, 3):
            handle.write("{:25.20f}{:25.20f}{:25.20f}\n".format(float(row[0]), float(row[1]), float(row[2])))

    # Keep a tiny save/checkpoint marker to mirror qchem signature.
    with open(savfile, "w") as handle:
        handle.write("# amber-mlips qchem shim checkpoint marker\n")


def _evaluate_via_server(socket_path, symbols, coords_q, charge, spinmult):
    """Evaluate using the persistent model server."""
    energy_ev, forces_ev_ang, _ = client_evaluate(
        socket_path=socket_path,
        symbols=symbols,
        coords_ang=coords_q,
        charge=int(charge),
        multiplicity=int(spinmult),
        need_forces=True,
        need_hessian=False,
        hessian_mode="analytical",
        hessian_step=1.0e-3,
    )
    return energy_ev, forces_ev_ang


def _evaluate_direct(kw_args, symbols, coords_q, charge, spinmult):
    """Evaluate by loading model in-process (fallback, slow)."""
    evaluator = _create_evaluator(kw_args)
    energy_ev, forces_ev_ang, _ = evaluator.evaluate(
        symbols=symbols,
        coords_ang=coords_q,
        charge=int(charge),
        multiplicity=int(spinmult),
        need_forces=True,
        need_hessian=False,
        hessian_mode="analytical",
        hessian_step=1.0e-3,
    )
    return energy_ev, forces_ev_ang


def run_qchem(inpfile, logfile, savfile, backend, ml_keywords, env_debug=False, server_socket=None):
    kw_args = _parse_keywords(backend=backend, ml_keywords=ml_keywords)
    debug = bool(env_debug) or bool(getattr(kw_args, "debug", False))

    (
        charge,
        spinmult,
        symbols,
        coords_q,
        coords_m,
        charges_m,
    ) = _parse_qchem_input(inpfile)

    nq = int(coords_q.shape[0])
    ncl = int(coords_m.shape[0])

    if server_socket and os.path.exists(server_socket):
        if debug:
            print("[amber-mlips-qchem] Using model server: {}".format(server_socket), file=sys.stderr, flush=True)
        energy_ev, forces_ev_ang = _evaluate_via_server(
            socket_path=server_socket,
            symbols=symbols, coords_q=coords_q,
            charge=charge, spinmult=spinmult,
        )
    else:
        if debug:
            print("[amber-mlips-qchem] No server, loading model directly (slow).", file=sys.stderr, flush=True)
        energy_ev, forces_ev_ang = _evaluate_direct(
            kw_args=kw_args,
            symbols=symbols, coords_q=coords_q,
            charge=charge, spinmult=spinmult,
        )

    forces_q = np.asarray(forces_ev_ang, dtype=np.float64).reshape(nq, 3)
    forces_m = np.zeros((ncl, 3), dtype=np.float64)

    if bool(kw_args.embedcharge) and ncl > 0:
        try:
            de_embed, df_full, _ = delta_embedcharge_minus_noembed(
                symbols=symbols,
                coords_q_ang=coords_q,
                mm_coords_ang=coords_m,
                mm_charges=charges_m,
                charge=int(charge),
                multiplicity=int(spinmult),
                need_forces=True,
                need_hessian=False,
                xtb_cmd=kw_args.xtb_cmd,
                xtb_acc=float(kw_args.xtb_acc),
                xtb_workdir=kw_args.xtb_workdir,
                xtb_keep_files=bool(kw_args.xtb_keep_files),
                ncores=int(kw_args.xtb_ncores),
                hessian_step=1.0e-3,
            )
        except XTBError as exc:
            raise QCShimError(
                "Embedcharge correction failed. "
                "Set --xtb-cmd correctly or disable --embedcharge. Details: {}".format(exc)
            ) from exc

        energy_ev += float(de_embed)
        df_full = np.asarray(df_full, dtype=np.float64).reshape(nq + ncl, 3)
        forces_q = forces_q + df_full[:nq, :]
        if ncl > 0:
            forces_m = forces_m + df_full[nq:, :]

    energy_ha = float(ev_to_ha(energy_ev))

    grad_q = np.asarray(
        forces_ev_ang_to_gradient_ha_bohr(forces_q),
        dtype=np.float64,
    ).reshape(nq, 3)

    grad_m = np.asarray(
        forces_ev_ang_to_gradient_ha_bohr(forces_m),
        dtype=np.float64,
    ).reshape(ncl, 3)

    efield_mm = np.zeros((ncl, 3), dtype=np.float64)
    for i in range(ncl):
        q = float(charges_m[i])
        if abs(q) > 1.0e-14:
            # AMBER later converts back to MM gradient using: grad = -E * q.
            efield_mm[i, :] = -grad_m[i, :] / q

    if debug:
        print(
            "[amber-mlips-qchem] backend={} nq={} ncl={} E={:.10f} Eh".format(
                backend,
                nq,
                ncl,
                energy_ha,
            ),
            file=sys.stderr,
            flush=True,
        )

    _write_outputs(
        logfile=logfile,
        savfile=savfile,
        energy_ha=energy_ha,
        efield_mm_au=efield_mm,
        grad_qm_ha_bohr=grad_q,
    )


def main(argv=None):
    parser = argparse.ArgumentParser(
        prog="qchem",
        description="amber-mlips non-MPI qchem shim",
    )
    parser.add_argument("inpfile")
    parser.add_argument("logfile")
    parser.add_argument("savfile")

    args = parser.parse_args(argv)

    backend = os.environ.get("AMBER_MLIPS_BACKEND", "").strip().lower()
    if not backend:
        print(
            "[amber-mlips-qchem] ERROR: AMBER_MLIPS_BACKEND is not set.",
            file=sys.stderr,
        )
        return 2

    ml_keywords = os.environ.get("AMBER_MLIPS_ML_KEYWORDS", "")
    server_socket = os.environ.get("AMBER_MLIPS_SERVER_SOCKET", "").strip() or None
    env_debug = str(os.environ.get("AMBER_MLIPS_DEBUG", "")).strip() in {
        "1",
        "true",
        "TRUE",
        "yes",
        "on",
    }

    try:
        run_qchem(
            inpfile=args.inpfile,
            logfile=args.logfile,
            savfile=args.savfile,
            backend=backend,
            ml_keywords=ml_keywords,
            env_debug=env_debug,
            server_socket=server_socket,
        )
        return 0
    except (QCShimError, BackendError, ServerError, OSError, ValueError) as exc:
        print("[amber-mlips-qchem] ERROR: {}".format(exc), file=sys.stderr)
        return 1


if __name__ == "__main__":
    raise SystemExit(main())
