Source code for blissoda.tomo.flint_tomo_imshow

from __future__ import annotations

import logging
import time
from collections import OrderedDict
from pathlib import Path
from typing import Optional

import numpy
import silx

from blissoda.tomo.utils import apply_labels
from blissoda.tomo.utils import compute_axes

from ..bliss_globals import current_session
from ..flint import capture_errors
from ..flint.plotter import BasePlotter
from ..import_utils import unavailable_class

try:
    from flint.viewers.custom_image.client import ImageView
except ImportError as ex:
    ImageView = unavailable_class(ex)

_logger = logging.getLogger(__name__)




class SingleSliceImshow(BasePlotter):
    """
    Manage a Flint window showing the most recent reconstructed slice with enhanced controls,
    using the flint DataPlot–based ImageView.
    """

    TITLE = "Last Reconstructed Slice"

    def __init__(self, history: int = 1) -> None:
        super().__init__(max_plots=history)
        self._cache: OrderedDict[str, numpy.ndarray] = OrderedDict()

    def _load_and_show(self, img_path):
        with silx.io.open(img_path) as h5In:
            img = h5In["entry0000/reconstruction/results/data"][:]
        self.set_image(numpy.squeeze(img))



    def monitor_and_display_slice(self, output_path: str) -> None:
        """
        Start monitoring the slice output directory and display the first reconstructed
        slice as soon as an HDF5 file becomes available.
        """

        output_path = Path(output_path)
        output_dir = output_path.parent.parent / "slices"
        file_prefix = output_path.stem
        if len(list(output_dir.glob(f"{file_prefix}*_xy_*.h*5"))) > 0:
            _logger.warning("Slice file already exists in %s", output_dir)
        img_path = self._wait_for_slice_file(output_dir, file_prefix=file_prefix)
        if img_path is not None:
            self._load_and_show(img_path)


    def _wait_for_slice_file(
        self,
        output_dir: Path,
        file_prefix: str,
        poll_s: float = 0.5,
        timeout_s: float = 600.0,
    ) -> Optional[str]:
        """Wait for an HDF5 slice file to appear in the output directory."""
        deadline = time.monotonic() + timeout_s

        while time.monotonic() < deadline:
            files = list(output_dir.glob(f"{file_prefix}*_xy_*.hdf5"))
            time.sleep(poll_s)
            if files:
                return str(files[0])

        _logger.warning("No slice HDF5 found in %s after %.1fs", output_dir, timeout_s)
        return None



    @capture_errors
    def set_image(self, image: numpy.ndarray) -> None:
        """
        Display a new image with physical axis limits based on pixel size and image center.
        """
        widget = self._get_plot(self.TITLE, ImageView)
        self._set_title(widget)

        x_axis, y_axis = compute_axes(image)
        origin = (float(x_axis[0]), float(y_axis[0]))

        dx = float(x_axis[1] - x_axis[0]) if len(x_axis) > 1 else 1.0
        dy = float(y_axis[1] - y_axis[0]) if len(y_axis) > 1 else 1.0
        scale = (dx, dy)

        widget.set_data(image, origin=origin, scale=scale)
        apply_labels(widget)

        self._cache[self.TITLE] = image
        self.purge_tasks()
        self._purge()


    def _set_title(self, widget: ImageView) -> None:
        """
        Set the plot window title using the current data filename from scan saving.
        """
        widget.title = current_session.scan_saving.data_filename

    def _purge(self) -> None:
        """
        Remove oldest images beyond history.
        """
        while len(self._cache) > self._max_plots:
            self._cache.popitem(last=False)