#!/usr/bin/python3

"""Lightweight GRRM-to-server bridge.

The GRRM-facing file names and the two-byte server protocol are intentionally
kept identical to the original implementation.
"""

import os
import socket
import struct
import sys
import time

SOCKET_POINTER_FILE = ".socketfile"
RESPONSE_SIZE = 2
CONNECT_TIMEOUT_SECONDS = 30.0
CONNECT_RETRY_INTERVAL_SECONDS = 0.1
FIN_TIMEOUT_SECONDS = 1000.0
POLL_INTERVAL_SECONDS = 0.1


def get_socket_file():
    """Read the Unix-domain socket path used by the persistent server."""
    if not os.path.exists(SOCKET_POINTER_FILE):
        print("ERROR: .socketfile not found", file=sys.stderr)
        sys.exit(8)

    with open(SOCKET_POINTER_FILE, encoding="utf-8") as f:
        return f.read().strip()


def recv_exact(sock, size):
    """Receive exactly *size* bytes or raise if the connection closes early."""
    chunks = []
    remaining = size
    while remaining:
        chunk = sock.recv(remaining)
        if not chunk:
            raise ConnectionError(
                f"Connection closed while receiving {size} bytes"
            )
        chunks.append(chunk)
        remaining -= len(chunk)
    return b"".join(chunks)


def connect_with_retry(socket_file):
    """Connect to the Unix-domain server, retrying transient startup errors."""
    deadline = time.monotonic() + CONNECT_TIMEOUT_SECONDS
    last_error = None

    while True:
        sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
        try:
            sock.connect(socket_file)
            return sock
        except (FileNotFoundError, ConnectionRefusedError) as exc:
            sock.close()
            last_error = exc
            if time.monotonic() >= deadline:
                raise TimeoutError(
                    f"Server did not become ready within "
                    f"{CONNECT_TIMEOUT_SECONDS:.1f} s: {last_error}"
                ) from exc
            time.sleep(CONNECT_RETRY_INTERVAL_SECONDS)
        except Exception:
            sock.close()
            raise


def wait_for_and_remove_fin(finfile):
    """Wait for GRRM's FIN4GEN flag file and remove it once observed."""
    deadline = time.monotonic() + FIN_TIMEOUT_SECONDS
    while time.monotonic() < deadline:
        if os.path.exists(finfile):
            os.remove(finfile)
            return True
        time.sleep(POLL_INTERVAL_SECONDS)
    return False


def main():
    if len(sys.argv) != 2:
        print(f"Usage: {sys.argv[0]} ARGNAME", file=sys.stderr)
        sys.exit(1)

    argname = sys.argv[1]
    socket_file = get_socket_file()

    inpfile = f"{argname}_INP4GEN.rrm"
    if not os.path.exists(inpfile):
        print(f"Input file {inpfile} not found.", file=sys.stderr)
        sys.exit(2)

    try:
        with connect_with_retry(socket_file) as sock:
            msg = argname.encode()
            sock.sendall(struct.pack(">H", len(msg)) + msg)
            resp = recv_exact(sock, RESPONSE_SIZE)
    except Exception as exc:
        print(f"ERROR: Could not connect to server: {exc}", file=sys.stderr)
        sys.exit(3)

    if resp == b"OK":
        finfile = f"{argname}_FIN4GEN.rrm"
        if wait_for_and_remove_fin(finfile):
            sys.exit(0)
        print(
            "ERROR: Server responded OK but output file not found after waiting.",
            file=sys.stderr,
        )
        sys.exit(4)

    if resp == b"AL":
        print("ERROR: Job already in progress.", file=sys.stderr)
        sys.exit(5)
    if resp == b"ER":
        print("ERROR: Server error.", file=sys.stderr)
        sys.exit(6)

    print(f"ERROR: Unexpected server response: {resp}", file=sys.stderr)
    sys.exit(7)


if __name__ == "__main__":
    main()

#!/usr/bin/python3

"""Persistent Lennard-Jones backend for GRRM external calculations.

GRRM-facing input/output file names and output formatting are intentionally
kept identical to the original implementation.  The server keeps Python and
NumPy resident and receives lightweight requests over a Unix-domain socket.
"""

import os
import socket
import struct
import sys
import threading
from concurrent.futures import ThreadPoolExecutor

import numpy as np

# Lennard-Jones parameters
EPSILON = 0.001
SIGMA = 1.6
BOHR_PER_ANGSTROM = 1.8897259886
REQUEST_LENGTH_BYTES = 2
DEFAULT_MAX_WORKERS = 40


# --- Lennard-Jones calculation functions ---
def lj_potential_energy(r):
    with np.errstate(divide="ignore", invalid="ignore", over="ignore"):
        value = 4.0 * EPSILON * ((SIGMA / r) ** 12 - (SIGMA / r) ** 6)
    return np.nan_to_num(value, nan=0.0)


def lj_potential_gradient(r):
    with np.errstate(divide="ignore", invalid="ignore", over="ignore"):
        value = -4.0 * EPSILON * (
            12 * (SIGMA**12 / r**13) - 6 * (SIGMA**6 / r**7)
        )
    return np.nan_to_num(value, nan=0.0)


def distance_matrix_np(coords):
    coords = np.asarray(coords, dtype=float)
    diff = coords[:, None, :] - coords[None, :, :]
    return np.sqrt(np.sum(diff**2, axis=-1))


def calc_gradient_vectors(distance_list, coord_list, grad_matrix):
    coords = np.asarray(coord_list, dtype=float)
    vector_list = coords[:, None, :] - coords[None, :, :]
    distance_array = np.asarray(distance_list)
    mask = distance_array > 1e-4
    grad = np.asarray(grad_matrix)

    with np.errstate(divide="ignore", invalid="ignore"):
        gradient_contrib = np.where(
            mask[:, :, None],
            vector_list * grad[:, :, None] / distance_array[:, :, None],
            0.0,
        )

    return np.sum(gradient_contrib, axis=1)


def read_grrm_input(inp_filename):
    """Parse only the fields required by the LJ test backend."""
    ene_mode = 0
    grad_mode = 0
    hess_mode = 0
    atomnum = 0
    coord_list = []
    atomname_list = []

    with open(inp_filename, encoding="utf-8") as f:
        lines = f.readlines()

    for idx, line in enumerate(lines):
        if "TASK:" in line:
            if "ENERGY" in line:
                ene_mode = 1
            if "GRADIENT" in line:
                grad_mode = 1
            if "HESSIAN" in line:
                hess_mode = 1

        if "NACTIVEATOM / NATOM:" in line:
            elem = line.split()
            atomnum = int(elem[3])
            for j in range(atomnum):
                atomline = lines[idx + 1 + j]
                elems = atomline.split()
                atomname_list.append(elems[0])
                coord_list.append(
                    [float(elems[1]), float(elems[2]), float(elems[3])]
                )
            break

    return ene_mode, grad_mode, hess_mode, atomnum, coord_list, atomname_list


def write_grrm_output(
    out_filename,
    atomnum,
    coord_list,
    atomname_list,
    optene,
    gradient_vector_list,
):
    """Write the GRRM output in exactly the same format as the original."""
    with open(out_filename, mode="w", encoding="utf-8") as fout:
        fout.write("RESULTS\n")
        fout.write("CURRENT COORDINATE\n")
        for i in range(atomnum):
            fout.write(
                f"{atomname_list[i]:2s} {coord_list[i][0]:22.12f} "
                f"{coord_list[i][1]:22.12f} {coord_list[i][2]:22.12f}\n"
            )
        fout.write(
            f"ENERGY = {optene:17.12f}    0.000000000000    0.000000000000\n"
        )
        fout.write(
            "       =    0.000000000000    0.000000000000    0.000000000000\n"
            "S**2   =    0.000000000000\n"
        )
        fout.write("GRADIENT\n")
        if gradient_vector_list is not None:
            for vec in gradient_vector_list:
                for val in vec:
                    fout.write(f"  {val / BOHR_PER_ANGSTROM:17.12f}\n")
        else:
            for _ in range(atomnum * 3):
                fout.write("  0.000000000000\n")
        fout.write(
            "DIPOLE =    0.000000000000         0.000000000000          "
            "0.000000000000\n"
        )
        fout.write("HESSIAN\n")
        for i in range((3 * atomnum) // 5 + 1):
            tmp_linenum = (3 * atomnum) - i * 5
            if tmp_linenum <= 0:
                tmp_linenum = (3 * atomnum) % 5
            for j in range(tmp_linenum):
                if j == 0:
                    fout.write("  0.000000000\n")
                if j == 1:
                    fout.write("  0.000000000   0.000000000\n")
                if j == 2:
                    fout.write(
                        "  0.000000000   0.000000000   0.000000000\n"
                    )
                if j == 3:
                    fout.write(
                        "  0.000000000   0.000000000   0.000000000   "
                        "0.000000000\n"
                    )
                if j >= 4:
                    fout.write(
                        "  0.000000000   0.000000000   0.000000000   "
                        "0.000000000   0.000000000\n"
                    )
        fout.write("DIPOLE DERIVATIVES\n")
        for _ in range(3 * atomnum):
            fout.write(
                "   0.000000000000          0.000000000000          "
                "0.000000000000\n"
            )
        fout.write("POLARIZABILITY\n")
        fout.write("   0.000000000000\n")
        fout.write("   0.000000000000          0.000000000000\n")
        fout.write(
            "   0.000000000000          0.000000000000          "
            "0.000000000000\n"
        )


def run_lj_model(inp_filename, out_filename, fin_filename):
    (
        ene_mode,
        grad_mode,
        _hess_mode,
        atomnum,
        coord_list,
        atomname_list,
    ) = read_grrm_input(inp_filename)

    coord_np = np.asarray(coord_list, dtype=float)
    distance_list = distance_matrix_np(coord_np)

    optene = 0.0
    if ene_mode:
        ene_matrix = lj_potential_energy(distance_list)
        optene = np.triu(ene_matrix, k=1).sum()

    gradient_vector_list = None
    if grad_mode:
        grad_matrix = lj_potential_gradient(distance_list)
        gradient_vector_list = calc_gradient_vectors(
            distance_list, coord_np, grad_matrix
        )

    write_grrm_output(
        out_filename,
        atomnum,
        coord_list,
        atomname_list,
        optene,
        gradient_vector_list,
    )

    with open(fin_filename, mode="w", encoding="utf-8") as ffin:
        ffin.write("")


def recv_exact(sock, size):
    """Receive exactly *size* bytes or raise if the peer closes early."""
    chunks = []
    remaining = size
    while remaining:
        chunk = sock.recv(remaining)
        if not chunk:
            raise ConnectionError(
                f"Connection closed while receiving {size} bytes"
            )
        chunks.append(chunk)
        remaining -= len(chunk)
    return b"".join(chunks)


# --- Server class ---
class RRMServer:
    def __init__(self, workdir, socket_file, max_workers=DEFAULT_MAX_WORKERS):
        self.workdir = os.path.abspath(workdir)
        self.socket_file = os.path.abspath(socket_file)
        self.executor = ThreadPoolExecutor(max_workers=max_workers)
        self.jobs_in_progress = set()
        self.lock = threading.Lock()

    def serve(self):
        if os.path.exists(self.socket_file):
            os.unlink(self.socket_file)

        with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as serversock:
            serversock.bind(self.socket_file)
            serversock.listen(8)
            print(f"Listening on {self.socket_file}", flush=True)

            try:
                while True:
                    conn, _ = serversock.accept()
                    threading.Thread(
                        target=self.handle_client,
                        args=(conn,),
                        daemon=True,
                    ).start()
            finally:
                self.executor.shutdown(wait=True)
                if os.path.exists(self.socket_file):
                    os.unlink(self.socket_file)

    def handle_client(self, conn):
        try:
            length = struct.unpack(">H", recv_exact(conn, REQUEST_LENGTH_BYTES))[0]
            argname = recv_exact(conn, length).decode()

            inp = os.path.join(self.workdir, f"{argname}_INP4GEN.rrm")
            out = os.path.join(self.workdir, f"{argname}_OUT4GEN.rrm")
            fin = os.path.join(self.workdir, f"{argname}_FIN4GEN.rrm")

            print("filepath:", inp, out, flush=True)

            with self.lock:
                if inp in self.jobs_in_progress:
                    conn.sendall(b"AL")
                    return
                self.jobs_in_progress.add(inp)

            # Wait on the Future rather than polling FIN.  This preserves the
            # external protocol while correctly propagating backend failures.
            future = self.executor.submit(self.do_job, inp, out, fin)
            future.result()
            conn.sendall(b"OK")

        except Exception as exc:
            print("ERROR:", exc, flush=True)
            try:
                conn.sendall(b"ER")
            except Exception:
                pass
        finally:
            conn.close()

    def do_job(self, inp, out, fin):
        try:
            run_lj_model(inp, out, fin)
        finally:
            with self.lock:
                self.jobs_in_progress.discard(inp)


if __name__ == "__main__":
    if len(sys.argv) < 3:
        print(f"Usage: {sys.argv[0]} WORKDIR SOCKET_FILE", file=sys.stderr)
        sys.exit(1)

    workdir = sys.argv[1]
    socket_file = sys.argv[2]
    RRMServer(workdir, socket_file, max_workers=DEFAULT_MAX_WORKERS).serve()