"""
This module provides a custom QWidget class (GTWidget) for use with the napari viewer.
It includes various functionalities to display synapse images, edit and create point and
label annotations, interconvert between points and labels, and save data in as .zarr.
"""
from typing import TYPE_CHECKING

from qtpy.QtWidgets import QLabel, QCheckBox, QSpinBox, QDoubleSpinBox, QGroupBox, QHBoxLayout, QVBoxLayout, QGridLayout, QPushButton, QWidget, QFileDialog, QComboBox, QLineEdit, QCompleter
from scipy.ndimage import gaussian_filter, distance_transform_edt, center_of_mass, map_coordinates
from skimage.feature import peak_local_max
from skimage.measure import label, regionprops
from skimage.segmentation import watershed
from skimage.filters import threshold_triangle, threshold_multiotsu
from napari.utils.notifications import show_error, show_info
from qtpy.QtGui import QFontMetrics
from napari.qt.threading import thread_worker

from ._widget_utils import _setup_spin, _limitStretch, create_resolution_group
import os
import numpy as np
import zarr
import tifffile

if TYPE_CHECKING:
    import napari
                
class GTWidget(QWidget):
    # your QWidget.__init__ can optionally request the napari viewer instance
    # use a type annotation of 'napari.viewer.Viewer' for any parameter
    def __init__(self, viewer: "napari.viewer.Viewer", parent_widget=None):
        super().__init__()
        self.viewer = viewer
        self.parent_widget = parent_widget
        self.init_ui()

    def init_ui(self):
        self.setLayout(QVBoxLayout())
        self.setup_image_box()
        self.setup_points_box()
        self.setup_labels_box()
        self.setup_save_box()

        self.update_layer_choices()
        self.viewer.layers.events.inserted.connect(self.update_layer_choices)
        self.viewer.layers.events.removed.connect(self.update_layer_choices)
        self.update_possible_functions()

    def setup_image_box(self):
        res_group, self.xyresbox, self.zresbox, self.z_scale = create_resolution_group(self)
        self.layout().addWidget(res_group)

        self.active_image = QComboBox()        
        box2 = QGroupBox('Image tools')
        _limitStretch(self.active_image)
        
        self.active_image.currentTextChanged.connect(lambda: self._read_res())
        self.active_image.currentTextChanged.connect(lambda: self.threshbox.setMaximum(
            int(self.viewer.layers[self.active_image.currentText()].data.max())) if self.active_image.currentText() in self.viewer.layers else None)
        self.active_image.currentTextChanged.connect(self.update_possible_functions)
        
        image_gbox = QVBoxLayout()
        layer_select = QHBoxLayout()
        layer_select.addWidget(QLabel('presynaptic layer:'), stretch=1) ; layer_select.addWidget(self.active_image, stretch=3)
        image_gbox.addLayout(layer_select)

        box2.setLayout(image_gbox)

        self.layout().addWidget(box2)

    def setup_points_box(self):
        # Point tools box ################################################################
        self.active_points = QComboBox()
        self.active_points.currentTextChanged.connect(self._connect_point_events)
        self.active_points.currentTextChanged.connect(self.update_possible_functions)
        
        box3 = QGroupBox('Points tools')
        ptsbtn = QPushButton("New points layer")
        self.pksbtn = QPushButton("Find peaks above:")
        self.threshbox = QSpinBox(); 
        self.guessbtn = QPushButton("Guess")

        self.snapcheck = QCheckBox("Snap to max"); self.snap_to_max = True; self.snapcheck.setChecked(self.snap_to_max)
        self.snapbox = QSpinBox(); 
        _setup_spin(self, self.snapbox,   minval=0, suff=' px', val=3, attrname='snap_rad')
        
        self.snapcheck.stateChanged.connect(lambda state: self.snapbox.setEnabled(state==2))
        self.snapcheck.stateChanged.connect(lambda state: setattr(self, 'snap_to_max', state==2))
        
        _setup_spin(self, self.threshbox, minval=0, maxval=65535, val=10000, attrname='thresh')
        ptsbtn.clicked.connect(self._new_pts)
        self.pksbtn.clicked.connect(self._find_peaks)
        self.guessbtn.clicked.connect(self._calc_pk_thresh)

        points_gbox = QGridLayout()
        points_gbox.addWidget(ptsbtn, 0, 0, 1, 2)
        points_gbox.addWidget(QLabel('points layer:'), 1, 0) ; points_gbox.addWidget(self.active_points, 1, 1)
        pks_gbox = QHBoxLayout()
        pks_gbox.addWidget(self.pksbtn)
        pks_gbox.addWidget(self.threshbox)
        pks_gbox.addWidget(self.guessbtn)
        points_gbox.addLayout(pks_gbox, 2, 0, 1, 2)

        points_gbox.addWidget(self.snapcheck, 3, 0, 1, 1)
        points_gbox.addWidget(self.snapbox, 3, 1, 1, 1)

        box3.setLayout(points_gbox)

        # Points2labels box ##############################################################
        box5 = QGroupBox('Points to labels')

        self.p2lbtn = QPushButton("Points to labels")
        advbtn = QPushButton("Advanced settings"); advbtn.setCheckable(True)
                
        # Advanced settings
        box5b = QGroupBox('Advanced settings')
        self.radxybox = QSpinBox(); 
        self.radzbox = QSpinBox(); 
        self.mradxybox = QSpinBox() ; 
        self.mradzbox = QSpinBox(); 
        self.sigxybox = QDoubleSpinBox(); 
        self.sigzbox = QDoubleSpinBox(); 
        self.solidbox = QDoubleSpinBox();
        self.thresholdbox = QDoubleSpinBox();
        self.wshedcombo = QComboBox();
        self.wshed_type = 'Image'

        _setup_spin(self, self.radxybox,  minval=1, suff=' px', val=6, attrname='rad_xy')
        _setup_spin(self, self.radzbox,   minval=0, suff=' px', val=4, attrname='rad_z')
        _setup_spin(self, self.mradxybox, minval=0, suff=' px', val=1, attrname='max_rad_xy')
        _setup_spin(self, self.mradzbox,  minval=0, suff=' px', val=1, attrname='max_rad_z')
        _setup_spin(self, self.sigxybox,  minval=0, suff=' px', val=0.5, step=0.1, attrname='blur_sig_xy', dtype=float)
        _setup_spin(self, self.sigzbox,   minval=0, suff=' px', val=0.2, step=0.1, attrname='blur_sig_z', dtype=float)
        _setup_spin(self, self.solidbox,  minval=0, maxval=1, val=0.6, step=0.05, attrname='solidity_thresh', dtype=float)
        _setup_spin(self, self.thresholdbox,  minval=0, maxval=1, val=0.5, step=0.05, attrname='threshold', dtype=float)
        self.wshedcombo.addItem('Image'); self.wshedcombo.addItem('Distance')
        self.wshedcombo.currentTextChanged.connect(lambda name: setattr(self, 'wshed_type', name))

        
        self.p2lbtn.clicked.connect(self._points2labels)
        
        box5b.setVisible(False)
        advbtn.toggled.connect(box5b.setVisible)
        
        gbox5b = QGridLayout()
        gbox5b.addWidget(QLabel('threshold:'), 0, 0); gbox5b.addWidget(self.thresholdbox, 0, 1)
        gbox5b.addWidget(QLabel('segment xy rad:'), 1, 0); gbox5b.addWidget(self.radxybox, 1, 1)
        gbox5b.addWidget(QLabel('segment z rad:'), 2, 0); gbox5b.addWidget(self.radzbox, 2, 1)
        gbox5b.addWidget(QLabel('local max xy rad:'), 3, 0); gbox5b.addWidget(self.mradxybox, 3, 1)
        gbox5b.addWidget(QLabel('local max z rad:'), 4, 0); gbox5b.addWidget(self.mradzbox, 4, 1)
        gbox5b.addWidget(QLabel('gaussian xy rad:'), 5, 0); gbox5b.addWidget(self.sigxybox, 5, 1)
        gbox5b.addWidget(QLabel('gaussian z rad:'), 6, 0); gbox5b.addWidget(self.sigzbox, 6, 1)
        gbox5b.addWidget(QLabel('solidity:'), 7, 0); gbox5b.addWidget(self.solidbox, 7, 1)
        gbox5b.addWidget(QLabel('watershed type:'), 8, 0); gbox5b.addWidget(self.wshedcombo, 8, 1)
        box5b.setLayout(gbox5b)

        p2l_gbox = QGridLayout()
        p2l_gbox.addWidget(self.p2lbtn, 4, 0, 1, 2)  
        p2l_gbox.addWidget(advbtn, 5, 0, 1, 2)
        p2l_gbox.addWidget(box5b, 6, 0, 1, 2)
        
        box5.setLayout(p2l_gbox)
        
        self.layout().addWidget(box3)
        self.layout().addWidget(box5)

    def setup_labels_box(self):
        self.active_label = QComboBox()
    
        # Label tools box ################################################################
        box4 = QGroupBox('Labels tools')
        self.labcheck = QCheckBox("Make labels editable")
        self.labcheck.setTristate(False); self.labcheck.setCheckState(False)
        self.labcheck.stateChanged.connect(self._set_editable)
        self.active_label.currentTextChanged.connect(self._set_editable)
        self.labelbox = QSpinBox(); 
        self.addbtn = QPushButton("New label")
        self.addbtn.clicked.connect(self._add_label)
        self.removebtn = QPushButton("Remove label")
        self.maxlab = QLabel("max label: ")
        self.active_label.currentTextChanged.connect(self._set_max_label)
        self.active_label.currentTextChanged.connect(self._connect_label_events)
        self.active_label.currentTextChanged.connect(self.update_possible_functions)

        self.active_merge_label = QComboBox(); 
        self.mlsbtn = QPushButton("Merge labels")        
        self.l2pbtn = QPushButton("Labels to points")
        self.selectLabelsbtn = QPushButton("Keep labels from points")

        self.mlsbtn.clicked.connect(self._merge_labels)
        self.mlsbtn.clicked.connect(self._set_max_label)
        self.removebtn.clicked.connect(self._remove_label)
        self.l2pbtn.clicked.connect(self._labels2points)
        self.selectLabelsbtn.clicked.connect(self._remove_labels_wo_points)
        _setup_spin(self, self.labelbox, minval=1, maxval=1000, val=1, attrname='rem_label')

        labels_gbox = QGridLayout()
        labels_gbox.addWidget(QLabel('labels layer:'), 0, 0) ; labels_gbox.addWidget(self.active_label, 0, 1)
        labels_gbox.addWidget(self.labcheck, 1, 0); labels_gbox.addWidget(self.addbtn, 1, 1)
        labels_gbox.addWidget(self.labelbox, 2, 0)
        labels_gbox.addWidget(self.removebtn, 2, 1)
        labels_gbox.addWidget(QLabel('merge labels:'), 4, 0); 
        labels_gbox.addWidget(self.active_merge_label, 4, 1); 
        labels_gbox.addWidget(self.mlsbtn, 5, 0, 1, 2)
        labels_gbox.addWidget(self.l2pbtn, 6, 0, 1, 2)
        labels_gbox.addWidget(self.selectLabelsbtn, 7, 0, 1, 2)

        box4.setLayout(labels_gbox)        
        self.layout().addWidget(box4)

    def update_possible_functions(self, event=None):
        # IMAGE FUNCS
        if self.active_image.currentText() == '':
            self.pksbtn.setEnabled(False)
            self.guessbtn.setEnabled(False)
            self.from_source_button.setEnabled(False)
            self.threshbox.setEnabled(False)
        else:
            self.pksbtn.setEnabled(True)
            self.guessbtn.setEnabled(True)
            self.from_source_button.setEnabled(True)
            self.pksbtn.setEnabled(True)
            self.threshbox.setEnabled(True)

        # IMAGE + POINTS FUNCS
        if self.active_image.currentText() == '' or self.active_points.currentText() == '':
            self.p2lbtn.setEnabled(False)
        else:
            self.p2lbtn.setEnabled(True)
        
        # LABEL FUNCS
        if self.active_label.currentText() == '':
            self.labcheck.setEnabled(False)
            self.l2pbtn.setEnabled(False)
            self.addbtn.setEnabled(False)
            self.removebtn.setEnabled(False)
            self.labelbox.setEnabled(False)
        else:
            self.labcheck.setEnabled(True)
            self.l2pbtn.setEnabled(True)
            self.addbtn.setEnabled(True)
            self.removebtn.setEnabled(True)
            self.labelbox.setEnabled(True)

        # LABEL + POINTS FUNCS
        if self.active_label.currentText() == '' or self.active_points.currentText() == '':
            self.selectLabelsbtn.setEnabled(False)
        else:
            self.selectLabelsbtn.setEnabled(True)

        # SAVE FUNCS
        if self.active_image.currentText() == '' and self.active_label.currentText() == '':
            self.save3Dbtn.setEnabled(False)
        else:
            self.save3Dbtn.setEnabled(True)
        return 

    def update_layer_choices(self, event=None):
        img_layers = [l.name for l in self.viewer.layers if l.__class__.__name__ == "Image"]
        label_layers = [l.name for l in self.viewer.layers if l.__class__.__name__ == "Labels"]
        point_layers = [l.name for l in self.viewer.layers if l.__class__.__name__ == "Points"]
        
        img_choice = self.active_image.currentText()
        pts_choice = self.active_points.currentText()
        label_choice = self.active_label.currentText()
        if len(label_layers) <= 1:
            self.active_merge_label.setEnabled(False)
            self.mlsbtn.setEnabled(False)
        else:
            self.active_merge_label.setEnabled(True)
            self.mlsbtn.setEnabled(True)
        
        merge_label_choice = self.active_merge_label.currentText()
        
        for combo in (self.active_image, self.active_points, self.active_label, self.active_merge_label):
            combo.clear()
        
        self.active_image.addItems(img_layers)
        self.active_points.addItems(point_layers)
        self.active_label.addItems(label_layers)
        self.active_merge_label.addItems(label_layers)

        if img_choice in img_layers:
            self.active_image.setCurrentText(img_choice)
        if pts_choice in point_layers:
            self.active_points.setCurrentText(pts_choice)
        if label_choice in label_layers:
            self.active_label.setCurrentText(label_choice)
        if merge_label_choice in label_layers:
            self.active_merge_label.setCurrentText(merge_label_choice)

        for layer in self.viewer.layers:
            # disconnect first:
            try:
                layer.events.name.disconnect(self.update_layer_choices)
            except (TypeError, RuntimeError, ValueError):
                pass
            layer.events.name.connect(self.update_layer_choices)

    def _add_label(self):
        ''' set selected label to the next available number '''
        lab = self.active_label.currentText()
        if lab == '':
            return
        
        n = self.viewer.layers[lab].data.max()
        self.viewer.layers[lab].selected_label = (n + 1)
    
    def _connect_point_events(self):
        # remove connection first:
        for layer in self.viewer.layers:
            if layer.__class__.__name__ == "Points" and layer.name != self.active_points.currentText():
                if self._mouse_click in layer.mouse_drag_callbacks:
                    layer.mouse_drag_callbacks.remove(self._mouse_click)
        # add only to active points layer:
        pt_layer = self.viewer.layers[self.active_points.currentText()] if self.active_points.currentText() in self.viewer.layers else None
        if pt_layer is None:
            return
        pt_layer.mouse_drag_callbacks.append(self._mouse_click)

    def _mouse_click(self, layer, event):
        # ADAPTED FROM https://github.com/bburan/napari-synaptogram
        if layer.mode != "pan_zoom":
            # A tool (e.g., add, remove, select) is being used. Don't interfere
            # with what's already going on otherwise we may end up with
            # duplicate points or delete two points.
            return
        if event.buttons[0] == 2:
            if "Shift" in event.modifiers:
                self._remove_point(layer, event)
            else:
                self._add_point(layer, event)

    def _remove_point(self, layer, event):
        # This seems to handle hit testing. Not thrilled about using an
        # internal API, but it's styled after
        # napari.layers.points._points_mouse_bindings.
        point_index = layer._get_value_(
            position=event.position,
            view_direction=event.view_direction,
            dims_displayed=event.dims_displayed,
            world=True,
        )
        if point_index is not None:
            layer.pop(point_index)

    def _add_point(self, layer, event):
        image_layer = self.active_image.currentText()
        image_layer = self.viewer.layers[image_layer] if image_layer in self.viewer.layers else None
        if image_layer is None:
            show_error("Load image layer first.")
            return
                
        if self.viewer.dims.ndisplay == 2:
            # Logic for handling 2D view.
            near_point = list(event.position)
            far_point = list(event.position)
    
            # Find the axis to project the ray along.
            ray_axis = ({0, 1, 2} - set(self.viewer.dims.displayed)).pop()
    
            # Get the thickness of the view. The thickness is the full range
            # (lower to upper), but point click is in the center.
            thickness = self.viewer.dims.thickness[ray_axis] / 2
            near_point[ray_axis] += thickness
            far_point[ray_axis] -= thickness
            near_point = image_layer.world_to_data(near_point)
            far_point = image_layer.world_to_data(far_point)
                        
            if thickness == 0:
                layer.add(near_point)
                return
        else:
            # Logic for handling 3D view.
            # Find coordinates where ray enters/exists layer bounding box.
            near_point, far_point = image_layer.get_ray_intersections(
                event.position, event.view_direction, event.dims_displayed
            )
            if (near_point is None) or (far_point is None):
                show_error("Ray intersection failed - no valid near/far points")
                return
            
    
        try:
            num_samples = 25
            ray = np.linspace(near_point, far_point, num_samples, endpoint=True)
            
            ray_clamped = np.copy(ray)
            for i in range(ray.shape[1]):
                ray_clamped[:, i] = np.clip(ray[:, i], 0, image_layer.data.shape[i] - 1)
            
            intensities = map_coordinates(
                image_layer.data,
                ray_clamped.T,
                order=1,  
                mode="constant",
                cval=0,
                prefilter=False,  
            )
            
            max_idx = intensities.argmax()
            max_point = ray[max_idx]
            
            if self.snap_to_max and self.snap_rad > 0:
                initial_point = np.round(max_point).astype(int)
                search_region = []
                for dim in range(3):
                    start = max(0, initial_point[dim] - self.snap_rad)
                    end = min(image_layer.data.shape[dim], initial_point[dim] + self.snap_rad + 1)
                    search_region.append(slice(start, end))
                search_data = image_layer.data[tuple(search_region)]
                local_max_pt = np.unravel_index(np.argmax(search_data), search_data.shape)
                max_point = np.array([local_max_pt[dim] + search_region[dim].start for dim in range(3)])
            layer.add(max_point)
            
        except Exception as e:
            show_error(f"Error in _add_point: {e}")
            import traceback
            traceback.print_exc()


    def setup_save_box(self):
        box6 = QGroupBox('Save zarr')
        self.save3Dbtn = QPushButton("Save zarr")
        #save23Dbtn = QPushButton("Save 2D and 3D")
        browse_dir_button = QPushButton("\uD83D\uDCC1"); browse_dir_button.setToolTip("Browse for directory")
        browse_zarr_button = QPushButton("\uD83D\uDD0D"); browse_zarr_button.setToolTip("Find existing zarr")
        self.from_source_button = QPushButton("from source"); self.from_source_button.setToolTip("Select location from raw source");
        
        self.file_path_input = QLineEdit(self)
        self.file_name_input = QLineEdit(self)

        browse_dir_button.clicked.connect(self._browse_for_path)
        browse_zarr_button.clicked.connect(self._browse_for_zarr)
        self.file_name_input.editingFinished.connect(self._ensure_zarr_extension)
        self.file_path_input.textChanged.connect(self._update_completer)
        self.save3Dbtn.clicked.connect(lambda: self._save_zarr(threeD=False, twoD=False))
        #save23Dbtn.clicked.connect(lambda: self._save_zarr(threeD=True, twoD=True))
        self.from_source_button.clicked.connect(self._path_from_raw_source)

        box6.setLayout(QVBoxLayout())
        box6a = QGroupBox('File path'); box6a.setLayout(QHBoxLayout())
        box6a.layout().addWidget(self.file_path_input)
        box6a.layout().addWidget(browse_dir_button)
        box6.layout().addWidget(box6a)
        box6b = QGroupBox('File name'); box6b.setLayout(QHBoxLayout())
        box6b.layout().addWidget(self.file_name_input)
        box6b.layout().addWidget(browse_zarr_button)
        box6.layout().addWidget(box6b)        
        box6.layout().addWidget(self.from_source_button)
        box6.layout().addWidget(self.save3Dbtn)
        #box6.layout().addWidget(save23Dbtn, 3, 1, 1, 1)

        #self.layout().addWidget(box1)
        self.layout().addWidget(box6)
    # def _limitStretch(self, widget, max_chars=16):
    #     font_metrics = QFontMetrics(widget.font())
    #     width = font_metrics.horizontalAdvance('M') * max_chars
    #     widget.setMaximumWidth(width)

    def _path_from_raw_source(self):
        raw_path = self.viewer.layers[self.active_image.currentText()].source.path

        if raw_path.rfind('.zarr')>0:
            raw_path = raw_path[0:raw_path.rfind('.zarr')+5]
        
        directory, zarrfi = os.path.split(raw_path)

        zarrfi = os.path.splitext(zarrfi)[0]+'.zarr'

        self.file_path_input.setText(directory)
        self.file_name_input.setText(zarrfi)

    def _browse_for_path(self):
        directory = QFileDialog.getExistingDirectory(self, "Select Directory")
        if directory:
            self.file_path_input.setText(directory)

    def _browse_for_zarr(self):
        start_from = self.file_path_input.text()
        directory = QFileDialog.getExistingDirectory(self, "Select Directory", start_from)
        if directory:
            directory, zarrfi = os.path.split(directory)
            self.file_path_input.setText(directory)
            self.file_name_input.setText(zarrfi)

    def _ensure_zarr_extension(self):
        file_name = self.file_name_input.text()
        if not file_name.endswith(".zarr"):
            self.file_name_input.setText(file_name + ".zarr")
    
    def _update_completer(self):
        directory = self.file_path_input.text()
        if os.path.isdir(directory):
            files_in_dir = os.listdir(directory)
            completer = QCompleter(files_in_dir, self.file_name_input)
            #completer.setCaseSensitivity(Qt.CaseInsensitive)
            self.file_name_input.setCompleter(completer)

    def _save_zarr(self, threeD=True, twoD=False):

        fileName = os.path.join(self.file_path_input.text(), self.file_name_input.text())

        # TODO: add dialog box to warn if overwriting a file
        # TODO: ensure fileName ends with .zarr

        zarrfi = zarr.open(fileName)
        try:
            raw = self.viewer.layers[self.active_image.currentText()].data
        except:
            show_info("Image layer not defined. Image will not be saved.")
            raw = None
        try:
            labels = self.viewer.layers[self.active_label.currentText()].data
        except:
            show_info("Labels layer not defined. Labels will not be saved.")
            labels = None

        for (name, data, process_func) in (('raw', raw, self._process_raw), ('labeled', labels, self._process_labels)):
            if data is None:
                continue
            is_dask=True
            try:
                data.chunks
            except:
                is_dask=False

            if is_dask:
                zarrfi[f'{name}'] = process_func(data.compute())
            else:
                zarrfi[f'{name}'] = process_func(data)
            zarrfi[f'{name}'].attrs['offset'] = [0,]*3
            zarrfi[f'{name}'].attrs['resolution'] = [self.zres*1000, self.xyres*1000, self.xyres*1000] # saved in nm
            show_info(f'Saved {name} to {self.file_name_input.text()}')

    def _process_raw(self, rawdata):
        return normalize(rawdata.astype(np.uint16), 
                maxval=(2**16-1)).astype(np.uint16)

    def _process_labels(self, labeldata):
        return labeldata.astype(np.int64)

    def _convert_dask(self, layer_name):
        try:
            self.viewer.layers[layer_name].data.chunks
        except:
            # not a dask
            return
        show_info('converting layer '+layer_name+' to numpy array')
        self.viewer.layers[layer_name].data = self.viewer.layers[layer_name].data.compute()
        self.labcheck.setCheckState(2)

    def _set_editable(self):
        if self.labcheck.isChecked():
            try:
                label = self.active_label.currentText()
            except:
                return
            self._convert_dask(self.active_label.currentText())

    def _find_peaks(self):
        try:
            img = self.viewer.layers[self.active_image.currentText()].data
        except:
            show_error("Image layer not defined. No peaks detected.")
            return
        
        img_filtered = gaussian_filter(img, sigma=(0.7, 1, 1))

        thresh = self.thresh
        @thread_worker
        def _find_peaks_thread(img, thresh):
             return peak_local_max(img, threshold_abs=thresh, min_distance=2)
        
        def on_done(peaks):
            if len(peaks) == 0:
                show_info("No peaks found with the current threshold.")
                return
            self._new_pts()
            pts = self.viewer.layers[-1]    
            pts.data = np.array(peaks, dtype=np.float32)
            pts.refresh()
        
        worker = _find_peaks_thread(img_filtered, thresh)
        worker.returned.connect(on_done)
        worker.start()

    def _calc_pk_thresh(self):
        try:
            img = self.viewer.layers[self.active_image.currentText()].data
        except:
            show_error("Image layer not defined. Cannot calculate peak threshold.")
            return
        try:
            img = img.compute()
        except:
            pass
        img_max = np.max(img, axis=0)
        _, thresh = threshold_multiotsu(img_max)

        self.thresh = thresh
        self.threshbox.setValue(self.thresh)
        self._find_peaks()
    
    def _new_pts(self):
        self.viewer.add_points(
                ndim=3, 
                face_color='magenta', 
                border_color='white',
                size=15,
                out_of_slice_display=True,
                opacity=0.7,
                symbol='x')
        # set active points layer to the new layer
        self.active_points.setCurrentText(self.viewer.layers[-1].name)

    def _read_res(self):
        try:
            img = self.viewer.layers[self.active_image.currentText()]
        except:
            return
        imgpath = img.source.path

        if imgpath is not None and imgpath.endswith('.tif'):
            [z, y, x] = self._read_tiff_voxel_size(imgpath)
            self._update_xy_res(x)
            self._update_z_res(z)
            if self.z_scale.isChecked(): 
                self._scale_layers()
            # Emit to parent to sync other widgets
            if self.parent_widget:
                self.parent_widget.xy_res_changed.emit(x)
                self.parent_widget.z_res_changed.emit(z)
        if imgpath is not None and '.zarr' in imgpath:
            [z, y, x] = self._read_zarr_voxel_size(imgpath)
            self._update_xy_res(x)
            self._update_z_res(z)
            if self.z_scale.isChecked(): 
                self._scale_layers()
            # Emit to parent to sync other widgets
            if self.parent_widget:
                self.parent_widget.xy_res_changed.emit(x)
                self.parent_widget.z_res_changed.emit(z)

        # TODO: add functionality for .czi or other formats?
    
    def _read_zarr_voxel_size(self, file_path):
        res = zarr.open(file_path, mode='r').attrs.get('resolution', [1., 1., 1.])
        res = [r/1000 for r in res] # convert from nm to um
        return res

    def _read_tiff_voxel_size(self, file_path):
        """
        Implemented based on information found in https://pypi.org/project/tifffile
        """

        def _xy_voxel_size(tags, key):
            assert key in ['XResolution', 'YResolution']
            if key in tags:
                num_pixels, units = tags[key].value
                return units / num_pixels
            # return default
            return 1.

        with tifffile.TiffFile(file_path) as tiff:
            image_metadata = tiff.imagej_metadata
            if image_metadata is not None:
                z = image_metadata.get('spacing', 1.)
            else:
                # default voxel size
                z = 1.

            tags = tiff.pages[0].tags
            # parse X, Y resolution
            y = _xy_voxel_size(tags, 'YResolution')
            x = _xy_voxel_size(tags, 'XResolution')
            # return voxel size
            return [z, y, x]

    def _rxy(self):
        self.viewer.camera.angles = [self.viewer.camera.angles[0], 0, 90]
    
    def _auto_z(self):
        should_break=False
        try:
            pts = self.viewer.layers[self.active_points.currentText()]
        except:
            show_error("Points layer not defined. Cannot auto-adjust Z positions.")
            should_break=True
        try:
            img = self.viewer.layers[self.active_image.currentText()]
        except:
            show_error("Image layer not defined. Cannot auto-adjust Z positions.")
            should_break=True

        if not should_break:
            pts = self.viewer.layers[self.active_points.currentText()]
            img = self.viewer.layers[self.active_image.currentText()]
            # if len(pts.selected_data)>0:
            #     pt_list = pts.selected_data
            # else:
            pt_list = [i for i in range(len(pts.data))]
            for pt_id in pt_list: #s.data:
                coords = tuple(pts.data[pt_id].astype(int))
                maxz = np.argmax(img.data[:, coords[1], coords[2]])
                pts.data[pt_id, 0] = maxz
            pts.refresh()

    def _change_z(self, z):
        should_break=False
        try:
            self.viewer.layers[self.active_points.currentText()]
        except:
            show_error("Points layer not defined. Cannot change Z positions.")
            should_break=True
        if not should_break:
            for pt in self.viewer.layers[self.active_points.currentText()].selected_data:
                self.viewer.layers[self.active_points.currentText()].data[pt, 0] = self.viewer.layers[self.active_points.currentText()].data[pt, 0]+z
        
        self.viewer.layers[self.active_points.currentText()].refresh()

    def _connect_label_events(self):
        layer_name = self.active_label.currentText()
        if layer_name:
            try:
                self.viewer.layers[layer_name].events.labels_update.disconnect(self._set_max_label)
                self.viewer.layers[layer_name].events.selected_label.disconnect(self._on_selected_label)
            except:
                pass
            self.viewer.layers[layer_name].events.labels_update.connect(self._set_max_label)
            self.viewer.layers[layer_name].events.selected_label.connect(self._on_selected_label)

    def _on_selected_label(self, event):
        value = None
        if event is not None and hasattr(event, "value"):
            try:
                value = int(event.value)
            except (TypeError, ValueError):
                value = None
        if value is None:
            layer_name = self.active_label.currentText()
            if layer_name and layer_name in self.viewer.layers:
                try:
                    value = int(self.viewer.layers[layer_name].selected_label)
                except (TypeError, ValueError):
                    value = None
        if value is not None:
            self.labelbox.setValue(value)

    def _set_max_label(self):
        should_break=False
        try:
            self.viewer.layers[self.active_label.currentText()]
        except: # AttributeError:
            should_break=True
        if not should_break:
            try: #check if dask
                maxL = self.viewer.layers[self.active_label.currentText()].data.max().compute()
            except: #if not dask
                maxL = self.viewer.layers[self.active_label.currentText()].data.max()
            #maxL = self.viewer.layers[self.active_label.currentText()].data.max()
            self.labelbox.setMaximum(maxL)
            self.maxlab.setText("max label: "+str(maxL))
        
    def _remove_labels_wo_points(self):
        should_break=False
        try:
            labels = self.viewer.layers[self.active_label.currentText()]
        except:
            show_error("Labels layer not defined. Cannot remove labels.")
            should_break = True
        try:
            points = self.viewer.layers[self.active_points.currentText()]
        except:
            show_error("Points layer not defined. Labels to be removed undefined.")
            should_break = True
        if should_break:
            return
        
        mask = np.zeros_like(labels.data)
        for pt in points.data:
            selected_label = labels.data[tuple([int(i) for i in pt])]
            mask[labels.data==selected_label] = selected_label

        self.viewer.add_labels(mask, name='selected '+self.active_label.currentText())


    def _remove_label(self):
        should_break=False
        try:
            labels = self.viewer.layers[self.active_label.currentText()]
        except:
            show_error("Labels layer not defined. Cannot remove label.")
            should_break = True
        if should_break:
            return
        
        self._convert_dask(self.active_label.currentText())
        
        mask = labels.data == self.rem_label
        labels.data[mask] = 0
        labels.refresh()
        self._set_max_label()

    def _merge_labels(self):
        should_break=False
        try:
            labels = self.viewer.layers[self.active_label.currentText()]
            self._convert_dask(self.active_label.currentText())
        except:
            show_error("Labels layer not defined. Cannot merge labels.")
            should_break = True
        try:
            labels2 = self.viewer.layers[self.active_merge_label.currentText()]
            self._convert_dask(self.active_merge_label.currentText())
        except:
            show_error("Merge labels layer not defined. Cannot merge labels.")
            should_break = True
        if labels == labels2:
            show_error("Cannot merge labels with itself.")
            should_break = True
        if should_break:
            return

        max0 = labels.data.max()
        mask = labels2.data>0
        min0 = labels2.data[mask].min()
        if min0 <= max0:
            const = max0+1-min0
            labels2.data[mask] = labels2.data[mask]+const
        labels.data[mask] = labels2.data[mask][:]
        self._renumber_labels(self.active_label.currentText())
        labels.refresh()
        self._set_max_label()
        self.viewer.layers.remove(self.active_merge_label.currentText())
        self.active_merge_label.setCurrentIndex(-1)
        try:
            n = self.viewer.layers[self.active_points.currentText()].data.shape[0]
            self.viewer.layers[self.active_points.currentText()].selected_data = [i for i in range(n)]
            self.viewer.layers[self.active_points.currentText()].remove_selected()
            #self.viewer.layers.remove(self.active_points.currentText())
            #del self.active_points.currentText()
        except:
            pass

    def _renumber_labels(self, layer_name):
        labels_layer = self.viewer.layers[layer_name]
        self._convert_dask(layer_name)
        labels_data = labels_layer.data
        props = regionprops(labels_data)
        new_labels = np.zeros_like(labels_data)
        for new_label, prop in enumerate(props, start=1):
            new_labels[labels_data == prop.label] = new_label
        labels_layer.data = new_labels
        
    def _labels2points(self):
        try:
            labels = self.viewer.layers[self.active_label.currentText()].data
        except:
            show_error("Labels layer not defined. Labels to points function exited.")
            return
        # cannot operate on dask...
        self._convert_dask(self.active_label.currentText())
        labels = self.viewer.layers[self.active_label.currentText()].data

        mask = labels>0
        com = center_of_mass(mask, labels=labels, index=np.unique(labels[mask]))
        self.viewer.add_points(com,             
                ndim=3, 
                face_color='green', 
                border_color='white',
                size=5,
                out_of_slice_display=True,
                opacity=0.7,
                symbol='x',
                name='points from '+self.active_label.currentText())
            
    def _points2labels(self):
        should_break=False
        try:
            pts = self.viewer.layers[self.active_points.currentText()].data
        except:
            show_error("Points layer not defined. Cannot convert points to labels.")
            should_break=True
        try:
            img = self.viewer.layers[self.active_image.currentText()].data
        except:
            show_error("Image layer not defined. Cannot convert points to labels.")
            should_break=True
        try:
            labels = self.viewer.layers[self.active_label.currentText()].data
            # check if dask:
            try:
                curr_n = labels.max().compute()
            except:
                curr_n = labels.max()
        except:
            curr_n = 0
        
        if should_break:
            show_error("Points to labels function exited.")
            return
        
        blur_sig = [self.blur_sig_z, self.blur_sig_xy, self.blur_sig_xy]
        
        if np.any(blur_sig):
            img_inv = gaussian_filter(np.min(img) + np.max(img) - img, blur_sig)
        else:
            img_inv = np.min(img) + np.max(img) - img

        markers = np.zeros(img.shape, dtype='int')
        mask = np.zeros(img.shape, dtype='bool')

        # make markers:
        count = 1
        rem_list = []
        for pos in pts: 
            pos = np.round(pos).astype('int')
            try:
                markers[tuple(pos)] = count+curr_n
            except:
                show_info("Point "+str(pos)+" is outside image bounds. Skipping.")
                rem_list.append(count-1)
            count += 1
        if len(rem_list)>0:
            pts = np.delete(pts, rem_list, axis=0)
            # make invalid points red and large:
            self.viewer.layers[self.active_points.currentText()].face_color[rem_list] = [1,0,0,1]
            self.viewer.layers[self.active_points.currentText()].size[rem_list] = 5*self.viewer.layers[self.active_points.currentText()].size[rem_list]
            self.viewer.layers[self.active_points.currentText()].refresh()

        # make mask:
        count = 1
        for pos in pts:
            pos = np.round(pos).astype('int')
            pointIntensity = img_inv[tuple(pos)]
            
            # find local min (inverted max) value:
            zrange, yrange, xrange, rel_pos = self._get_slices(self.max_rad_xy, self.max_rad_z, pos, img.shape)
            subim = img_inv[zrange, yrange, xrange]        
            local_min = np.min(subim) 
            # get local region to threshold, find local min value:
            zrange, yrange, xrange, rel_pos = self._get_slices(self.rad_xy, self.rad_z, pos, img.shape)
            subim = img_inv[zrange, yrange, xrange]
            local_max = np.max(subim) # background
            
            # threshold:
            thresh = self.threshold*local_min+(1-self.threshold)*local_max 
            #print(local_min, local_max, thresh)
            if thresh < pointIntensity:
                #print("threshold overriden for spot "+str(j)+" "+str(thresh)+" "+str(pointIntensity))
                thresh = self.threshold*pointIntensity+(1-self.threshold)*local_max 
            subim_mask = subim <= thresh
            
            # check for multiple objects:
            sublabels = label(subim_mask)
            if sublabels.max() > 1:
                wantLabel = sublabels[tuple(rel_pos)]
                subim_mask = sublabels == wantLabel
                
                # recheck max:
                thresh2 = self.threshold*np.min(subim[subim_mask])+(1-self.threshold)*local_max
                if thresh < thresh2:
                    subim_mask = subim <= thresh2
                    sublabels = label(subim_mask)
                    wantLabel = sublabels[tuple(rel_pos)]
                    subim_mask = sublabels == wantLabel
            
            if self.solidity_thresh > 0:
                pt_solidity = regionprops(subim_mask.astype('int'))[0].solidity
            
                if pt_solidity < self.solidity_thresh:
                    subim_mask = self._dist_watershed_sep(subim_mask, rel_pos)
            submask = mask[zrange, yrange, xrange]
            submask = np.logical_or(submask, subim_mask)

            mask[zrange, yrange, xrange] = submask
            
        if self.wshed_type=='Image':
            outlabels = watershed(img_inv, markers=markers, mask=mask)
        elif self.wshed_type=='Distance':
            outlabels = watershed(
                    distance_transform_edt(markers==0, sampling=[3,1,1]),
                    markers=markers,
                    mask=mask)
        self.viewer.add_labels(outlabels, name='labels from '+self.active_points.currentText())
        
        self.active_merge_label.setCurrentText(self.viewer.layers[-1].name)


    def _get_slices(self, rad_xy, rad_z, loc, shape):
        '''
        Get slices for a box around loc with specified radii, staying 
        within shape bounds. 
        '''
        x1 = max(loc[2] - rad_xy, 0) 
        x2 = min(loc[2] + rad_xy+1, shape[2]) 
        y1 = max(loc[1] - rad_xy, 0) 
        y2 = min(loc[1] + rad_xy+1, shape[1]) 
        z1 = max(loc[0] - rad_z, 0) 
        z2 = min(loc[0] + rad_z+1, shape[0]) 
        relx = loc[2] - x1 
        rely = loc[1] - y1 
        relz = loc[0] - z1 
        
        return slice(z1,z2), slice(y1,y2), slice(x1,x2), [relz, rely, relx]
    
    def _dist_watershed_sep(self, mask, loc):
        """
        Watershed separation based on distance transform. After separation, 
        returns the mask region overlapping with location loc.
            mask : 3D binary mask 
            loc : (z, y, x) coordinates of location to keep 
        Returns:
            mask_out: 3D binary mask 
        """
        sampling = [self.zres/self.xyres, 1, 1]
        dists = distance_transform_edt(mask, sampling=sampling)
        pk_idx = peak_local_max(dists, labels=mask)
        pks = np.zeros_like(dists, dtype=bool)
        for pk in pk_idx:
            pks[tuple(pk)] = True
        pk_labels = label(pks)
        if pk_labels.max()>1:
            merged_peaks = center_of_mass(pks, pk_labels, index=range(1, pk_labels.max()+1))
            merged_peaks = np.round(merged_peaks).astype('int')
            
            markers = np.zeros_like(mask, dtype='int')
            for i in range(merged_peaks.shape[0]):
                markers[merged_peaks[i,0], merged_peaks[i,1], merged_peaks[i,2]] = i+1

            labels = watershed(-dists, markers=markers, mask=mask)
            wantLabel = labels[loc[0], loc[1], loc[2]] 
            mask_out = labels == wantLabel
        
        else:
            mask_out = mask
        
        return mask_out


    def to_settings(self):
        """
        Serialize current widget settings to a dictionary.
        
        Returns:
            Dictionary of all persistent settings.
        """
        return {
            'z_scale': self.z_scale.isChecked(),
            'thresh': self.thresh,
            'snap_to_max': self.snap_to_max,
            'snap_rad': self.snap_rad,
            'rad_xy': self.rad_xy,
            'rad_z': self.rad_z,
            'max_rad_xy': self.max_rad_xy,
            'max_rad_z': self.max_rad_z,
            'blur_sig_xy': self.blur_sig_xy,
            'blur_sig_z': self.blur_sig_z,
            'solidity_thresh': self.solidity_thresh,
            'threshold': self.threshold,
            'wshed_type': self.wshed_type,
            'file_path_input': self.file_path_input.text()
        }
    
    def apply_settings(self, settings):
        """
        Restore widget settings from a dictionary.
        
        Args:
            settings: Dictionary of settings to apply.
        """
        if 'z_scale' in settings:
            self.z_scale.setChecked(settings['z_scale'])
        if 'thresh' in settings:
            self.thresh = settings['thresh']
            self.threshbox.setValue(self.thresh)
        if 'snap_to_max' in settings:
            self.snap_to_max = settings['snap_to_max']
            self.snapcheck.setChecked(self.snap_to_max)
        if 'snap_rad' in settings:
            self.snap_rad = settings['snap_rad']
            self.snapbox.setValue(self.snap_rad)
        if 'rad_xy' in settings:
            self.rad_xy = settings['rad_xy']
            self.radxybox.setValue(settings['rad_xy'])
        if 'rad_z' in settings:
            self.rad_z = settings['rad_z']
            self.radzbox.setValue(settings['rad_z'])
        if 'max_rad_xy' in settings:
            self.max_rad_xy = settings['max_rad_xy']
            self.mradxybox.setValue(settings['max_rad_xy'])
        if 'max_rad_z' in settings:
            self.max_rad_z = settings['max_rad_z']
            self.mradzbox.setValue(settings['max_rad_z'])
        if 'blur_sig_xy' in settings:
            self.blur_sig_xy = settings['blur_sig_xy']
            self.sigxybox.setValue(settings['blur_sig_xy'])
        if 'blur_sig_z' in settings:
            self.blur_sig_z = settings['blur_sig_z']
            self.sigzbox.setValue(settings['blur_sig_z'])
        if 'solidity_thresh' in settings:
            self.solidity_thresh = settings['solidity_thresh']
            self.solidbox.setValue(settings['solidity_thresh'])
        if 'threshold' in settings:
            self.threshold = settings['threshold']
            self.thresholdbox.setValue(settings['threshold'])
        if 'wshed_type' in settings:
            self.wshed_type = settings['wshed_type']
            idx = self.wshedcombo.findText(settings['wshed_type'])
            if idx >= 0:
                self.wshedcombo.setCurrentIndex(idx)
        if 'file_path_input' in settings:
            self.file_path_input.setText(settings['file_path_input'])


def normalize(data, maxval=1., dtype=np.uint16):
    '''
    Min-max data normalization to range [0, maxval].
    Returns:
        Normalized data as numpy array of type dtype.
    '''
    data = data.astype(dtype)
    data_norm = data - data.min()
    scale_fact = maxval/data_norm.max()
    data_norm = data_norm * scale_fact
    return data_norm.astype(dtype)
