From 2319fa9cbf0a8c28e08a165e6ca873171457363e Mon Sep 17 00:00:00 2001
From: Alexander Steppke <alexander.steppke@psi.ch>
Date: Mon, 17 Jun 2024 17:03:17 +0200
Subject: [PATCH] collaborative helper extended

---
 AUTHORS.rst                       |  1 +
 src/cristallina/jupyter_helper.py | 23 ++++++++++++++++++++++-
 src/cristallina/plot.py           |  7 +++++--
 src/cristallina/utils.py          |  2 +-
 4 files changed, 29 insertions(+), 4 deletions(-)

diff --git a/AUTHORS.rst b/AUTHORS.rst
index 38632d0..b8adc09 100644
--- a/AUTHORS.rst
+++ b/AUTHORS.rst
@@ -3,3 +3,4 @@ Contributors
 ============
 
 * Alexander Steppke <alexander.steppke@psi.ch>
+* Jakub Vonka <jakub.vonka@psi.ch>
diff --git a/src/cristallina/jupyter_helper.py b/src/cristallina/jupyter_helper.py
index 95fb098..5dda451 100644
--- a/src/cristallina/jupyter_helper.py
+++ b/src/cristallina/jupyter_helper.py
@@ -13,6 +13,7 @@ This should work on the local consoles and also on RA (yet untested).
 
 from packaging.version import Version
 import random
+import os
 
 import jupyter_collaboration
 
@@ -27,5 +28,25 @@ class MySQLiteYStore(SQLiteYStore):
     """
      Custom SQL location for jupyter collaboration to avoid NFS locking issues. 
     """
-    suffix = random.randint(0, 1E9)
+    suffix = random.randint(1E9, 9E9)
     db_path = f"/tmp/ystore_{suffix}.db"
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        
+        self.db_path = MySQLiteYStore.db_path
+        self.log.info(f"Using custom SQLiteYStore with path {self.db_path}")
+
+        if self.document_ttl is not None:
+            self.log.info(f"Using database_ttl: {self.document_ttl}")
+            self.document_ttl = int(self.document_ttl)
+
+
+    def __del__(self):
+        # Delete the temporary database file upon shutdown
+        try:
+            if os.path.exists(self.db_path):
+                os.remove(self.db_path)
+        except OSError as e:
+            self.log.warning(f"Error when removing temporary database: {e.strerror}")
+
diff --git a/src/cristallina/plot.py b/src/cristallina/plot.py
index be31c1c..7d628fb 100644
--- a/src/cristallina/plot.py
+++ b/src/cristallina/plot.py
@@ -157,7 +157,8 @@ def plot_2d_channel(data: SFDataFiles, channel_name, ax=None):
     axis_styling(ax, channel_name, description)
 
 
-def plot_detector_image(image_data, title=None, comment=None, ax=None, rois=None, norms=None, log_colorscale=False, show_legend=True, **fig_kw):
+def plot_detector_image(image_data, title=None, comment=None, ax=None, rois=None, norms=None, 
+                        log_colorscale=False, show_legend=True, ax_colormap=matplotlib.colormaps['viridis'], **fig_kw):
     """
     Plots channel data for a channel that contains an image (2d array) of numeric values per pulse.
     Optional:
@@ -166,6 +167,7 @@ def plot_detector_image(image_data, title=None, comment=None, ax=None, rois=None
     - log_colorscale: True for a logarithmic colormap
     - title: Title of the plot
     - show_legend: True if the legend box should be drawn
+    - ax_colormap: a matplotlib colormap (viridis by default)
     """
 
     im = image_data
@@ -187,7 +189,8 @@ def plot_detector_image(image_data, title=None, comment=None, ax=None, rois=None
     else:
         norm = matplotlib.colors.Normalize(vmin=norms[0], vmax=norms[1])
 
-    ax.imshow(im, norm=norm)
+    ax.imshow(im, norm=norm, cmap=ax_colormap)
+
     ax.invert_yaxis()
 
     if rois is not None:
diff --git a/src/cristallina/utils.py b/src/cristallina/utils.py
index a2a46ed..906c939 100644
--- a/src/cristallina/utils.py
+++ b/src/cristallina/utils.py
@@ -355,7 +355,7 @@ def wait_for_run(run_number, total_num_steps, snd_file_path="/tmp/CantinaBand3.w
     """ Busy wait loop until run has completed all steps. Plays sound 
         when all files are written to disk.
     """
-    base_path = cr.utils.heuristic_extract_base_path()
+    base_path = heuristic_extract_base_path()
     data_path = Path(base_path) / f"run{run_number:0>4}/data"
 
     while True: