object-segmentation/src/gui/tabs/annotation_tab.py

"""
Annotation tab for the microscopy object detection application.
Manual annotation with pen tool and object class management.
"""

from PySide6.QtWidgets import (
    QWidget,
    QVBoxLayout,
    QHBoxLayout,
    QLabel,
    QGroupBox,
    QPushButton,
    QFileDialog,
    QMessageBox,
    QSplitter,
)
from PySide6.QtCore import Qt, QSettings
from pathlib import Path

from src.database.db_manager import DatabaseManager
from src.utils.config_manager import ConfigManager
from src.utils.image import Image, ImageLoadError
from src.utils.logger import get_logger
from src.gui.widgets import AnnotationCanvasWidget, AnnotationToolsWidget

logger = get_logger(__name__)


class AnnotationTab(QWidget):
    """Annotation tab for manual image annotation."""

    def __init__(
        self, db_manager: DatabaseManager, config_manager: ConfigManager, parent=None
    ):
        super().__init__(parent)
        self.db_manager = db_manager
        self.config_manager = config_manager
        self.current_image = None
        self.current_image_path = None
        self.current_image_id = None
        self.current_annotations = []

        self._setup_ui()

    def _setup_ui(self):
        """Setup user interface."""
        layout = QVBoxLayout()

        # Main horizontal splitter to divide left (image) and right (controls)
        self.main_splitter = QSplitter(Qt.Horizontal)
        self.main_splitter.setHandleWidth(10)

        # { Left splitter for image display and zoom info
        self.left_splitter = QSplitter(Qt.Vertical)
        self.left_splitter.setHandleWidth(10)

        # Annotation canvas section
        canvas_group = QGroupBox("Annotation Canvas")
        canvas_layout = QVBoxLayout()

        # Use the AnnotationCanvasWidget
        self.annotation_canvas = AnnotationCanvasWidget()
        self.annotation_canvas.zoom_changed.connect(self._on_zoom_changed)
        self.annotation_canvas.annotation_drawn.connect(self._on_annotation_drawn)
        canvas_layout.addWidget(self.annotation_canvas)

        canvas_group.setLayout(canvas_layout)
        self.left_splitter.addWidget(canvas_group)

        # Controls info
        controls_info = QLabel(
            "Zoom: Mouse wheel or +/- keys | Drawing: Enable pen and drag mouse"
        )
        controls_info.setStyleSheet("QLabel { color: #888; font-style: italic; }")
        self.left_splitter.addWidget(controls_info)
        # }

        # { Right splitter for annotation tools and controls
        self.right_splitter = QSplitter(Qt.Vertical)
        self.right_splitter.setHandleWidth(10)

        # Annotation tools section
        self.annotation_tools = AnnotationToolsWidget(self.db_manager)
        self.annotation_tools.polyline_enabled_changed.connect(
            self.annotation_canvas.set_polyline_enabled
        )
        self.annotation_tools.polyline_pen_color_changed.connect(
            self.annotation_canvas.set_polyline_pen_color
        )
        self.annotation_tools.polyline_pen_width_changed.connect(
            self.annotation_canvas.set_polyline_pen_width
        )
        # Show / hide bounding boxes
        self.annotation_tools.show_bboxes_changed.connect(
            self.annotation_canvas.set_show_bboxes
        )
        # RDP simplification controls
        self.annotation_tools.simplify_on_finish_changed.connect(
            self._on_simplify_on_finish_changed
        )
        self.annotation_tools.simplify_epsilon_changed.connect(
            self._on_simplify_epsilon_changed
        )
        # Class selection and class-color changes
        self.annotation_tools.class_selected.connect(self._on_class_selected)
        self.annotation_tools.class_color_changed.connect(self._on_class_color_changed)
        self.annotation_tools.clear_annotations_requested.connect(
            self._on_clear_annotations
        )
        self.right_splitter.addWidget(self.annotation_tools)

        # Image loading section
        load_group = QGroupBox("Image Loading")
        load_layout = QVBoxLayout()

        # Load image button
        button_layout = QHBoxLayout()
        self.load_image_btn = QPushButton("Load Image")
        self.load_image_btn.clicked.connect(self._load_image)
        button_layout.addWidget(self.load_image_btn)
        button_layout.addStretch()
        load_layout.addLayout(button_layout)

        # Image info label
        self.image_info_label = QLabel("No image loaded")
        load_layout.addWidget(self.image_info_label)

        load_group.setLayout(load_layout)
        self.right_splitter.addWidget(load_group)
        # }

        # Add both splitters to the main horizontal splitter
        self.main_splitter.addWidget(self.left_splitter)
        self.main_splitter.addWidget(self.right_splitter)

        # Set initial sizes: 75% for left (image), 25% for right (controls)
        self.main_splitter.setSizes([750, 250])

        layout.addWidget(self.main_splitter)
        self.setLayout(layout)

        # Restore splitter positions from settings
        self._restore_state()

    def _load_image(self):
        """Load and display an image file."""
        # Get last opened directory from QSettings
        settings = QSettings("microscopy_app", "object_detection")
        last_dir = settings.value("annotation_tab/last_directory", None)

        # Fallback to image repository path or home directory
        if last_dir and Path(last_dir).exists():
            start_dir = last_dir
        else:
            repo_path = self.config_manager.get_image_repository_path()
            start_dir = repo_path if repo_path else str(Path.home())

        # Open file dialog
        file_path, _ = QFileDialog.getOpenFileName(
            self,
            "Select Image",
            start_dir,
            "Images (*.jpg *.jpeg *.png *.tif *.tiff *.bmp)",
        )

        if not file_path:
            return

        try:
            # Load image using Image class
            self.current_image = Image(file_path)
            self.current_image_path = file_path

            # Store the directory for next time
            settings.setValue(
                "annotation_tab/last_directory", str(Path(file_path).parent)
            )

            # Get or create image in database
            relative_path = str(Path(file_path).name)  # Simplified for now
            self.current_image_id = self.db_manager.get_or_create_image(
                relative_path,
                Path(file_path).name,
                self.current_image.width,
                self.current_image.height,
            )

            # Display image using the AnnotationCanvasWidget
            self.annotation_canvas.load_image(self.current_image)

            # Load and display any existing annotations for this image
            self._load_annotations_for_current_image()

            # Update info label
            self._update_image_info()

            logger.info(f"Loaded image: {file_path} (DB ID: {self.current_image_id})")

        except ImageLoadError as e:
            logger.error(f"Failed to load image: {e}")
            QMessageBox.critical(
                self, "Error Loading Image", f"Failed to load image:\n{str(e)}"
            )
        except Exception as e:
            logger.error(f"Unexpected error loading image: {e}")
            QMessageBox.critical(self, "Error", f"Unexpected error:\n{str(e)}")

    def _update_image_info(self):
        """Update the image info label with current image details."""
        if self.current_image is None:
            self.image_info_label.setText("No image loaded")
            return

        zoom_percentage = self.annotation_canvas.get_zoom_percentage()
        info_text = (
            f"File: {Path(self.current_image_path).name}\n"
            f"Size: {self.current_image.width}x{self.current_image.height} pixels\n"
            f"Channels: {self.current_image.channels}\n"
            f"Data type: {self.current_image.dtype}\n"
            f"Format: {self.current_image.format.upper()}\n"
            f"File size: {self.current_image.size_mb:.2f} MB\n"
            f"Zoom: {zoom_percentage}%"
        )
        self.image_info_label.setText(info_text)

    def _on_zoom_changed(self, zoom_scale: float):
        """Handle zoom level changes from the annotation canvas."""
        self._update_image_info()

    def _on_annotation_drawn(self, points: list):
        """
        Handle when an annotation stroke is drawn.

        Saves the new annotation directly to the database and refreshes the
        on-canvas display of annotations for the current image.
        """
        # Ensure we have an image loaded and in the DB
        if not self.current_image or not self.current_image_id:
            logger.warning("Annotation drawn but no image loaded")
            QMessageBox.warning(
                self,
                "No Image",
                "Please load an image before drawing annotations.",
            )
            return

        current_class = self.annotation_tools.get_current_class()

        if not current_class:
            logger.warning("Annotation drawn but no object class selected")
            QMessageBox.warning(
                self,
                "No Class Selected",
                "Please select an object class before drawing annotations.",
            )
            return

        if not points:
            logger.warning("Annotation drawn with no points, ignoring")
            return

        # points are [(x_norm, y_norm), ...]
        xs = [p[0] for p in points]
        ys = [p[1] for p in points]
        x_min, x_max = min(xs), max(xs)
        y_min, y_max = min(ys), max(ys)

        # Store segmentation mask in [y_norm, x_norm] format to match DB
        db_polyline = [[float(y), float(x)] for (x, y) in points]

        try:
            annotation_id = self.db_manager.add_annotation(
                image_id=self.current_image_id,
                class_id=current_class["id"],
                bbox=(x_min, y_min, x_max, y_max),
                annotator="manual",
                segmentation_mask=db_polyline,
                verified=False,
            )

            logger.info(
                f"Saved annotation (ID: {annotation_id}) for class "
                f"'{current_class['class_name']}' "
                f"Bounding box: ({x_min:.3f}, {y_min:.3f}) to ({x_max:.3f}, {y_max:.3f})\n"
                f"with {len(points)} polyline points"
            )

            # Reload annotations from DB and redraw (respecting current class filter)
            self._load_annotations_for_current_image()

        except Exception as e:
            logger.error(f"Failed to save annotation: {e}")
            QMessageBox.critical(self, "Error", f"Failed to save annotation:\n{str(e)}")

    def _on_simplify_on_finish_changed(self, enabled: bool):
        """Update canvas simplify-on-finish flag from tools widget."""
        self.annotation_canvas.simplify_on_finish = enabled
        logger.debug(f"Annotation simplification on finish set to {enabled}")

    def _on_simplify_epsilon_changed(self, epsilon: float):
        """Update canvas RDP epsilon from tools widget."""
        self.annotation_canvas.simplify_epsilon = float(epsilon)
        logger.debug(f"Annotation simplification epsilon set to {epsilon}")

    def _on_class_color_changed(self):
        """
        Handle changes to the selected object's class color.

        When the user updates a class color in the tools widget, reload the
        annotations for the current image so that all polylines are redrawn
        using the updated per-class colors.
        """
        if not self.current_image_id:
            return

        logger.debug(
            f"Class color changed; reloading annotations for image ID {self.current_image_id}"
        )
        self._load_annotations_for_current_image()

    def _on_class_selected(self, class_data):
        """
        Handle when an object class is selected or cleared.

        When a specific class is selected, only annotations of that class are drawn.
        When the selection is cleared (\"-- Select Class --\"), all annotations are shown.
        """
        if class_data:
            logger.debug(f"Object class selected: {class_data['class_name']}")
        else:
            logger.debug(
                'No class selected ("-- Select Class --"), showing all annotations'
            )

        # Whenever the selection changes, update which annotations are visible
        self._redraw_annotations_for_current_filter()

    def _on_clear_annotations(self):
        """Handle clearing all annotations."""
        self.annotation_canvas.clear_annotations()
        logger.info("Cleared all annotations")

    def _load_annotations_for_current_image(self):
        """
        Load all annotations for the current image from the database and
        redraw them on the canvas, honoring the currently selected class
        filter (if any).
        """
        if not self.current_image_id:
            self.current_annotations = []
            self.annotation_canvas.clear_annotations()
            return

        try:
            self.current_annotations = self.db_manager.get_annotations_for_image(
                self.current_image_id
            )
            self._redraw_annotations_for_current_filter()
        except Exception as e:
            logger.error(
                f"Failed to load annotations for image {self.current_image_id}: {e}"
            )
            QMessageBox.critical(
                self,
                "Error",
                f"Failed to load annotations for this image:\n{str(e)}",
            )

    def _redraw_annotations_for_current_filter(self):
        """
        Redraw annotations for the current image, optionally filtered by the
        currently selected object class.
        """
        # Clear current on-canvas annotations but keep the image
        self.annotation_canvas.clear_annotations()

        if not self.current_annotations:
            return

        current_class = self.annotation_tools.get_current_class()
        selected_class_id = current_class["id"] if current_class else None

        drawn_count = 0
        for ann in self.current_annotations:
            # Filter by class if one is selected
            if (
                selected_class_id is not None
                and ann.get("class_id") != selected_class_id
            ):
                continue

            if ann.get("segmentation_mask"):
                polyline = ann["segmentation_mask"]
                color = ann.get("class_color", "#FF0000")

                self.annotation_canvas.draw_saved_polyline(polyline, color, width=3)
                self.annotation_canvas.draw_saved_bbox(
                    [ann["x_min"], ann["y_min"], ann["x_max"], ann["y_max"]],
                    color,
                    width=3,
                )
                drawn_count += 1

        logger.info(
            f"Displayed {drawn_count} annotation(s) for current image with "
            f"{'no class filter' if selected_class_id is None else f'class_id={selected_class_id}'}"
        )

    def _restore_state(self):
        """Restore splitter positions from settings."""
        settings = QSettings("microscopy_app", "object_detection")

        # Restore main splitter state
        main_state = settings.value("annotation_tab/main_splitter_state")
        if main_state:
            self.main_splitter.restoreState(main_state)
            logger.debug("Restored main splitter state")

        # Restore left splitter state
        left_state = settings.value("annotation_tab/left_splitter_state")
        if left_state:
            self.left_splitter.restoreState(left_state)
            logger.debug("Restored left splitter state")

        # Restore right splitter state
        right_state = settings.value("annotation_tab/right_splitter_state")
        if right_state:
            self.right_splitter.restoreState(right_state)
            logger.debug("Restored right splitter state")

    def save_state(self):
        """Save splitter positions to settings."""
        settings = QSettings("microscopy_app", "object_detection")

        # Save main splitter state
        settings.setValue(
            "annotation_tab/main_splitter_state", self.main_splitter.saveState()
        )

        # Save left splitter state
        settings.setValue(
            "annotation_tab/left_splitter_state", self.left_splitter.saveState()
        )

        # Save right splitter state
        settings.setValue(
            "annotation_tab/right_splitter_state", self.right_splitter.saveState()
        )

        logger.debug("Saved annotation tab splitter states")

    def refresh(self):
        """Refresh the tab."""
        pass