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._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.pen_enabled_changed.connect(
            self.annotation_canvas.set_pen_enabled
        )
        self.annotation_tools.pen_color_changed.connect(
            self.annotation_canvas.set_pen_color
        )
        self.annotation_tools.pen_width_changed.connect(
            self.annotation_canvas.set_pen_width
        )
        self.annotation_tools.class_selected.connect(self._on_class_selected)
        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)

            # 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."""
        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

        logger.info(
            f"Annotation drawn with {len(points)} points for class: {current_class['class_name']}"
        )
        # Future: Save annotation to database or export

    def _on_class_selected(self, class_data: dict):
        """Handle when an object class is selected."""
        logger.debug(f"Object class selected: {class_data['class_name']}")

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

    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