From 8b42ef36b3603b6f90d45d77352a0171b0f09639 Mon Sep 17 00:00:00 2001
From: Adam Tyson <code@adamltyson.com>
Date: Fri, 2 Aug 2024 14:48:24 +0100
Subject: [PATCH] Export cells to brainrender from brainmapper widget

---
 .../brainmapper/transform_widget.py           | 44 ++++++++++++++++++-
 .../test_brainmapper/test_transform_widget.py | 25 +++++++++++
 2 files changed, 68 insertions(+), 1 deletion(-)

diff --git a/brainglobe_utils/brainmapper/transform_widget.py b/brainglobe_utils/brainmapper/transform_widget.py
index 2fcceb5..24c4d24 100644
--- a/brainglobe_utils/brainmapper/transform_widget.py
+++ b/brainglobe_utils/brainmapper/transform_widget.py
@@ -21,6 +21,7 @@
 from brainglobe_utils.brainmapper.analysis import (
     summarise_points_by_atlas_region,
 )
+from brainglobe_utils.brainmapper.export import export_points_to_brainrender
 from brainglobe_utils.brainreg.transform import (
     transform_points_from_downsampled_to_atlas_space,
 )
@@ -146,7 +147,7 @@ def setup_main_layout(self) -> None:
         self.add_points_combobox(row=1, column=0)
         self.add_raw_data_combobox(row=2, column=0)
         self.add_transform_button(row=3, column=0)
-
+        self.add_brainrender_export_button(row=3, column=1)
         self.add_points_summary_table(row=4, column=0)
         self.add_save_all_points_button(row=6, column=0)
         self.add_save_points_summary_button(row=6, column=1)
@@ -229,6 +230,28 @@ def add_transform_button(self, row: int, column: int) -> None:
             tooltip="Transform points layer to atlas space",
         )
 
+    def add_brainrender_export_button(self, row: int, column: int) -> None:
+        """
+        Add a button to export the points in atlas space in the brainrender
+        format.
+
+        Parameters
+        ----------
+        row : int
+            Row in the grid layout.
+        column : int
+            Column in the grid layout.
+        """
+        self.brainrender_export_button = add_button(
+            "Export to brainrender",
+            self.layout,
+            self.export_points_to_brainrender,
+            row=row,
+            column=column,
+            visibility=False,
+            tooltip="Export points in atlas space to brainrender",
+        )
+
     def add_points_summary_table(self, row: int, column: int) -> None:
         """
         Add a table to display the summary of points per atlas region.
@@ -565,6 +588,7 @@ def analyse_points(self) -> None:
         )
 
         self.populate_summary_table()
+        self.brainrender_export_button.setVisible(True)
         self.save_all_points_button.setVisible(True)
         self.save_points_summary_button.setVisible(True)
 
@@ -596,6 +620,24 @@ def populate_summary_table(
         self.points_per_region_table_title.setVisible(True)
         self.points_per_region_table.setVisible(True)
 
+    def export_points_to_brainrender(self) -> None:
+        """
+        Export points in the format required for brainrender
+        N.B. assumes atlas is isotropic
+        """
+        path, _ = QFileDialog.getSaveFileName(
+            self,
+            "Choose filename",
+            "",
+            "NumPy Files (*.npy)",
+        )
+
+        if path:
+            path = ensure_extension(path, ".npy")
+        export_points_to_brainrender(
+            self.points_in_atlas_space, self.atlas.resolution[0], path
+        )
+
     def save_all_points_csv(self) -> None:
         """
         Save the coordinate and atlas region
diff --git a/tests/tests/test_brainmapper/test_transform_widget.py b/tests/tests/test_brainmapper/test_transform_widget.py
index c4f8523..01ce648 100644
--- a/tests/tests/test_brainmapper/test_transform_widget.py
+++ b/tests/tests/test_brainmapper/test_transform_widget.py
@@ -494,6 +494,31 @@ def test_analysis(transformation_widget_with_transformed_points):
     )
 
 
+def test_export_to_brainmapper(
+    mocker,
+    transformation_widget_with_transformed_points,
+    tmp_path,
+):
+    mock_get_save_file_name = mocker.patch(
+        "brainglobe_utils.brainmapper.transform_widget.QFileDialog.getSaveFileName"
+    )
+
+    save_path = tmp_path / "points.npy"
+    mock_get_save_file_name.return_value = (save_path, "NumPy Files (*.npy)")
+
+    transformation_widget_with_transformed_points.export_points_to_brainrender()
+
+    # Ensure the file dialog was called
+    mock_get_save_file_name.assert_called_once_with(
+        transformation_widget_with_transformed_points,
+        "Choose filename",
+        "",
+        "NumPy Files (*.npy)",
+    )
+
+    assert save_path.exists()
+
+
 def test_save_df_to_csv(
     mocker,
     transformation_widget_with_transformed_points,