From 7dde9b12448cd5314546a37f9dec996b9b44ba23 Mon Sep 17 00:00:00 2001
From: Calvin <calvinpieters@gmail.com>
Date: Fri, 27 Dec 2024 22:37:09 +0200
Subject: [PATCH] Auto stash before cherry pick of "Update file cleanup logic
 to exclude additional output files and set working directory"

---
 arc/job/adapters/ts/heuristics.py | 129 +++++++++++++++++++++++++++++-
 1 file changed, 126 insertions(+), 3 deletions(-)

diff --git a/arc/job/adapters/ts/heuristics.py b/arc/job/adapters/ts/heuristics.py
index 635f7874c7..5599d17ed2 100644
--- a/arc/job/adapters/ts/heuristics.py
+++ b/arc/job/adapters/ts/heuristics.py
@@ -1154,9 +1154,10 @@ def crest_ts_conformer_search(xyz_guess: dict, a_atom: int, h_atom: int, b_atom:
     # Prepare the command
     commands = [
         f'{CREST_PATH}',
-        f'{path}/coords.ref',
-        f'--cinp {path}/constraints.inp',
-        '--noreftopo'
+        'coords.ref',
+        '--cinp constraints.inp',
+        '--noreftopo',
+        f' -T {SERVERS["local"].get("cpus", 8)}',
     ]
     command = ' '.join(commands)
     command = f"{CREST_ENV_PATH} && {command}" if CREST_ENV_PATH else command
@@ -1177,6 +1178,128 @@ def crest_ts_conformer_search(xyz_guess: dict, a_atom: int, h_atom: int, b_atom:
         print(f"Standard Error: {stderr.decode()}")
         return None
 
+    if SERVERS.get('local') is not None:
+        if SERVERS['local']['cluster_soft'].lower() in ['condor', 'htcondor']:
+            # Write job submission scripts
+
+            sub_job = submit_scripts['local']['crest']
+            format_params = {
+                "name": f"crest_{xyz_crest_int}",
+                "cpus": SERVERS['local'].get('cpus', 8),
+                "memory": SERVERS['local'].get('memory', 32.0) * 1024,
+            }
+            sub_job = sub_job.format(**format_params)
+
+            with open(os.path.join(path, settings['submit_filenames']['HTCondor']), 'w') as f:
+                f.write(sub_job)
+
+            # Write the crest job
+            crest_job = submit_scripts['local']['crest_job']
+            format_params = {
+                "path": path,
+                "activation_line": activation_line,
+                "commands": command,
+            }
+            crest_job = crest_job.format(**format_params)
+
+            with open(os.path.join(path, 'job.sh'), 'w') as f:
+                f.write(crest_job)
+        elif SERVERS['local']['cluster_soft'].lower() == 'pbs':
+            # Write job submission scripts
+            sub_job = submit_scripts['local']['crest']
+            format_params = {
+                "queue": "alon_q",
+                "name": f"crest_{xyz_crest_int}",
+                "cpus": SERVERS['local'].get('cpus', 8),
+                "memory": SERVERS['local'].get('memory', 32.0)  if SERVERS['local'].get('memory', 32.0) < 60.0 else 40.0,
+                "activation_line": activation_line,
+                "commands": command,
+            }
+            sub_job = sub_job.format(**format_params)
+
+            with open(os.path.join(path, settings['submit_filenames']['PBS']), 'w') as f:
+                f.write(sub_job)
+
+    return path
+
+def submit_crest_jobs(crest_paths: List[str]) -> None:
+    """
+    Submit CREST jobs to the server.
+
+    Args:
+        crest_paths (List[str]): List of paths to the CREST directories.
+
+    Returns:
+        dict: A dictionary containing job IDs as keys and their statuses as values.
+    """
+    crest_jobs = {}
+    for crest_path in crest_paths:
+        job_status, job_id = submit_job(path=crest_path)
+        crest_jobs[job_id] = {"path": crest_path, "status": job_status}
+    return crest_jobs
+
+
+def monitor_crest_jobs(crest_jobs: dict, check_interval: int = 300) -> None:
+    """
+    Monitor CREST jobs until they are complete.
+
+    Args:
+        crest_jobs (dict): Dictionary containing job information (job ID, path, and status).
+        check_interval (int): Time interval (in seconds) to wait between status checks.
+    """
+    while True:
+        all_done = True
+        for job_id, job_info in crest_jobs.items():
+            if job_info["status"] not in ["done", "failed"]:
+                try:
+                    job_info["status"] = check_job_status(job_id)  # Update job status
+                except Exception as e:
+                    logger.error(f"Error checking job status for job {job_id}: {e}")
+                    job_info["status"] = "failed"
+                if job_info["status"] not in ["done", "failed"]:
+                    all_done = False
+        if all_done:
+            break
+        time.sleep(min(check_interval, 100))
+
+def process_completed_jobs(crest_jobs: dict, xyz_guesses: list) -> None:
+    """
+    Process the completed CREST jobs and update XYZ guesses.
+
+    Args:
+        crest_jobs (dict): Dictionary containing job information.
+        xyz_guesses (list): List to store the resulting XYZ guesses.
+    """
+    for job_id, job_info in crest_jobs.items():
+        crest_path = job_info["path"]
+        if job_info["status"] == "done":
+            crest_best_path = os.path.join(crest_path, "crest_best.xyz")
+            if os.path.exists(crest_best_path):
+                with open(crest_best_path, "r") as f:
+                    content = f.read()
+                xyz_guess = str_to_xyz(content)
+                xyz_guesses.append(xyz_guess)
+            else:
+                logger.error(f"crest_best.xyz not found in {crest_path}")
+        elif job_info["status"] == "failed":
+            logger.error(f"CREST job failed for {crest_path}")
+
+    return xyz_guesses
+
+def extract_digits(s: str) -> int:
+    """
+    Extract the first integer from a string
+
+    Args:
+        s (str): The string to extract the integer from
+
+    Returns:
+        int: The first integer in the string
+
+    """
+    return int(re.sub(r"[^\d]", "", s))
+
+
 def convert_xyz_to_df(xyz: dict) -> pd.DataFrame:
     """
     Convert a dictionary of xyz coords to a pandas DataFrame with bond distances