From 02fd6f568f69ff63c2df3dd51f31113465ef56d3 Mon Sep 17 00:00:00 2001
From: Carifio24 <carifio24@gmail.com>
Date: Thu, 8 Aug 2024 01:10:32 -0400
Subject: [PATCH] Fix bad function call in the Jupyter tool.

---
 glue_ar/jupyter/export_tool.py | 38 ++++++++++++++++++++++++++++++----
 1 file changed, 34 insertions(+), 4 deletions(-)

diff --git a/glue_ar/jupyter/export_tool.py b/glue_ar/jupyter/export_tool.py
index 485de30..b580f02 100644
--- a/glue_ar/jupyter/export_tool.py
+++ b/glue_ar/jupyter/export_tool.py
@@ -1,11 +1,16 @@
 from os import getcwd
-from os.path import exists
+from os.path import exists, splitext
+from typing import Dict, Iterable, Tuple
 
 from glue.config import viewer_tool
+from glue.core.state_objects import State
+from glue.viewers.common.state import LayerState
 from glue.viewers.common.tool import Tool
+from glue_vispy_viewers.volume.qt.volume_viewer import VispyVolumeViewerMixin
 
-from glue_ar.common.export import export_to_ar
-from glue_ar.utils import AR_ICON
+from glue_ar.common.export import export_viewer
+from glue_ar.common.export_options import ar_layer_export
+from glue_ar.utils import AR_ICON, export_label_for_layer, xyz_bounds
 
 import ipyvuetify as v  # noqa
 from ipywidgets import HBox, Layout  # noqa
@@ -88,5 +93,30 @@ def on_no_click(button, event, data):
             self.save_figure(filepath)
             self.viewer.output_widget.clear_output()
 
+    def _state_dictionary(self,
+                          layers: Iterable[LayerState],
+                          filetype: str) -> Dict[str, Tuple[str, State]]:
+        state_dict = {}
+        for layer in layers:
+            label = export_label_for_layer(layer)
+            layer_state_cls = type(layer.state)
+            states = ar_layer_export.export_state_classes(layer_state_cls)
+            method_names = ar_layer_export.method_names(layer_state_cls, filetype)
+            method = method_names[0]
+            state_cls = next(t[1] for t in states if t[0] == method)
+            state = state_cls()
+            state_dict[label] = (method, state)
+
+        return state_dict
+
     def save_figure(self, filepath):
-        export_to_ar(self.viewer, filepath, state_dict={}, compression="draco")
+        bounds = xyz_bounds(self.viewer.state, with_resolution=isinstance(self.viewer, VispyVolumeViewerMixin))
+        layers = [layer for layer in self.viewer.layers if layer.enabled and layer.state.visible]
+        filetype = splitext(filepath)[1][1:]
+        state_dict = self._state_dictionary(layers, filetype)
+        export_viewer(viewer_state=self.viewer.state,
+                      layer_states=[layer.state for layer in self.viewer.layers
+                                    if layer.enabled and layer.state.visible],
+                      bounds=bounds,
+                      state_dictionary=state_dict,
+                      filepath=filepath)