from functools import partial
from typing import TYPE_CHECKING

import numpy as np
from magicgui import magic_factory
from napari import layers
from napari.utils.notifications import show_info
from qtpy.QtCore import QItemSelectionModel, Qt, Signal
from qtpy.QtWidgets import (
    QAbstractItemView,
    QAction,
    QCheckBox,
    QComboBox,
    QFileDialog,
    QFormLayout,
    QGridLayout,
    QLabel,
    QLineEdit,
    QListView,
    QMenu,
    QMessageBox,
    QProgressDialog,
    QPushButton,
    QSpinBox,
    QTableView,
    QTabWidget,
    QVBoxLayout,
    QWidget,
)

import napari_cci_annotator._annotations_handler as _ann_handler
import napari_cci_annotator._backend_detect as backend
import napari_cci_annotator._image_handler as _image_handler
from napari_cci_annotator._settigns import Settings

if TYPE_CHECKING:
    import napari


class CciAnnotatorQWidget(QWidget):
    # your QWidget.__init__ can optionally request the napari viewer instance
    # use a type annotation of 'napari.viewer.Viewer' for any parameter

    add_labels_signal = Signal(np.ndarray, str)
    close_progress_signal = Signal()
    update_progress_signal = Signal(int, int)
    show_result_dialog_signal = Signal(str, str)

    def __init__(self, viewer: "napari.viewer.Viewer"):
        super().__init__()
        self.viewer = viewer

        self.NETWORK_IMG_SIZE = 1024

        self.setLayout(QVBoxLayout())
        
        self.settings = Settings()

        self.backend_label = QLabel("Selected/detected segmentation backend")
        self.backend_selector = QComboBox()
        self.backend_selector.addItems(backend.get_list_of_backends())
        self.cell_type_label = QLabel("Selected cell type")
        self.cell_type_selector = QComboBox()
        self.cell_type_selector.addItem("Myelin")
        self.cell_type_selector.addItem("Axons")
        self.cell_type_selector.addItem("Unmyelinated axons")
        self.layout().addWidget(self.backend_label)
        self.layout().addWidget(self.backend_selector)

        self.layout().addWidget(self.cell_type_label)
        self.layout().addWidget(self.cell_type_selector)

        self.tabWidget = QTabWidget()
        self.layout().addWidget(self.tabWidget)
        self.tabWidget.addTab(self._createFileListTab(), "Dataset")
        self.tabWidget.addTab(
            self._createLargeImageAnnotateTab(), "Auto Annotate"
        )
        self.tabWidget.addTab(self._createMetricsTab(), "Ann. Metrics")
        self.tabWidget.addTab(self._createGridTab(), "Grid Generator")
        self.tabWidget.addTab(self._createSettingsTab(), "Settings")
        self.tabWidget.currentChanged.connect(self._currentTabChanged)

        self.annotationInProgress = False
        self.imgHandler = _image_handler.ImageHandler()
        self.img_file_view.setModel(self.imgHandler.getImgModel())
        self.ann_file_view.setModel(self.imgHandler.getAnnModel())

        self.ann_handler = _ann_handler.AnnotationsHandler()
        self.annotations_view.setModel(self.ann_handler.model())
        self.annotations_view.selectionModel().currentChanged.connect(
            self._on_ann_selection_changed
        )
        self.annotations_view.setSortingEnabled(True)

        self.imgDirSet = False
        self.annDirSet = False

        # self._setup_ann_menu()
        self.ann_file_view.setContextMenuPolicy(Qt.CustomContextMenu)
        self.ann_file_view.customContextMenuRequested.connect(
            self._show_context_menu
        )

        self.add_labels_signal.connect(
            self._add_labels_slot, Qt.QueuedConnection
        )
        self.show_result_dialog_signal.connect(
            self._show_result_dialog, Qt.QueuedConnection
        )
        self.close_progress_signal.connect(
            self._close_progress_dialog, Qt.QueuedConnection
        )
        self.update_progress_signal.connect(
            self._update_progress_dialog, Qt.QueuedConnection
        )

        self.progress_dialog = QProgressDialog("", "", 0, 0)
        self.progress_dialog.setAutoClose(True)
        self.progress_dialog.reset()

        # #DEBUG STUFF!!!!
        # idir = "/home/xfolka/Projects/gisela_workflow/dl/myelin/images"
        # self.set_image_directory(idir, False)
        # adir = "/home/xfolka/Projects/gisela_workflow/dl/myelin/annotations"
        # self.set_ann_directory(adir)

    def _createMetricsTab(self):
        mainWidget = QWidget()
        mainLayout = QVBoxLayout()
        mainWidget.setLayout(mainLayout)

        self.annotations_view = QTableView()
        self.annotations_view.setEnabled(False)
        mainLayout.addWidget(self.annotations_view)
        self.annotations_view.setSelectionBehavior(
            QAbstractItemView.SelectRows
        )
        self.annotations_view.setSelectionMode(
            QAbstractItemView.SingleSelection
        )

        read_ann_btn = QPushButton("Read annotations")
        read_ann_btn.clicked.connect(self._on_read_ann_btn_click)

        mainLayout.addWidget(read_ann_btn)
        self.annotations_view.clicked.connect(self._click_annotation_id)
        self.annotations_view.activated.connect(self._click_annotation_id)

        self.del_ann_btn = QPushButton("delete annotation")
        self.del_ann_btn.clicked.connect(self._del_annotation_clicked)
        self.del_ann_btn.setEnabled(False)
        mainLayout.addWidget(self.del_ann_btn)

        self.save_morpho_btn = QPushButton("Save Morpho xls")
        self.save_morpho_btn.setEnabled(False)
        self.save_morpho_btn.setToolTip(
            "Read annotations to enable this button..."
        )
        self.save_morpho_btn.clicked.connect(self._save_morpho_xls)
        mainLayout.addWidget(self.save_morpho_btn)

        return mainWidget

    def _createGridTab(self):
        mainWidget = QWidget()
        layout = QGridLayout(mainWidget)

        layout.addWidget(
            QLabel(f"Neural Network Image Size {self.NETWORK_IMG_SIZE}"),
            0,
            0,
            1,
            2,
            Qt.AlignTop,
        )

        self.grid_size_spinbox = QSpinBox()
        self.grid_size_spinbox.setRange(100, 2000)
        self.grid_size_spinbox.setValue(
            self.NETWORK_IMG_SIZE - (2 * self.overlap_slider.value())
        )
        # self.cols_spinbox = QSpinBox()
        # self.cols_spinbox.setRange(1, 100)

        layout.addWidget(QLabel("Grid Size:"), 1, 0, Qt.AlignTop)
        layout.addWidget(self.grid_size_spinbox, 1, 1, Qt.AlignTop)
        # layout.addWidget(QLabel("Columns:"), 1, 0)
        # layout.addWidget(self.cols_spinbox, 1, 1)

        generate_button = QPushButton("Generate Grid")
        generate_button.clicked.connect(self.generate_grid)
        layout.addWidget(generate_button, 2, 0, 1, 2, Qt.AlignTop)
        layout.setVerticalSpacing(10)

        return mainWidget

    def _createLargeImageAnnotateTab(self):
        mainWidget = QWidget()
        mainLayout = QVBoxLayout()
        mainLayout.setSpacing(20)
        mainWidget.setLayout(mainLayout)

        #  large_img_btn = QPushButton("Select Image")
        #  mainLayout.addWidget(large_img_btn)

        sliderWidget = QWidget()
        sliderLayout = QFormLayout(sliderWidget)
        # self.radius_slider = QSpinBox()
        # self.radius_slider.setMinimum(1)
        # self.radius_slider.setMaximum(20)
        # radius_slider.setTickPosition(QSlider.TickBothSides)
        # sliderLayout.addRow("iso opening radius:", self.radius_slider)

        self.overlap_slider = QSpinBox()
        self.overlap_slider.setMinimum(100)
        self.overlap_slider.setMaximum(500)
        self.overlap_slider.setValue(200)
        sliderLayout.addRow("Overlap:", self.overlap_slider)

        self.crop_segment_check = QCheckBox()
        sliderLayout.addRow(
            "Crop size of segmentation to match model field of view:",
            self.crop_segment_check,
        )

        self.clear_border_check = QCheckBox()
        self.clear_border_check.setChecked(True)
        sliderLayout.addRow(
            "Remove segmentations on border:", self.clear_border_check
        )

        mainLayout.addWidget(sliderWidget)

        # self.cell_label = QLabel("Select type of cell to segment")
        # self.cell_selector = QComboBox()
        # self.cell_selector.addItem("Myelin")
        # self.cell_selector.addItem("Axons")
        # mainLayout.addWidget(self.cell_label)
        # mainLayout.addWidget(self.cell_selector)

        self.large_img_btn = QPushButton("Start Annotation")
        self.large_img_btn.clicked.connect(self._on_large_img_btn_clicked)
        # large_img_btn.setEnabled(False)
        mainLayout.addWidget(self.large_img_btn)

        # self.morpho_btn = QPushButton("Morphometrics")
        # self.morpho_btn.clicked.connect(self._on_morphometrics_clicked)
        # mainLayout.addWidget(self.morpho_btn)

        self.backend_btn = QPushButton("Check AI Backend")
        self.backend_btn.clicked.connect(self._on_backend_clicked)
        # mainLayout.addWidget(self.backend_btn)

        mainLayout.insertStretch(-1, 1)

        return mainWidget

    def _createFileListTab(self):
        mainWidget = QWidget()
        mainLayout = QGridLayout()
        mainWidget.setLayout(mainLayout)

        img_btn = QPushButton("Select image dir")
        img_btn.clicked.connect(self._on_img_btn_click)
        self.img_dir_label = QLabel("None selected")
        self.img_file_view = QListView()
        self.img_file_view.setEnabled(False)

        mainLayout.addWidget(img_btn, 0, 0)
        mainLayout.addWidget(self.img_dir_label, 1, 0)

        ann_btn = QPushButton("Select annotations dir")
        ann_btn.clicked.connect(self._on_ann_btn_click)
        self.ann_dir_label = QLabel("None selected")
        self.ann_file_view = QListView()
        self.ann_file_view.setEnabled(False)
        self.ann_file_view.setSelectionMode(QAbstractItemView.SingleSelection)

        mainLayout.addWidget(ann_btn, 0, 1)
        mainLayout.addWidget(self.ann_dir_label, 1, 1)

        mainLayout.addWidget(self.img_file_view, 2, 0)
        mainLayout.addWidget(self.ann_file_view, 2, 1)

        self.img_file_view.doubleClicked.connect(self._dbl_click_image_file)
        self.img_file_view.clicked.connect(self._click_image_file)

        self.start_ann_btn = QPushButton("Start annotation")
        self.start_ann_btn.clicked.connect(self._start_annotation_clicked)
        self.start_ann_btn.setEnabled(False)

        self.next_ann_btn = QPushButton("Next (save)")
        self.next_ann_btn.clicked.connect(self._next_annotation_clicked)
        self.next_ann_btn.setEnabled(False)

        self.done_ann_btn = QPushButton("Done (save)")
        self.done_ann_btn.clicked.connect(self._done_annotation_clicked)
        self.done_ann_btn.setEnabled(False)

        self.auto_annotate_btn = QPushButton("Auto Annotate")
        self.auto_annotate_btn.clicked.connect(self._auto_annotate_clicked)
        self.auto_annotate_btn.setEnabled(False)

        self.new_labels_check = QCheckBox()

        self.save_annotation_btn = QPushButton("Save Annotation")
        self.save_annotation_btn.clicked.connect(self._save_annotation_clicked)
        self.save_annotation_btn.setEnabled(False)

        mainLayout.addWidget(self.start_ann_btn, 3, 0)
        mainLayout.addWidget(
            QLabel(
                "Starts annotation from the selected image (or first)\nWhen you are done with that image press next\nto start annotating the next image."
            ),
            3,
            1,
        )

        mainLayout.addWidget(self.next_ann_btn, 4, 0)
        mainLayout.addWidget(
            QLabel(
                "Saves the annotated image and loads the next image and annotation."
            ),
            4,
            1,
        )

        mainLayout.addWidget(self.done_ann_btn, 5, 0)
        mainLayout.addWidget(
            QLabel("Stops the annotation session (and saves the annotation)."),
            5,
            1,
        )

        mainLayout.addWidget(self.auto_annotate_btn, 6, 0)
        mainLayout.addWidget(
            QLabel("Automatically annotates the current image.")
        )

        # mainLayout.addWidget(self.new_labels_check,7,0)
        # mainLayout.addWidget(QLabel("Create new layer when auto annotating"))

        mainLayout.addWidget(self.save_annotation_btn, 7, 0)
        mainLayout.addWidget(QLabel("Saves the current annotation."))

        return mainWidget

    def _createSettingsTab(self):
        mainWidget = QWidget()
        mainLayout = QVBoxLayout()
        mainLayout.setSpacing(20)
        mainWidget.setLayout(mainLayout)

        self.path_edit = QLineEdit(self)
        self.path_edit.setReadOnly(True)
        self.path_edit.setPlaceholderText("No model path set")

        self.clear_button = QPushButton("Clear", self)
        self.clear_button.clicked.connect(self._clear_model_path)

        # --- Layout ---
        mainLayout.addWidget(self.path_edit)
        mainLayout.addWidget(self.clear_button)
        #mainLayout.setContentsMargins(4, 4, 4, 4)
        
        self._refreshSettings()
        #self.setLayout(layout)

        #  large_img_btn = QPushButton("Select Image")
        #  mainLayout.addWidget(large_img_btn)

        return mainWidget

    def _refreshSettings(self):
        """Update the displayed path."""
        model_path = self.settings.model_path
        if model_path:
            self.path_edit.setText(str(model_path))
        else:
            self.path_edit.clear()
            self.path_edit.setPlaceholderText("No model path set")

    def _clear_model_path(self):
        """Clear the stored model path after user confirmation."""
        if QMessageBox.question(
            self,
            "Clear Model Path",
            "Are you sure you want to remove the stored model path?",
            QMessageBox.Yes | QMessageBox.No,
        ) == QMessageBox.Yes:
            self.settings.remove_model_path()
            self.path_edit.clear()
            self.path_edit.setPlaceholderText("No model path set")
            #QMessageBox.information(self, "Model Path Removed", "Model path has been cleared.")

    def _currentTabChanged(self, index):
        self.grid_size_spinbox.setValue(
            self.NETWORK_IMG_SIZE - (2 * self.overlap_slider.value())
        )

    def getRadius(self):
        return self.radius_slider.value()

    def checkEnableLists(self):
        self.img_file_view.setEnabled(self.annDirSet and self.imgDirSet)
        self.ann_file_view.setEnabled(self.annDirSet and self.imgDirSet)

        self.start_ann_btn.setEnabled(self.annDirSet and self.imgDirSet)

    def generate_grid(self):
        selected = self.viewer.layers.selection.active
        if not selected:
            return QMessageBox.information(
                None, "No layer selected", "Select a layer to generate grid on"
            )

        gridSize = self.grid_size_spinbox.value()

        # Define the size of the image
        shape = selected.data.shape
        height = shape[0]
        width = shape[1]

        # Create a transparent RGBA image
        grid = np.zeros((width, height, 4), dtype=np.uint8)  # (H, W, RGBA)

        # Define red color for grid lines (R=255, G=0, B=0, A=255)
        red_color = [255, 0, 0, 255]

        cols = height // gridSize
        rows = width // gridSize
        # Draw horizontal lines
        for row in range(rows):
            r = row * gridSize
            grid[r : r + 3, :, :] = red_color

        # Draw vertical lines
        for col in range(cols):
            c = col * gridSize
            grid[:, c : c + 3, :] = red_color

        # Add the grid as a new image layer
        self.viewer.add_image(grid, name=f"Grid {rows}x{cols}", rgb=True)

    def _show_context_menu(self, qpoint):
        globalPos = self.ann_file_view.mapToGlobal(qpoint)
        index = self.ann_file_view.indexAt(qpoint)
        nn_menu = QMenu()
        delAct = QAction("Delete")
        delAct.triggered.connect(lambda: self._delete_annotation(index))
        nn_menu.addAction(delAct)

        nn_menu.exec(globalPos)
        # nn_menu.show()

    def set_image_directory(self, img_dir, ask_for_ann_dir=True):
        self.img_dir_label.setText(img_dir)
        rIdx = self.imgHandler.setImgRootPath(img_dir)

        self.img_file_view.setRootIndex(rIdx)

        if ask_for_ann_dir and self.imgHandler.checkForMatchingAnnotationsDir(
            img_dir
        ):
            btn = QMessageBox.question(
                None,
                "Found annotations directory",
                "Annotations directory found next to images. Do you want to load the annotations from there?",
            )
            if btn == QMessageBox.Yes:
                annPath = self.imgHandler.getMatchingAnnotationsDir(img_dir)
                self.set_ann_directory(annPath)
        self.imgDirSet = True

        self.checkEnableLists()

    def _delete_annotation(self, index):
        btn = QMessageBox.warning(
            None,
            "Delete annotation?",
            "Are you sure you want to delete the annotation?",
            buttons=QMessageBox.Yes | QMessageBox.No,
        )
        if btn != QMessageBox.Yes:
            return

        self.imgHandler.delete_annotation(index, self.viewer)

        return

    def set_ann_directory(self, ann_dir):
        self.ann_dir_label.setText(ann_dir)
        rIdx = self.imgHandler.setAnnRootPath(ann_dir)
        self.ann_file_view.setRootIndex(rIdx)
        self.annDirSet = True
        self.checkEnableLists()

    def _match_ann_view_by_filename(self, fileName):
        aindx = self.imgHandler.getAnnIndexForFileName(fileName)
        self.ann_file_view.setCurrentIndex(aindx)

    def _on_img_btn_click(self):
        img_dir = QFileDialog.getExistingDirectory(
            caption="Get image directory"
        )
        self.set_image_directory(img_dir)
        show_info(f"Img dir: {img_dir} selected")

    def _add_labels_slot(self, data, name_):
        self.viewer.add_labels(data, name=name_)

    def _show_result_dialog(self, title=None, msg=None):
        QMessageBox.information(None, title, msg, buttons=QMessageBox.Ok)

    def _close_progress_dialog(self):
        self.progress_dialog.reset()

    def show_progress_dialog(self, label, cancel):
        self.progress_dialog.setLabelText(label)
        self.progress_dialog.setMinimum(0)
        self.progress_dialog.setCancelButtonText(cancel)
        self.progress_dialog.open()

    def _update_progress_dialog(self, current_value, max_value):
        self.progress_dialog.setMaximum(max_value)
        self.progress_dialog.setValue(current_value)
        self.progress_dialog.setLabelText(
            f"Processing {current_value}/{max_value}"
        )

    def _update_progress_callback(self, current_value, max_value):
        self.update_progress_signal.emit(current_value, max_value)

    def _future_done_callback(
        self,
        future,
        success_msg=None,
        error_msg=None,
        extra_func=None,
        show_exception=True,
    ):
        self.close_progress_signal.emit()

        msg = success_msg
        title = "Success"
        # The task is completed
        if (
            future.done()
            and future.exception() is not None
            and error_msg is not None
        ):
            msg = error_msg
            title = "Error during task"
            if show_exception:
                msg += ": " + str(future.exception())

        if msg is not None:
            self.show_result_dialog_signal.emit(title, msg)

        if extra_func:
            extra_func()

    def _on_large_img_btn_clicked(self):
        overlap = self.overlap_slider.value()

        be_type = self.backend_selector.currentText()

        if not be_type:
            QMessageBox.information(
                None,
                "No AI Backend  available",
                "No AI backend available for annotation.",
            )
            return

        cell_type_name = self.cell_type_selector.currentText()
        crop_image = self.crop_segment_check.isChecked()
        do_clear_border = self.clear_border_check.isChecked()
        try:
            res, future = self.imgHandler.annotate_selected_layer(
                overlap,
                self.viewer,
                be_type,
                cell_type_name,
                crop_image,
                do_clear_border,
            )
        except FileNotFoundError as e:
            QMessageBox.information(None, "Model not found", "Change the model path in settings.\n")
            return
        
        if not res or not future:
            QMessageBox.information(
                None,
                "No layer selected",
                "Wrong or no layer to segment selected!",
            )
            return
        self.large_img_btn.setEnabled(False)

        def onFinished():
            return (
                    self.large_img_btn.setEnabled(True),
                    self.add_labels_signal.emit(
                        future.result(), f"{cell_type_name.lower()}-segmentation"
                    ),
                )

        f_cb = partial(
            self._future_done_callback,
            error_msg="Segmentation error",
            extra_func=onFinished,
        )
        future.add_done_callback(f_cb)
        self.show_progress_dialog(f"Segmenting using {be_type}", "cancel")

    def _click_annotation_id(self, index):
        coord = self.ann_handler.get_coordinates_for_row(index.row())
        desired_coordinate = (0, coord[0], coord[1])
        desired_zoom_level = 2

        self.viewer.camera.center = desired_coordinate
        self.viewer.camera.zoom = desired_zoom_level

        # select the label in the contros
        l_layer = self._get_first_labels_layer_by_cell_type_if_any("myelin")
        if not l_layer:
            return

        label = self.ann_handler.get_label_for_row(index.row())
        l_layer.selected_label = label
        self.viewer.layers.selection.clear()
        self.viewer.layers.selection.add(l_layer)

        return

    def _del_annotation_clicked(self):
        labelindex = self.annotations_view.currentIndex()
        label_to_delete = self.ann_handler.get_label_for_row(labelindex.row())

        label_layer = self._get_first_labels_layer_if_any()
        if not label_layer:
            self.show_message_box(
                "No labels layer", "No labels layer available"
            )
            return False

        # labels_layer = viewer.layers['my_labels_layer']

        # Modify the layer's data to remove the specific label
        label_layer.data[label_layer.data == label_to_delete] = 0
        label_layer.data = label_layer.data
        self.ann_handler.remove_row(labelindex.row())

    def _on_ann_selection_changed(self, selected, deselected):
        self._click_annotation_id(selected)

    def _save_morpho_xls(self):
        if self.ann_handler.count() == 0:
            self.show_message_box("Empty", "No data to save")
            return

        fileName = QFileDialog.getSaveFileName(
            caption="Save File",
            directory="./morpho_report.xlsx",
            filter="*.xlsx",
        )
        future = self.ann_handler.generate_xls_report(fileName[0])

        future.add_done_callback(self._future_done_callback)
        self.show_progress_dialog("Saving morphometrics...", "cancel")

    def show_message_box(self, title, msg):
        QMessageBox.information(None, title, msg, buttons=QMessageBox.Ok)

    def _check_layers_for_morpho(self):
        if (
            self._count_label_layers_by_cell_type("myelin") != 1
            or self._count_label_layers_by_cell_type("axon") != 1
        ):
            self.show_message_box(
                "Wrong number of layers",
                "Make sure there are exactly one myelin layer and one axon layer\nDelete all unwanted label layers and try again.",
            )
            return False, None, None, None

        myelin_label_layer = self._get_first_labels_layer_by_cell_type_if_any(
            "myelin"
        )
        if not myelin_label_layer:
            self.show_message_box(
                "No myelin labels layer", "No myelin labels layer available"
            )
            return False, None, None, None

        axon_label_layer = self._get_first_labels_layer_by_cell_type_if_any(
            "axon"
        )
        if not axon_label_layer:
            self.show_message_box(
                "No axon labels layer", "No axon labels layer available"
            )
            return False, None, None, None

        image_layer = self._get_first_image_layer_if_any()
        if not image_layer:
            self.show_message_box("No image layer", "No image layer available")
            return False, None, None, None

        return True, myelin_label_layer, axon_label_layer, image_layer

    def _on_read_ann_btn_click(self):
        res, myelin_label_layer, axon_label_layer, image_layer = (
            self._check_layers_for_morpho()
        )
        if not res:
            return False

        future = self.ann_handler.add_annotations_to_model(
            myelin_label_layer.data,
            axon_label_layer.data,
            image_layer.data,
            self._update_progress_callback,
        )

        def onFinished():
            return (
                    self.annotations_view.setEnabled(self.ann_handler.count() > 0),
                    self.del_ann_btn.setEnabled(True),
                    self.save_morpho_btn.setEnabled(True),
                )

        f_cb = partial(self._future_done_callback, extra_func=onFinished)
        future.add_done_callback(f_cb)
        self.show_progress_dialog("Adding annotations...", "cancel")

    def _on_morphometrics_clicked(self):
        res, myelin_label_layer, axon_label_layer, image_layer = (
            self._check_layers_for_morpho()
        )
        if not res:
            return False

        xls_dir = QFileDialog.getExistingDirectory(
            caption="Select morpho table directory"
        )
        if xls_dir == "":
            return
        show_info(f"Save excel file in dir: {xls_dir}")
        fname = xls_dir + "/" + image_layer.name + "_morpho.xlsx"

        # start progress diablog
        f_cb = partial(
            self._future_done_callback,
            success_msg=f"Morphometrics saved in: {fname}",
            error_msg="Exception during morphometrics",
        )
        future = self.imgHandler.calulate_morphometrics(
            fname,
            myelin_label_layer.data,
            axon_label_layer.data,
            image_layer.data,
        )
        future.add_done_callback(f_cb)
        self.show_progress_dialog("working morphometrics...", "Cancel")

    def _on_backend_clicked(self):
        import napari_cci_annotator._backend_detect as backend

        _, be_type = backend.select_backend()
        QMessageBox.information(
            None, "Detected backend", f"Detected backend is: {be_type}"
        )
        # else:
        #     QMessageBox.information(None,"GPU available","GPU available for annotation!")

    def _on_ann_btn_click(self):
        ann_dir = QFileDialog.getExistingDirectory(
            caption="Get annotation directory"
        )
        if ann_dir == "":
            return
        self.set_ann_directory(ann_dir)
        show_info(f"Annotations dir: {ann_dir} selected")

    def _dbl_click_image_file(self, index):
        self._remove_all_layers()
        # data = self.imgHandler.getImgData(index)
        self.imgHandler.loadImageAndAnnotation(index, self.viewer)
        self.save_annotation_btn.setEnabled(True)
        self.auto_annotate_btn.setEnabled(True)

    def _select_image_and_corresponding_annotation(self, index):
        data = self.imgHandler.getImgData(index)
        self._match_ann_view_by_filename(data)
        self.img_file_view.setCurrentIndex(index)

    def _get_first_labels_layer_if_any(self):
        for layer in self.viewer.layers:
            if isinstance(layer, layers.Labels):
                return layer
        return None

    def _get_first_labels_layer_by_cell_type_if_any(self, cell_type: str):
        for layer in self.viewer.layers:
            if isinstance(
                layer, layers.Labels
            ) and layer.name.lower().startswith(cell_type.lower()):
                return layer
        return None

    def _get_first_image_layer_if_any(self):
        for layer in self.viewer.layers:
            if isinstance(layer, layers.Image):
                return layer
        return None

    def _auto_annotate_clicked(self):
        be, be_type = backend.select_backend()
        if not be:
            QMessageBox.information(
                None,
                "No AI Backend  available",
                "No AI backend available for annotation.",
            )
            return

        cell_type = self.cell_type_selector.currentText()

        self.imgHandler.autoAnnotateImage(
            self.img_file_view.currentIndex(), self.viewer, cell_type, be_type
        )

    def _count_label_layers(self):
        cnt = 0
        for layer in self.viewer.layers:
            if isinstance(layer, layers.Labels):
                cnt += 1
        return cnt

    def _count_label_layers_by_cell_type(self, cell_type: str):
        cnt = 0
        for layer in self.viewer.layers:
            if isinstance(
                layer, layers.Labels
            ) and layer.name.lower().startswith(cell_type.lower()):
                cnt += 1
        return cnt

    def _count_image_layers(self):
        cnt = 0
        for layer in self.viewer.layers:
            if isinstance(layer, layers.Image):
                cnt += 1
        return cnt

    def _save_only_one_layer(self):
        if self._count_label_layers() > 1:
            QMessageBox.information(
                None,
                "More than one labels layer",
                "More than one labels layer exist.\nDelete all unwanted label layers and save again.",
                buttons=QMessageBox.Ok,
            )
            return False

        labels = self._get_first_labels_layer_if_any()
        imgName = self.imgHandler.getImgData(self.img_file_view.currentIndex())
        self.imgHandler.saveAnnotationUsingName(
            labels.name, imgName, self.viewer
        )
        return True

    def _save_annotation_clicked(self):
        self._save_only_one_layer()

    def _click_image_file(self, index):
        self._select_image_and_corresponding_annotation(index)
        data = self.imgHandler.getImgData(index)
        self.move_layer_to_top_if_it_exists(data)

    def _remove_all_layers(self):
        self.viewer.layers.clear()

    def _remove_images_from_layers(self, img_layer_name):
        imgLayer = None
        annLayer = None
        if img_layer_name in self.viewer.layers:
            imgLayer = self.viewer.layers[img_layer_name]
            self.viewer.layers.remove(imgLayer)
        else:
            return (imgLayer, annLayer)

        ann_layer_name = self.imgHandler.getAnnotationLayerNameFromImageName(
            img_layer_name
        )
        if ann_layer_name in self.viewer.layers:
            annLayer = self.viewer.layers[ann_layer_name]
            self.viewer.layers.remove(annLayer)

        return (imgLayer, annLayer)

    def move_layer_to_top_if_it_exists(self, img_layer_name):
        (imgLayer, annLayer) = self._remove_images_from_layers(img_layer_name)

        if imgLayer:
            self.viewer.layers.append(imgLayer)

        if annLayer:
            self.viewer.layers.append(annLayer)

    def _select_first_items(self):
        ridx = self.img_file_view.rootIndex()
        firstIndex = self.imgHandler.getImgModel().index(0, 0, ridx)
        if firstIndex.isValid():
            self.img_file_view.setCurrentIndex(firstIndex)
            self.img_file_view.selectionModel().select(
                firstIndex, QItemSelectionModel.Select
            )

    def _start_annotation_clicked(self):
        # show dialog warning for clearing layers before starting. User can cancel
        if len(self.viewer.layers) > 0:
            btn = QMessageBox.warning(
                None,
                "Erase progress",
                "Clicking OK and starting annotations will delete all current layers and destroy progress",
                buttons=QMessageBox.Ok | QMessageBox.Cancel,
            )
            if btn != QMessageBox.Ok:
                return

        self.start_ann_btn.setEnabled(False)
        self.next_ann_btn.setEnabled(True)
        self.done_ann_btn.setEnabled(True)
        self.img_file_view.setEnabled(False)
        self.ann_file_view.setEnabled(False)
        self.save_annotation_btn.setEnabled(True)
        self.auto_annotate_btn.setEnabled(True)

        if not self.img_file_view.currentIndex().isValid():
            self._select_first_items()

        self._remove_all_layers()
        self._select_image_and_corresponding_annotation(
            self.img_file_view.currentIndex()
        )
        self.imgHandler.loadImageAndAnnotation(
            self.img_file_view.currentIndex(), self.viewer
        )
        return

    def _next_annotation_clicked(self):
        # self.imgHandler.saveAnnotation(self.img_file_view.currentIndex(),self.viewer)
        if not self._save_only_one_layer():
            return

        self._remove_all_layers()

        nextIdx = self.imgHandler.nextImageIndex(
            self.img_file_view.currentIndex()
        )
        self._select_image_and_corresponding_annotation(nextIdx)
        self.imgHandler.loadImageAndAnnotation(
            self.img_file_view.currentIndex(), self.viewer
        )

        return

    def _done_annotation_clicked(self):
        if not self._save_only_one_layer():
            return

        self.start_ann_btn.setEnabled(True)
        self.next_ann_btn.setEnabled(False)
        self.done_ann_btn.setEnabled(False)
        self.img_file_view.setEnabled(True)
        self.ann_file_view.setEnabled(True)

        return
