From 38dd0a07d333b7930d0563af46737fe5173fedd4 Mon Sep 17 00:00:00 2001
From: Tim Dudgeon <tdudgeon@informaticsmatters.com>
Date: Tue, 17 Dec 2024 12:20:28 +0000
Subject: [PATCH] guess correct working dir

---
 src/xchemalign/aligner.py  | 21 +++++++++-------
 src/xchemalign/collator.py | 20 ++++------------
 src/xchemalign/utils.py    | 49 ++++++++++++++++++++++++++++++++++----
 3 files changed, 62 insertions(+), 28 deletions(-)

diff --git a/src/xchemalign/aligner.py b/src/xchemalign/aligner.py
index 85a7932..aaf608d 100644
--- a/src/xchemalign/aligner.py
+++ b/src/xchemalign/aligner.py
@@ -219,23 +219,26 @@ def __init__(self, dir, logger=None):
 
     def _find_version_dir(self, dir):
         if dir:
-            self.working_dir = Path(dir)
+            wd0 = Path(dir)
         else:
-            self.working_dir = Path.cwd()
+            wd0 = Path.cwd()
 
-        if not self.working_dir.is_dir():
-            self._log_error("Working dir {} does not exist".format(self.working_dir))
+        wd1 = utils._verify_working_dir(wd0)
+        if not wd1:
+            self._log_error("Working dir " + wd0 + " is not valid")
             exit(1)
 
+        self.working_dir = wd1
         current_dir = self.working_dir / 'upload-current'
-        if not current_dir.is_symlink():
-            self._log_error(
-                "Working dir {} does not seem well formed - missing 'upload_current' symlink".format(self.working_dir)
-            )
+
+        # check that we at least have an upload_1 dir
+        if not (current_dir / 'upload_1').is_dir():
+            self._log_error("Working dir " + dir + " does not contain and upload_? dirs")
             exit(1)
 
+        # now find the latest upload_? dir
         i = 0
-        while True:
+        while i < 100:
             i += 1
             version_dir = current_dir / ('upload_' + str(i))
             if version_dir.is_dir():
diff --git a/src/xchemalign/collator.py b/src/xchemalign/collator.py
index 5c70ca3..4279fa5 100644
--- a/src/xchemalign/collator.py
+++ b/src/xchemalign/collator.py
@@ -1378,20 +1378,6 @@ def get_dataset_event_maps(
         return closest_event_maps
 
 
-def _check_working_dir(working_dir):
-    print("Using", working_dir, "as working dir")
-
-    if not working_dir.is_dir():
-        print("Working dir {} does not exist".format(working_dir))
-        exit(1)
-
-    current_dir = working_dir / 'upload-current'
-    if current_dir.is_symlink():
-        return 0
-    else:
-        return 1
-
-
 def main():
     parser = argparse.ArgumentParser(description="collator")
 
@@ -1408,7 +1394,11 @@ def main():
     else:
         working_dir = Path.cwd()
 
-    if _check_working_dir(working_dir):
+    wd = utils._verify_working_dir(working_dir)
+
+    if wd:
+        working_dir = wd
+    else:
         print("Working dir does not seem to have been initialised - missing 'upload_current' symlink")
         inp = input("Do you want the working dir to be initialised? (Y/N)")
         if inp == "Y" or inp == "y":
diff --git a/src/xchemalign/utils.py b/src/xchemalign/utils.py
index 73cf92e..427e6f7 100644
--- a/src/xchemalign/utils.py
+++ b/src/xchemalign/utils.py
@@ -544,6 +544,45 @@ def check_data_format_version(ver_to_check):
         return 1
 
 
+def _verify_working_dir(working_dir):
+    """
+    Verifies that we can work out what the real working dir is. Preferably the user has specified it correctly,
+    but we allow then to specify one or 2 levels down from this.
+
+    :param working_dir: The working dir specified by the user
+    :return: The actual working dir, or None if we can't work it out
+    """
+    print("Trying", working_dir, "as working dir")
+
+    if not working_dir.is_dir():
+        print("Working dir {} does not exist".format(working_dir))
+        exit(1)
+
+    current_dir = working_dir / 'upload-current'
+    if current_dir.is_symlink():
+        return working_dir
+    else:
+        # check if working dir is already one or 2 levels under where it really needs to be
+        working_dir = working_dir.parent
+        current_dir = working_dir / 'upload-current'
+        if current_dir.is_symlink():
+            print(
+                "WARNING: you seem to have specified a directory one level under the true working dir. Using",
+                working_dir,
+            )
+            return working_dir
+        working_dir = working_dir.parent
+        current_dir = working_dir / 'upload-current'
+        if current_dir.is_symlink():
+            print(
+                "WARNING: you seem to have specified a directory two levels under the true working dir. Using",
+                working_dir,
+            )
+            return working_dir
+        else:
+            return None
+
+
 def main():
     # log = Logger(logfile="logfile.log", level=1)
     #
@@ -551,10 +590,12 @@ def main():
     # log.log("foo", "bar", "baz")
     # log.log("foo", 99, "apples", level=2)
 
-    mols = gen_mols_from_cif('data/Zx1674a.cif')
-    for mol in mols:
-        molfile = Chem.MolToMolBlock(mol)
-        print(molfile)
+    # mols = gen_mols_from_cif('data/Zx1674a.cif')
+    # for mol in mols:
+    #     molfile = Chem.MolToMolBlock(mol)
+    #     print(molfile)
+
+    print(_verify_working_dir(Path('data/std_test/lb32633-6_2024-11-22/upload-current/upload_1')))
 
 
 if __name__ == "__main__":