From 224383d11aad99008048e74720065ac5d7ac4ff9 Mon Sep 17 00:00:00 2001
From: alessandrofelder <alessandrofelder@users.noreply.github.com>
Date: Mon, 18 Sep 2023 16:23:20 +0100
Subject: [PATCH] works locally

---
 brainglobe_workflows/niftyreg_defaults.json |  12 ++
 brainglobe_workflows/registration_script.py | 128 ++++++++++++++------
 2 files changed, 106 insertions(+), 34 deletions(-)
 create mode 100644 brainglobe_workflows/niftyreg_defaults.json

diff --git a/brainglobe_workflows/niftyreg_defaults.json b/brainglobe_workflows/niftyreg_defaults.json
new file mode 100644
index 00000000..207a49f1
--- /dev/null
+++ b/brainglobe_workflows/niftyreg_defaults.json
@@ -0,0 +1,12 @@
+{
+    "affine_n_steps": 6,
+    "affine_use_n_steps": 5,
+    "freeform_n_steps": 6,
+    "freeform_use_n_steps": 4,
+    "bending_energy_weight": 0.95,
+    "grid_spacing": -10,
+    "smoothing_sigma_reference": -1.0,
+    "smoothing_sigma_floating": -1.0,
+    "histogram_n_bins_floating": 128,
+    "histogram_n_bins_reference": 128
+}
\ No newline at end of file
diff --git a/brainglobe_workflows/registration_script.py b/brainglobe_workflows/registration_script.py
index a60b10cb..46375425 100644
--- a/brainglobe_workflows/registration_script.py
+++ b/brainglobe_workflows/registration_script.py
@@ -1,44 +1,104 @@
+import logging
 import os
 from pathlib import Path
 import subprocess
 import json
-from typing import Dict, Union
+from typing import Dict, Union, Tuple
+
+from pathlib import Path
+
+from brainglobe_utils.general.numerical import check_positive_int
+from brainglobe_utils.general.system import ensure_directory_exists
+from fancylog import fancylog
+
+import brainreg as program_for_log
+from brainreg import __version__
+from brainreg.backend.niftyreg.parser import niftyreg_parse
+from brainreg.main import main as register
+from brainreg.paths import Paths
+from brainreg.utils.misc import get_arg_groups, log_metadata
+from brainreg.cli import prep_registration
+
+from dataclasses import dataclass, make_dataclass
 
 Pathlike = Union[str, bytes, os.PathLike]
-# TODO: use pydantic or attrs, and/or a dataclass, for this??
-def _default_config():
-	defaults = {
-        "sample_data_path"
-    }
-
-def _validate_config(config_dict: Dict[str]):
-	expected_keys = [
-		"sample_data_path",
-		"output_path",
-		"voxel size",
-		"orientation"
-		"atlas"
-		]
-	expected_types = [
-		Pathlike,
-    ]
-	for key, type in zip(expected_keys, expected_types):
-        assert key in config_dict.keys()
-
-
-def example_workflow():
-	input_config_path = Path(os.environ("BRAINGLOBE_REGISTRATION_CONFIG_PATH"))
-	if input_config_path.exists():
-		config = json.loads(input_config_path)
-	else:
-		config = _default_config()
-	input_data = read_from_config(config)
-	preprocessed_data = process(input_data) # if required
-	results = run_main_tool(prepocessed_data)
-	save(results)
+# TODO: use pydanti`c or attrs, and/or a dataclass, for this?? And grab default test data off pooch!
+
+@dataclass
+class BrainregConfig():
+    image_paths: Pathlike
+    brainreg_directory: Pathlike
+    voxel_sizes: Tuple[float]
+    orientation: str
+    atlas: str
+    debug: bool = False
+    image_paths : Pathlike
+    backend: str = "niftyreg"
+    sort_input_file: bool = False
+    save_original_orientation: bool = False
+    brain_geometry: str = "full"
+    n_free_cpus: int = 1
+    additional_images = False
+    preprocessing: Dict[str, str] = None,
+    niftyreg_backend_options = None
+
+def example_registration_script():
+    input_config_path = Path(os.environ["BRAINGLOBE_REGISTRATION_CONFIG_PATH"])
+    if input_config_path.exists():
+        print(f"Read config from: {input_config_path}")
+        with open(input_config_path) as config_file:
+            config_dict = json.load(config_file)
+        config = BrainregConfig(**config_dict)
+    else:
+        raise NotImplementedError
+        # TODO: define some defaults here? load from local json file or pooch
+        # config = BrainregConfig(**default_dict)
+    
+    with open(Path(__file__).parent / "niftyreg_defaults.json") as niftyreg_defaults_file:
+        niftyreg_defaults_dict = json.load(niftyreg_defaults_file)
+        NiftyregOptions = make_dataclass("NiftyregOptions", niftyreg_defaults_dict.keys())
+        config.niftyreg_backend_options = NiftyregOptions(**niftyreg_defaults_dict)
+
+    config.preprocessing = {"preprocessing": "default"}
+    
+    config, additional_images_downsample = prep_registration(config)
+
+    paths = Paths(config.brainreg_directory)
+
+    log_metadata(paths.metadata_path, config)
+
+    fancylog.start_logging(
+        paths.registration_output_folder,
+        program_for_log,
+        variables=[config],
+        verbose=config.debug,
+        log_header="BRAINREG LOG",
+        multiprocessing_aware=False,
+    )
+
+    logging.info("Starting registration")
+
+    register(
+        config.atlas,
+        config.orientation,
+        config.image_paths,
+        paths,
+        config.voxel_sizes,
+        config.niftyreg_backend_options,
+        None,
+        sort_input_file=config.sort_input_file,
+        n_free_cpus=config.n_free_cpus,
+        additional_images_downsample=additional_images_downsample,
+        backend=config.backend,
+        debug=config.debug,
+        save_original_orientation=config.save_original_orientation,
+        brain_geometry=config.brain_geometry,
+    )
+
+    # TODO save(results) # or maybe assert stuff around results???
 
 if __name__ == "__main__":
-	example_workflow()
+	example_registration_script()
 
 """
 To run brainreg, you need to pass:
@@ -51,4 +111,4 @@ def example_workflow():
 We put this all together in a single command:
 
 brainreg test_brain brainreg_output -v 50 40 40  --orientation psl --atlas allen_mouse_50um
-"""
+"""
\ No newline at end of file