From da6bc3e47672c3d8f518d68e61614b0ae45fe695 Mon Sep 17 00:00:00 2001
From: Kyle Conroy <kyleconroy@gmail.com>
Date: Tue, 23 Jul 2024 08:17:01 -0400
Subject: [PATCH] WIP: exclude spatial subsets from legend/plot options

---
 lcviz/plugins/plot_options/plot_options.py | 14 ++++++++++++++
 lcviz/viewers.py                           | 14 ++++++++++++++
 2 files changed, 28 insertions(+)

diff --git a/lcviz/plugins/plot_options/plot_options.py b/lcviz/plugins/plot_options/plot_options.py
index fee6892..5f79444 100644
--- a/lcviz/plugins/plot_options/plot_options.py
+++ b/lcviz/plugins/plot_options/plot_options.py
@@ -3,6 +3,9 @@
 from traitlets import observe
 from jdaviz.configs.default.plugins import PlotOptions
 from jdaviz.core.registries import tray_registry
+from jdaviz.utils import get_subset_type
+
+from lcviz.viewers import CubeView
 
 __all__ = ['PlotOptions']
 
@@ -37,6 +40,17 @@ def __init__(self, *args, **kwargs):
     def _update_docs_link(self, *args):
         self.docs_link = f"https://lcviz.readthedocs.io/en/{self.vdocs}/plugins.html#plot-options"
 
+        def not_spatial_subset_in_scatter_viewer(lyr):
+            # note: have to check the classname instead of isinstance to avoid circular import
+            if np.any([isinstance(viewer, CubeView)
+                       for viewer in self.layer.viewer_objs]):
+                return True
+            # at this point, NO selected viewers are TPF Cube viewers,
+            # so we want to exclude spatial subsets
+            return get_subset_type(lyr) != 'spatial'
+
+        self.layer.add_filter(not_spatial_subset_in_scatter_viewer)
+
     def _default_tpf_stretch(
             self, vmin_percentile=5, vmax_percentile=99, tpf_viewer_reference='image'
     ):
diff --git a/lcviz/viewers.py b/lcviz/viewers.py
index a2e7910..66ef89a 100644
--- a/lcviz/viewers.py
+++ b/lcviz/viewers.py
@@ -150,6 +150,13 @@ def _apply_layer_defaults(self, layer_state):
             # increased size of binned results, by default
             layer_state.size = 5
 
+    def _layer_included_in_legend(self, layer, subset_type):
+        print("***", layer.layer.label, subset_type)
+        if subset_type == 'spatial':
+            # do not show spatial subsets in time or phase viewers
+            return False
+        return super()._layer_included_in_legend(layer, subset_type)
+
     def set_plot_axes(self):
         # set which components should be plotted
         dc = self.jdaviz_app.data_collection
@@ -360,6 +367,13 @@ def _on_layers_update(self, layers=None):
             if hasattr(layer, 'attribute') and layer.attribute != flux_comp:
                 layer.attribute = flux_comp
 
+    def _layer_included_in_legend(self, layer, subset_type):
+        print("***", layer.layer.label, subset_type)
+        if subset_type == 'spectral':  # NOTE: spectral here means xrange (i.e. not spatial)
+            # ONLY show spatial subsets in image/cube viewer
+            return False
+        return super()._layer_included_in_legend(layer, subset_type)
+
     def data(self, cls=None):
         # TODO: generalize upstream in jdaviz.
         # This method is generalized from