From e3b0af94ed8aa1f474f38a6e5794d43bd9a3ac23 Mon Sep 17 00:00:00 2001
From: Ali Khan <alik@robarts.ca>
Date: Thu, 19 Dec 2024 22:48:07 -0500
Subject: [PATCH] removes --keep-work, adds --workdir

- WIP
- sets workdir to a temp directory by default
- the files won't be deleted by the workflow, but we could make use of
temp() to do so..
---
 hippunfold/config/snakebids.yml      |  9 +++--
 hippunfold/workflow/Snakefile        | 10 +++--
 hippunfold/workflow/rules/common.smk | 57 +---------------------------
 3 files changed, 13 insertions(+), 63 deletions(-)

diff --git a/hippunfold/config/snakebids.yml b/hippunfold/config/snakebids.yml
index 850a729d..c6f38d55 100644
--- a/hippunfold/config/snakebids.yml
+++ b/hippunfold/config/snakebids.yml
@@ -264,10 +264,10 @@ parse_args:
     nargs: '+'
     help: 'Run hipp (CA + subiculum) alone or include dentate (default: %(default)s)'
 
-  --keep_work:
-    help: 'Keep work folder intact instead of archiving it for each subject (default: %(default)s)'
-    default: False
-    action: 'store_true'
+  --workdir:
+    help: 'Folder for storing working files, if not specified, will make a temporary folder with tempfile.mkdtemp() (default: %(default)s)'
+    default: None
+    type: Path
 
   --force_nnunet_model:
     help: 'Force nnunet model to use (expert option). (default: %(default)s)'
@@ -693,6 +693,7 @@ t1_reg_template: False
 generate_myelin_map: False
 no_unfolded_reg: False
 root: results
+workdir: null
 use_template_seg: False
 template_seg_smoothing_factor: 2
 native_surf: False
diff --git a/hippunfold/workflow/Snakefile b/hippunfold/workflow/Snakefile
index c9952da9..9071bda0 100644
--- a/hippunfold/workflow/Snakefile
+++ b/hippunfold/workflow/Snakefile
@@ -2,7 +2,7 @@
 
 import snakebids
 from snakebids import bids, set_bids_spec
-
+from tempfile import mkdtemp
 
 configfile: "config/snakebids.yml"
 
@@ -94,9 +94,13 @@ wildcard_constraints:
     template="[a-zA-Z0-9]+",
 
 
-root = os.path.normpath(os.path.join(config["root"], "hippunfold"))
-work = os.path.normpath(os.path.join(config["root"], "work"))
+#root = os.path.normpath(os.path.join(config["root"], "hippunfold"))
+#work = os.path.normpath(os.path.join(config["root"], "work"))
 
+root = os.path.expandvars(config["root"])
+if config["workdir"] == None:
+    config["workdir"] = mkdtemp()
+work = os.path.expandvars(config["workdir"])
 
 include: "rules/common.smk"
 include: "rules/download.smk"
diff --git a/hippunfold/workflow/rules/common.smk b/hippunfold/workflow/rules/common.smk
index 91fd1c23..ae9a2aa4 100644
--- a/hippunfold/workflow/rules/common.smk
+++ b/hippunfold/workflow/rules/common.smk
@@ -332,7 +332,7 @@ def get_final_qc():
     return qc
 
 
-def get_final_subj_output():
+def get_final_output():
     subj_output = []
     subj_output.extend(get_final_spec())
     subj_output.extend(get_final_subfields())
@@ -343,33 +343,6 @@ def get_final_subj_output():
     return subj_output
 
 
-def get_final_output():
-    if config["keep_work"]:
-        subj_output = get_final_subj_output()
-    else:
-        subj_output = get_final_work_tar()
-
-    final_output = []
-
-    modality_suffix = get_modality_suffix(config["modality"])
-    modality_key = config["modality"]
-
-    # use a zip list for subject/session:
-    zip_list = inputs[modality_key].zip_lists
-    if "session" in zip_list:
-        zip_list = snakebids.filter_list(
-            zip_list,
-            {"session": inputs[config["modality"]].zip_lists["session"]},
-        )
-    final_output.extend(
-        expand(
-            expand(subj_output, zip, allow_missing=True, **zip_list),
-            modality_suffix=modality_suffix,
-        )
-    )
-    return final_output
-
-
 if "corobl" in ref_spaces:
 
     rule copy_coords_to_results:
@@ -403,18 +376,6 @@ if "corobl" in ref_spaces:
             "cp {input} {output}"
 
 
-def get_final_work_tar():
-    bids = bids_factory(specs.v0_0_0(subject_dir=False, session_dir=False))
-    return bids(root=work, suffix="work.tar.gz", **inputs.subj_wildcards)
-
-
-def get_work_dir(wildcards):
-    folder_with_file = inputs[config["modality"]].expand(
-        bids(root=work, **inputs.subj_wildcards), **wildcards
-    )
-    folder_without_file = os.path.dirname(folder_with_file[0])
-    return folder_without_file
-
 
 def get_download_dir():
     if "HIPPUNFOLD_CACHE_DIR" in os.environ.keys():
@@ -426,19 +387,3 @@ def get_download_dir():
     return download_dir
 
 
-rule archive_work_after_final:
-    input:
-        get_final_subj_output(),
-    params:
-        work_dir=get_work_dir,
-    output:
-        get_final_work_tar(),
-    group:
-        "subj"
-    shell:
-        #exit code 0 or 1 is acceptable (2 is fatal)
-        "tar -czf {output} {params.work_dir}; "
-        "if [ $? -le 1 ]; then "
-        "  rm -rf {params.work_dir}; "
-        "else exit 1; "
-        "fi"