From 175a8983cdd457de269660d33087a7a20df2e255 Mon Sep 17 00:00:00 2001
From: Pey Lian Lim <2090236+pllim@users.noreply.github.com>
Date: Wed, 23 Jun 2021 18:10:57 -0400
Subject: [PATCH] WIP: add_markers initial implementation

---
 jdaviz/configs/imviz/helper.py | 64 +++++++++++++++++++++++++++++++++-
 notebooks/ImvizExample.ipynb   | 23 ++++++++++++
 2 files changed, 86 insertions(+), 1 deletion(-)

diff --git a/jdaviz/configs/imviz/helper.py b/jdaviz/configs/imviz/helper.py
index a699d7f872..332c7e4e8f 100644
--- a/jdaviz/configs/imviz/helper.py
+++ b/jdaviz/configs/imviz/helper.py
@@ -7,7 +7,7 @@
 from astropy.wcs import NoConvergence
 from astropy.wcs.wcsapi import BaseHighLevelWCS
 from echo import delay_callback
-from glue.core import BaseData
+from glue.core import BaseData, Data
 
 from jdaviz.core.events import SnackbarMessage
 from jdaviz.core.helpers import ConfigHelper
@@ -15,12 +15,20 @@
 __all__ = ['Imviz']
 
 ASTROPY_LT_4_3 = not minversion('astropy', '4.3')
+RESERVED_MARKER_SET_NAMES = ['all']
 
 
 class Imviz(ConfigHelper):
     """Imviz Helper class"""
     _default_configuration = 'imviz'
 
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+        # Markers
+        self._marktags = set()
+        self._default_mark_tag_name = 'default-marker-name'
+
     def load_data(self, data, parser_reference=None, **kwargs):
         """Load data into Imviz.
 
@@ -185,6 +193,60 @@ def offset_to(self, dx, dy, skycoord_offset=False):
                 viewer.state.x_max = viewer.state.x_min + width
                 viewer.state.y_max = viewer.state.y_min + height
 
+    def _validate_marker_name(self, marker_name):
+        """Raise an error if the marker_name is not allowed."""
+        if marker_name in RESERVED_MARKER_SET_NAMES:
+            raise ValueError(
+                f"The marker name {marker_name} is not allowed. Any name is "
+                f"allowed except these: {', '.join(RESERVED_MARKER_SET_NAMES)}")
+
+    def add_markers(self, table, x_colname='x', y_colname='y',
+                    skycoord_colname='coord', use_skycoord=False,
+                    marker_name=None):
+        """Creates markers in the image at given points.
+
+        Parameters
+        ----------
+        table : `~astropy.table.Table`
+            Table containing marker locations.
+
+        x_colname, y_colname : str
+            Column names for X and Y.
+            Coordinates must be 0-indexed.
+
+        skycoord_colname : str
+            Column name with ``SkyCoord`` objects.
+
+        use_skycoord : bool
+            If `True`, use ``skycoord_colname`` to mark.
+            Otherwise, use ``x_colname`` and ``y_colname``.
+
+        marker_name : str, optional
+            Name to assign the markers in the table. Providing a name
+            allows markers to be removed by name at a later time.
+
+        """
+        if marker_name is None:
+            marker_name = self._default_mark_tag_name
+
+        self._validate_marker_name(marker_name)
+        self._marktags.add(marker_name)
+        viewer = self.app.get_viewer("viewer-1")
+
+        # TODO: Is this really the correct Data to use?
+        image = viewer.state.reference_data
+
+        if use_skycoord:
+            raise NotImplementedError
+        else:
+            # TODO: Doing it this way is not compatible with blink
+            t_glue = Data(marker_name, **table[x_colname, y_colname])
+            jglue = self.app.session.application
+            jglue.data_collection[marker_name] = t_glue
+            jglue.add_link(t_glue, x_colname, image, image.pixel_component_ids[1].label)
+            jglue.add_link(t_glue, y_colname, image, image.pixel_component_ids[0].label)
+            viewer.add_data(t_glue)
+
 
 def split_filename_with_fits_ext(filename):
     """Split a ``filename[ext]`` input into filename and FITS extension.
diff --git a/notebooks/ImvizExample.ipynb b/notebooks/ImvizExample.ipynb
index c536191d2c..6a83b4f62f 100644
--- a/notebooks/ImvizExample.ipynb
+++ b/notebooks/ImvizExample.ipynb
@@ -314,6 +314,29 @@
     "imviz.offset_to(0.5 * u.arcsec, -1.5 * u.arcsec, skycoord_offset=True)"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "id": "11fab067-7428-4ce4-bc9b-f5462fe52e2a",
+   "metadata": {},
+   "source": [
+    "It is also possible to programmatically add non-interactive markers to the image."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "270a6dd2-ce21-457f-845e-19cf5405fee7",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import numpy as np\n",
+    "from astropy.table import Table\n",
+    "\n",
+    "t = Table({'x': np.random.randint(0, 500, 100),\n",
+    "           'y': np.random.randint(0, 500, 100)})\n",
+    "imviz.add_markers(t)"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,