From c5b0c984311f9c0a38d01fff82ca907d27685d64 Mon Sep 17 00:00:00 2001
From: Kyle Conroy <kyleconroy@gmail.com>
Date: Fri, 26 Jan 2024 15:55:04 -0500
Subject: [PATCH] feature-flagged wireframe for cubeviz spec extraction plan
 (#2676)

* feature-flagged wireframe for cubeviz spec extraction plan
* changelog entry
* fix original mark visibility
---
 CHANGES.rst                                   |   4 +
 jdaviz/components/plugin_section_header.vue   |  13 +-
 .../spectral_extraction.py                    |  59 +++-
 .../spectral_extraction.vue                   | 282 +++++++++++-------
 .../spectral_extraction.vue                   |   6 +-
 jdaviz/core/marks.py                          |   3 +-
 jdaviz/core/template_mixin.py                 |  17 +-
 7 files changed, 248 insertions(+), 136 deletions(-)

diff --git a/CHANGES.rst b/CHANGES.rst
index f67fd24711..73590fb05b 100644
--- a/CHANGES.rst
+++ b/CHANGES.rst
@@ -24,6 +24,8 @@ Cubeviz
 
 - Single-pixel subset tool now shows spectrum-at-spaxel on hover. [#2647]
 
+- Spectral extraction plugin re-organized into subsections to be more consistent with specviz2d. [#2676]
+
 Imviz
 ^^^^^
 
@@ -39,6 +41,8 @@ Specviz
 Specviz2d
 ^^^^^^^^^
 
+- Spectral extraction plugin: highlighting of active header section. [#2676]
+
 API Changes
 -----------
 
diff --git a/jdaviz/components/plugin_section_header.vue b/jdaviz/components/plugin_section_header.vue
index 9be0e0aa07..5f0e849b48 100644
--- a/jdaviz/components/plugin_section_header.vue
+++ b/jdaviz/components/plugin_section_header.vue
@@ -1,11 +1,17 @@
 <template>
-  <div class="strike text--secondary">
+  <div :class="active ? 'strike strike-active text--secondary': 'strike text--secondary'">
      <span>
        <slot></slot>
      </span>
   </div>
 </template>
 
+<script>
+module.exports = {
+  props: ['active'],
+};
+</script>
+
 
 <style scoped>
 .strike {
@@ -34,6 +40,11 @@
     background: rgba(0, 0, 0, 0.15);
 }
 
+.strike-active > span:before,
+.strike-active > span:after {
+    background: #c75d2c;  /* active orange */
+}
+
 .strike > span:before {
     right: 100%;
     margin-right: 15px;
diff --git a/jdaviz/configs/cubeviz/plugins/spectral_extraction/spectral_extraction.py b/jdaviz/configs/cubeviz/plugins/spectral_extraction/spectral_extraction.py
index c8380a09e4..1c67e03434 100644
--- a/jdaviz/configs/cubeviz/plugins/spectral_extraction/spectral_extraction.py
+++ b/jdaviz/configs/cubeviz/plugins/spectral_extraction/spectral_extraction.py
@@ -9,7 +9,7 @@
 from astropy.nddata import (
     NDDataArray, StdDevUncertainty, NDUncertainty
 )
-from traitlets import Bool, Float, List, Unicode, observe
+from traitlets import Any, Bool, Float, List, Unicode, observe
 
 from jdaviz.core.custom_traitlets import FloatHandleEmpty
 from jdaviz.core.events import SnackbarMessage, SliceWavelengthUpdatedMessage
@@ -18,6 +18,7 @@
                                         DatasetSelectMixin,
                                         SelectPluginComponent,
                                         ApertureSubsetSelectMixin,
+                                        ApertureSubsetSelect,
                                         AddResultsMixin,
                                         with_spinner)
 from jdaviz.core.user_api import PluginUserApi
@@ -52,11 +53,23 @@ class SpectralExtraction(PluginTemplateMixin, ApertureSubsetSelectMixin,
     uses_active_status = Bool(True).tag(sync=True)
 
     # feature flag for cone support
-    dev_cone_support = Bool(False).tag(sync=True)
+    dev_cone_support = Bool(False).tag(sync=True)  # when enabling: add entries to docstring
+    dev_bg_support = Bool(False).tag(sync=True)  # when enabling: add entries to docstring
+    dev_subpixel_support = Bool(False).tag(sync=True)  # when enabling: add entries to docstring
+
+    active_step = Unicode().tag(sync=True)
+
     wavelength_dependent = Bool(False).tag(sync=True)
     reference_wavelength = FloatHandleEmpty().tag(sync=True)
     slice_wavelength = Float().tag(sync=True)
 
+    bg_items = List([]).tag(sync=True)
+    bg_selected = Any('').tag(sync=True)
+    bg_scale_factor = Float(1).tag(sync=True)
+    bg_wavelength_dependent = Bool(False).tag(sync=True)
+
+    subpixel = Bool(False).tag(sync=True)
+
     function_items = List().tag(sync=True)
     function_selected = Unicode('Sum').tag(sync=True)
     filename = Unicode().tag(sync=True)
@@ -84,6 +97,14 @@ def __init__(self, *args, **kwargs):
         self.aperture.items = [{"label": "Entire Cube"}]
         self.aperture.select_default()
 
+        self.background = ApertureSubsetSelect(self,
+                                               'bg_items',
+                                               'bg_selected',
+                                               'bg_scale_factor',
+                                               dataset='dataset',
+                                               multiselect=None,
+                                               default_text='None')
+
         self.function = SelectPluginComponent(
             self,
             items='function_items',
@@ -111,27 +132,37 @@ def __init__(self, *args, **kwargs):
 
     @property
     def user_api(self):
-        return PluginUserApi(
-            self,
-            expose=(
-                'function', 'spatial_subset', 'aperture',
-                'add_results', 'collapse_to_spectrum'
-            )
-        )
+        expose = ['function', 'spatial_subset', 'aperture',
+                  'add_results', 'collapse_to_spectrum']
+        if self.dev_cone_support:
+            expose += ['wavelength_dependent', 'reference_wavelength']
+        if self.dev_bg_support:
+            expose += ['background', 'bg_wavelength_dependent']
+        if self.dev_subpixel_support:
+            expose += ['subpixel']
+
+        return PluginUserApi(self, expose=expose)
 
     @property
     @deprecated(since="3.9", alternative="aperture")
     def spatial_subset(self):
         return self.user_api.aperture
 
+    @observe('active_step')
+    def _active_step_changed(self, *args):
+        self.aperture._set_mark_visiblities(self.active_step in ('', 'ap', 'ext'))
+        self.background._set_mark_visiblities(self.active_step == 'bg')
+
     @property
     def slice_plugin(self):
         return self.app._jdaviz_helper.plugins['Slice']
 
-    @observe('wavelength_dependent')
+    @observe('wavelength_dependent', 'bg_wavelength_dependent')
     def _wavelength_dependent_changed(self, *args):
         if self.wavelength_dependent:
             self.reference_wavelength = self.slice_plugin.wavelength
+        else:
+            self.bg_wavelength_dependent = False
         # NOTE: this can be redundant in the case where reference_wavelength changed and triggers
         # the observe, but we need to ensure it is updated if reference_wavelength is unchanged
         self._update_mark_scale()
@@ -149,8 +180,12 @@ def vue_adopt_slice_as_reference(self, *args):
     def _update_mark_scale(self, *args):
         if not self.wavelength_dependent:
             self.aperture.scale_factor = 1.0
-            return
-        self.aperture.scale_factor = self.slice_wavelength/self.reference_wavelength
+        else:
+            self.aperture.scale_factor = self.slice_wavelength/self.reference_wavelength
+        if not self.bg_wavelength_dependent:
+            self.background.scale_factor = 1.0
+        else:
+            self.background.scale_factor = self.slice_wavelength/self.reference_wavelength
 
     @with_spinner()
     def collapse_to_spectrum(self, add_data=True, **kwargs):
diff --git a/jdaviz/configs/cubeviz/plugins/spectral_extraction/spectral_extraction.vue b/jdaviz/configs/cubeviz/plugins/spectral_extraction/spectral_extraction.vue
index 41695b4510..e80426f1b4 100644
--- a/jdaviz/configs/cubeviz/plugins/spectral_extraction/spectral_extraction.vue
+++ b/jdaviz/configs/cubeviz/plugins/spectral_extraction/spectral_extraction.vue
@@ -8,134 +8,192 @@
     :popout_button="popout_button"
     :disabled_msg="disabled_msg">
 
-    <v-row>
-      <v-select
-        :menu-props="{ left: true }"
-        attach
-        :items="function_items.map(i => i.label)"
-        v-model="function_selected"
-        label="Function"
-        hint="Function for reducing dimensions of spectral cube."
-        persistent-hint
-      ></v-select>
-    </v-row>
-
-    <plugin-subset-select 
-      :items="aperture_items"
-      :selected.sync="aperture_selected"
-      :show_if_single_entry="true"
-      label="Spatial aperture"
-      hint="Select a spatial region to extract its spectrum."
-    />
-
-    <div v-if="aperture_selected !== 'Entire Cube' && dev_cone_support">
-      <v-alert type='warning'>cone support is under active development and hidden from users</v-alert>
-      <v-row>
-        <v-switch
-          v-model="wavelength_dependent"
-          label="Wavelength dependent"
-          hint="Vary aperture linearly with wavelength"
-          persistent-hint>
-        </v-switch>
-      </v-row>
-      <div v-if="wavelength_dependent">
-        <v-row justify="end">
-          <j-tooltip tooltipcontent="Adopt the current slice as the reference wavelength">
-            <plugin-action-button :results_isolated_to_plugin="true" @click="adopt_slice_as_reference">
-              Adopt Current Slice
-            </plugin-action-button>
-          </j-tooltip>
-        </v-row>
+    <div @mouseover="() => active_step='ap'">
+      <j-plugin-section-header :active="active_step==='ap'">Aperture</j-plugin-section-header>
+
+      <plugin-subset-select 
+        :items="aperture_items"
+        :selected.sync="aperture_selected"
+        :show_if_single_entry="true"
+        label="Spatial aperture"
+        hint="Select a spatial region to extract its spectrum."
+      />
+
+      <div v-if="aperture_selected !== 'Entire Cube' && dev_cone_support">
         <v-row>
-          <v-text-field
-            v-model.number="reference_wavelength"
-            type="number"
-            :step="0.1"
-            class="mt-0 pt-0"
-            label="Wavelength"
-            hint="Wavelength at which the aperture matches the selected subset."
-            persistent-hint
-          ></v-text-field>
-        </v-row>
-        <v-row justify="end">
-          <j-tooltip tooltipcontent="Select the slice nearest the reference wavelength">
-            <plugin-action-button :results_isolated_to_plugin="true" @click="goto_reference_wavelength">
-              Slice to Reference Wavelength
-            </plugin-action-button>
-          </j-tooltip>
+          <v-switch
+            v-model="wavelength_dependent"
+            label="Wavelength dependent"
+            hint="Vary aperture linearly with wavelength"
+            persistent-hint>
+          </v-switch>
         </v-row>
-      </div>
-
-    </div>
-
-    <plugin-add-results
-      :label.sync="results_label"
-      :label_default="results_label_default"
-      :label_auto.sync="results_label_auto"
-      :label_invalid_msg="results_label_invalid_msg"
-      :label_overwrite="results_label_overwrite"
-      label_hint="Label for the extracted spectrum"
-      :add_to_viewer_items="add_to_viewer_items"
-      :add_to_viewer_selected.sync="add_to_viewer_selected"
-      action_label="Extract"
-      action_tooltip="Run spectral extraction with error and mask propagation"
-      :action_spinner="spinner"
-      @click:action="spectral_extraction"
-    ></plugin-add-results>
-
-    <j-plugin-section-header v-if="extracted_spec_available && export_enabled">Results</j-plugin-section-header>
-
-    <div style="display: grid; position: relative"> <!-- overlay container -->
-      <div style="grid-area: 1/1">
-        <div v-if="extracted_spec_available && export_enabled">
-
+        <div v-if="wavelength_dependent">
+          <v-row justify="end">
+            <j-tooltip tooltipcontent="Adopt the current slice as the reference wavelength">
+              <plugin-action-button :results_isolated_to_plugin="true" @click="adopt_slice_as_reference">
+                Adopt Current Slice
+              </plugin-action-button>
+            </j-tooltip>
+          </v-row>
           <v-row>
             <v-text-field
-            v-model="filename"
-            label="Filename"
-            hint="Export the latest extracted spectrum."
-            :rules="[() => !!filename || 'This field is required']"
-            persistent-hint>
-            </v-text-field>
+              v-model.number="reference_wavelength"
+              type="number"
+              :step="0.1"
+              class="mt-0 pt-0"
+              label="Wavelength"
+              hint="Wavelength at which the aperture matches the selected subset."
+              persistent-hint
+            ></v-text-field>
           </v-row>
-
           <v-row justify="end">
-            <j-tooltip tipid='plugin-extract-save-fits'>
-              <v-btn color="primary" text @click="save_as_fits">Save as FITS</v-btn>
-
+            <j-tooltip tooltipcontent="Select the slice nearest the reference wavelength">
+              <plugin-action-button :results_isolated_to_plugin="true" @click="goto_reference_wavelength">
+                Slice to Reference Wavelength
+              </plugin-action-button>
             </j-tooltip>
           </v-row>
-
         </div>
       </div>
+    </div>
 
-      <v-overlay
-        absolute
-        opacity=1.0
-        :value="overwrite_warn && export_enabled"
-        :zIndex=3
-        style="grid-area: 1/1; 
-               margin-left: -24px;
-               margin-right: -24px">
-      
-      <v-card color="transparent" elevation=0 >
-        <v-card-text width="100%">
-          <div class="white--text">
-            A file with this name is already on disk. Overwrite?
-          </div>
-        </v-card-text>
+    <div v-if="dev_bg_support" @mouseover="() => active_step='bg'">
+      <j-plugin-section-header :active="active_step==='bg'">Background</j-plugin-section-header>
+      <plugin-subset-select 
+        :items="bg_items"
+        :selected.sync="bg_selected"
+        :show_if_single_entry="true"
+        label="Background"
+        hint="Select a spatial region to use for background subtraction."
+      />
+
+      <v-row v-if="aperture_selected === bg_selected">
+        <span class="v-messages v-messages__message text--secondary" style="color: red !important">
+            Background and aperture cannot be set to the same subset
+        </span>
+      </v-row>
 
-        <v-card-actions>
-          <v-row justify="end">
-            <v-btn tile small color="primary" class="mr-2" @click="overwrite_warn=false">Cancel</v-btn>
-            <v-btn tile small color="accent" class="mr-4" @click="overwrite_fits" >Overwrite</v-btn>
+      <div v-if="dev_cone_support && wavelength_dependent">
+        <v-row>
+          <v-switch
+            v-model="bg_wavelength_dependent"
+            label="Wavelength dependent"
+            hint="Vary background linearly with wavelength"
+            persistent-hint>
+          </v-switch>
+        </v-row>
+        <div v-if="bg_wavelength_dependent">
+          <v-row>
+            <v-text-field
+              v-model.number="reference_wavelength"
+              type="number"
+              :step="0.1"
+              class="mt-0 pt-0"
+              label="Wavelength"
+              hint="Wavelength at which the background matches the selected subset (fixed at same value as for aperture above)."
+              persistent-hint
+              disabled
+            ></v-text-field>
           </v-row>
-        </v-card-actions>
-      </v-card>
+        </div>
+      </div>
+    </div>
+
+    <div @mouseover="() => active_step='ext'">
+      <j-plugin-section-header :active="active_step==='ext'">Extract</j-plugin-section-header>
+      <div v-if="dev_subpixel_support">
+        <v-row>
+          <v-switch
+            v-model="subpixel"
+            label="Subpixel"
+            hint="Extract spectrum using a subpixel aperture in place of the subset mask"
+            persistent-hint>
+          </v-switch>
+        </v-row>
+      </div>
+
+      <v-row>
+        <v-select
+          :menu-props="{ left: true }"
+          attach
+          :items="function_items.map(i => i.label)"
+          v-model="function_selected"
+          label="Function"
+          hint="Function for reducing dimensions of spectral cube."
+          persistent-hint
+        ></v-select>
+      </v-row>
 
-      </v-overlay>
 
+      <plugin-add-results
+        :label.sync="results_label"
+        :label_default="results_label_default"
+        :label_auto.sync="results_label_auto"
+        :label_invalid_msg="results_label_invalid_msg"
+        :label_overwrite="results_label_overwrite"
+        label_hint="Label for the extracted spectrum"
+        :add_to_viewer_items="add_to_viewer_items"
+        :add_to_viewer_selected.sync="add_to_viewer_selected"
+        action_label="Extract"
+        action_tooltip="Run spectral extraction with error and mask propagation"
+        :action_spinner="spinner"
+        @click:action="spectral_extraction"
+      ></plugin-add-results>
+
+      <j-plugin-section-header v-if="extracted_spec_available && export_enabled">Results</j-plugin-section-header>
+
+      <div style="display: grid; position: relative"> <!-- overlay container -->
+        <div style="grid-area: 1/1">
+          <div v-if="extracted_spec_available && export_enabled">
+
+            <v-row>
+              <v-text-field
+              v-model="filename"
+              label="Filename"
+              hint="Export the latest extracted spectrum."
+              :rules="[() => !!filename || 'This field is required']"
+              persistent-hint>
+              </v-text-field>
+            </v-row>
+
+            <v-row justify="end">
+              <j-tooltip tipid='plugin-extract-save-fits'>
+                <v-btn color="primary" text @click="save_as_fits">Save as FITS</v-btn>
+
+              </j-tooltip>
+            </v-row>
 
+          </div>
+        </div>
+
+        <v-overlay
+          absolute
+          opacity=1.0
+          :value="overwrite_warn && export_enabled"
+          :zIndex=3
+          style="grid-area: 1/1; 
+                 margin-left: -24px;
+                 margin-right: -24px">
+        
+        <v-card color="transparent" elevation=0 >
+          <v-card-text width="100%">
+            <div class="white--text">
+              A file with this name is already on disk. Overwrite?
+            </div>
+          </v-card-text>
+
+          <v-card-actions>
+            <v-row justify="end">
+              <v-btn tile small color="primary" class="mr-2" @click="overwrite_warn=false">Cancel</v-btn>
+              <v-btn tile small color="accent" class="mr-4" @click="overwrite_fits" >Overwrite</v-btn>
+            </v-row>
+          </v-card-actions>
+        </v-card>
+
+        </v-overlay>
+      </div>
     </div>
+
   </j-tray-plugin>
 </template>
diff --git a/jdaviz/configs/specviz2d/plugins/spectral_extraction/spectral_extraction.vue b/jdaviz/configs/specviz2d/plugins/spectral_extraction/spectral_extraction.vue
index 47cb055262..51b0f00ef8 100644
--- a/jdaviz/configs/specviz2d/plugins/spectral_extraction/spectral_extraction.vue
+++ b/jdaviz/configs/specviz2d/plugins/spectral_extraction/spectral_extraction.vue
@@ -29,7 +29,7 @@
 
 
     <div @mouseover="() => active_step='trace'">
-      <j-plugin-section-header>Trace</j-plugin-section-header>
+      <j-plugin-section-header :active="active_step==='trace'">Trace</j-plugin-section-header>
       <v-row>
         <j-docs-link>
           Create a trace for the spectrum.  See the <j-external-link link='https://specreduce.readthedocs.io/en/latest/#module-specreduce.tracing' linktext='specreduce docs'></j-external-link> for more details on the available trace types.
@@ -189,7 +189,7 @@
     </div>
 
     <div @mouseover="() => active_step='bg'">
-      <j-plugin-section-header>Background</j-plugin-section-header>
+      <j-plugin-section-header :active="active_step==='bg'">Background</j-plugin-section-header>
       <v-row>
         <j-docs-link>Create a background and/or background-subtracted image.</j-docs-link>
       </v-row>
@@ -352,7 +352,7 @@
     </div>
 
     <div @mouseover="() => active_step='ext'">
-      <j-plugin-section-header>Extraction</j-plugin-section-header>
+      <j-plugin-section-header :active="active_step==='ext'">Extraction</j-plugin-section-header>
       <v-row>
         <j-docs-link>
           Extract a 1D spectrum from a 2D spectrum.  See the <j-external-link link='https://specreduce.readthedocs.io/en/latest/#module-specreduce.extract' linktext='specreduce docs'></j-external-link> for more details on the available extraction methods.
diff --git a/jdaviz/core/marks.py b/jdaviz/core/marks.py
index b6eccb0ef3..535c0547d8 100644
--- a/jdaviz/core/marks.py
+++ b/jdaviz/core/marks.py
@@ -654,7 +654,8 @@ def overlay(self):
 
 
 class ApertureMark(PluginLine):
-    def __init__(self, viewer, **kwargs):
+    def __init__(self, viewer, id, **kwargs):
+        self._id = id
         super().__init__(viewer, **kwargs)
 
 
diff --git a/jdaviz/core/template_mixin.py b/jdaviz/core/template_mixin.py
index 152cc879b4..747c59d2c4 100644
--- a/jdaviz/core/template_mixin.py
+++ b/jdaviz/core/template_mixin.py
@@ -2006,7 +2006,7 @@ class ApertureSubsetSelect(SubsetSelect):
 
     """
     def __init__(self, plugin, items, selected, scale_factor, multiselect=None,
-                 dataset=None, viewers=None):
+                 dataset=None, viewers=None, default_text=None):
         """
         Parameters
         ----------
@@ -2035,7 +2035,7 @@ def __init__(self, plugin, items, selected, scale_factor, multiselect=None,
                          filters=['is_spatial', 'is_not_composite', 'is_not_annulus'],
                          dataset=dataset,
                          viewers=viewers,
-                         default_text=None)
+                         default_text=default_text)
 
         self.add_traitlets(scale_factor=scale_factor)
 
@@ -2057,9 +2057,12 @@ def _on_dataset_selected_changed(self, event):
         super()._on_dataset_selected_changed(event)
         self._update_mark_coords()
 
-    def _plugin_active_changed(self, *args):
+    def _set_mark_visiblities(self, visible):
         for mark in self.marks:
-            mark.visible = self.plugin.is_active
+            mark.visible = visible
+
+    def _plugin_active_changed(self, *args):
+        self._set_mark_visiblities(self.plugin.is_active)
 
     @property
     def image_viewers(self):
@@ -2068,13 +2071,12 @@ def image_viewers(self):
 
     @property
     def marks(self):
-        # NOTE: this will require additional logic if we want to support multiple independent
-        # ApertureSubsetSelect instances per-plugin
         all_aperture_marks = []
         for viewer in self.image_viewers:
             # search for existing mark
             matches = [mark for mark in viewer.figure.marks
-                       if isinstance(mark, ApertureMark)]
+                       if (isinstance(mark, ApertureMark) and
+                           mark._id == self._plugin_traitlets['selected'])]
             if len(matches):
                 all_aperture_marks += matches
                 continue
@@ -2083,6 +2085,7 @@ def marks(self):
 
             mark = ApertureMark(
                 viewer,
+                id=self._plugin_traitlets['selected'],
                 x=x_coords,
                 y=y_coords,
                 colors=['#c75109'],