import os
import sys
from PyQt5.QtCore import Qt, QRectF, QPointF, QTimer, QPoint, QSizeF, QEvent
from PyQt5.QtGui import QPixmap, QImage, QPen, QColor, QCursor, QPainter, QTransform, QPolygonF, QMouseEvent, QPainterPath
from PyQt5.QtWidgets import (
    QApplication, QMainWindow, QWidget, QVBoxLayout, QHBoxLayout, QLabel, QPushButton,
    QCheckBox, QLineEdit, QListWidget, QFileDialog, QGraphicsView, QGraphicsScene, QGraphicsRectItem,
    QGraphicsLineItem, QInputDialog, QMessageBox, QGroupBox, QGraphicsEllipseItem, QComboBox, QGraphicsPolygonItem,
    QGraphicsItem, QDesktopWidget, QGraphicsPathItem
)

BASE_DIR = os.path.dirname(os.path.abspath(__file__))
IMAGES_DIR = os.path.join(BASE_DIR, "images")


class AnnotationTool(QMainWindow):
    def __init__(self):
        super().__init__()
        self.setWindowFlags(Qt.Window | Qt.WindowMinimizeButtonHint | Qt.CustomizeWindowHint)
        self.current_bbox_count = 0
        self.setWindowTitle('Annotation Tool')
        screen = QDesktopWidget().screenGeometry()
        self.setFixedSize(screen.width(), screen.height())
        self.showMaximized()
        self.setFixedSize(self.size())

        self.max_bboxes = 100
        self.current_shape = "Rectangle"
        self.class_to_id = {}  # Map label to class ID
        self.id_to_class = {}  # Map class ID to label
        self.resizing = False
        self.resize_handle = None
        self.resize_start_pos = None

        self.vertex_handles = []
        self.preview_path = QGraphicsPathItem()
        self.preview_path.setPen(QPen(QColor(255, 100, 100), 2, Qt.DashLine))
        self.preview_path.setZValue(100)
        self.preview_path.setPath(QPainterPath())

        # Polygon specific
        self.is_drawing_polygon = False
        self.polygon_points = []
        self.first_polygon_point = None
        self.close_polygon_threshold = 10

        self.setStyleSheet("""
            QMainWindow {
                background-color: #FFFFFF;
                color: #000000;
            }
            QPushButton {
                background-color: #000000;
                color: #FFFFFF;
                border: 2px solid #FFFFFF;
                padding: 10px;
                margin: 5px;
                border-radius: 5px;
                font: 25px 'Arial';
            }
            QPushButton:hover {
                background-color: #CCCCCC;
                color: #000000;
            }
            QPushButton:pressed {
                background-color: #AAAAAA;
                color: #000000;
            }
            QCheckBox {
                color: #000000;
            }
            QLineEdit {
                background-color: #000000;  
                color: #FFFFFF;
                border: 2px solid #FFFFFF;
                padding: 5px;
                margin: 5px;
                border-radius: 5px;
            }
            QLabel {
                color: #000000;
                margin: 5px;
            }
            QListWidget {
                background-color: #FFFFFF;
                color: #000000;
                border: 2px solid #000000;
                margin: 5px;
                border-radius: 5px;
            }
            QGraphicsView {
                border: 2px solid #000000;
            }
            QComboBox {
                background-color: #000000;
                color: #FFFFFF;
                border: 2px solid #FFFFFF;
                padding: 5px;
                margin: 5px;
                border-radius: 5px;
                font: 25px 'Arial';
            }
            QComboBox:hover {
                background-color: #CCCCCC;
                color: #000000;
            }
            QComboBox::drop-down {
                border: none;
                width: 30px;
            }
            QComboBox::down-arrow {
                image: none;
            }
        """)

        main_layout = QHBoxLayout()

        left_panel = QVBoxLayout()
        self.back_button = QPushButton('Back')
        self.back_button.setToolTip("Return to Data Annotation Options")
        left_panel.addWidget(self.back_button)
        self.folder_selector_button = QPushButton('Select Folder')
        left_panel.addWidget(self.folder_selector_button)
        self.save_directory_button = QPushButton('Save Dir Location')
        left_panel.addWidget(self.save_directory_button)
        self.shape_selector = QComboBox()
        self.shape_selector.addItems(["Rectangle", "Square", "Circle", "Polygon"])
        self.shape_selector.setCurrentText("Rectangle")
        left_panel.addWidget(self.shape_selector)

        self.bbox_widget = QWidget()
        self.bbox_widget.setFixedHeight(180)
        self.bbox_widget.setStyleSheet("""
            QWidget {
                background-color: #708090;
                border: 2px solid #FFFFFF;
                border-radius: 5px;
                font-size: 30;
            }
            QLabel {
                color: #FFFFFF;
                margin: 2px;
                font-size: 20px;
                wordwrap: true;
            }
        """)
        bbox_layout = QVBoxLayout(self.bbox_widget)
        self.bbox_count_label = QLabel("Number of bounding boxes:")
        self.bbox_count_label.setFixedHeight(30)
        self.bbox_count_label.setAlignment(Qt.AlignCenter)
        self.bbox_count_label.setStyleSheet("color: #000000")
        self.bbox_count_label.setWordWrap(True)
        bbox_layout.addWidget(self.bbox_count_label)
        bbox_input_layout = QHBoxLayout()
        self.bbox_input = QLineEdit()
        self.bbox_input.setFixedHeight(40)
        self.bbox_input.setFixedWidth(153)
        self.bbox_input.setPlaceholderText("Number of bounding boxes")
        self.bbox_input.setStyleSheet("color: #000000; font-size: 10px")
        bbox_input_layout.addWidget(self.bbox_input)
        self.bbox_submit_button = QPushButton("Submit")
        self.bbox_submit_button.setFixedHeight(90)
        self.bbox_submit_button.setFixedWidth(107)
        self.bbox_submit_button.setStyleSheet("""
            QPushButton {
                color: #FFFFFF;
                font-size: 22px;
                background-color: #000000;
                border: 2px solid #FFFFFF;
                padding: 10px;
                margin: 5px;
                border-radius: 5px;
            }
            QPushButton:hover {
                color: #00FF00;
            }
            QPushButton:pressed {
                color: #0000FF;
            }
            QPushButton:released {
                color: #FFFFFF;
            }
        """)
        bbox_input_layout.addWidget(self.bbox_submit_button)
        bbox_layout.addLayout(bbox_input_layout)
        left_panel.addWidget(self.bbox_widget)

        self.previous_image_button = QPushButton('Previous Image')
        left_panel.addWidget(self.previous_image_button)
        self.next_image_button = QPushButton('Next Image')
        left_panel.addWidget(self.next_image_button)
        self.delete_bounding_box_button = QPushButton('Delete Bounding Box')
        left_panel.addWidget(self.delete_bounding_box_button)
        self.zoom_in_button = QPushButton('Zoom In')
        left_panel.addWidget(self.zoom_in_button)
        self.zoom_out_button = QPushButton('Zoom Out')
        left_panel.addWidget(self.zoom_out_button)

        left_panel_widget = QWidget()
        left_panel_widget.setLayout(left_panel)
        left_panel_widget.setFixedWidth(300)

        self.image_viewer = QGraphicsView()
        self.image_scene = QGraphicsScene()
        self.image_viewer.setScene(self.image_scene)
        self.image_viewer.setRenderHint(QPainter.Antialiasing)
        self.image_viewer.setRenderHint(QPainter.SmoothPixmapTransform)
        self.image_viewer.setDragMode(QGraphicsView.NoDrag)
        self.image_viewer.setHorizontalScrollBarPolicy(Qt.ScrollBarAlwaysOff)
        self.image_viewer.setVerticalScrollBarPolicy(Qt.ScrollBarAlwaysOff)
        self.image_viewer.setContextMenuPolicy(Qt.NoContextMenu)

        right_panel = QVBoxLayout()
        default_label_group = QGroupBox("Default Label")
        default_label_layout = QVBoxLayout()
        self.default_label_checkbox = QCheckBox('Use Default Label')
        default_label_layout.addWidget(self.default_label_checkbox)
        self.default_label_textbox = QLineEdit()
        default_label_layout.addWidget(self.default_label_textbox)
        default_label_group.setLayout(default_label_layout)
        right_panel.addWidget(default_label_group)

        label_reflector_group = QGroupBox("Label Reflector")
        label_reflector_layout = QVBoxLayout()
        self.label_reflector = QListWidget()
        label_reflector_layout.addWidget(self.label_reflector)
        label_reflector_group.setLayout(label_reflector_layout)
        right_panel.addWidget(label_reflector_group)

        file_list_group = QGroupBox("File List")
        file_list_layout = QHBoxLayout()
        self.image_position_label = QLabel("0/0")
        self.image_position_label.setStyleSheet("""
            QLabel {
                color: #000000;
                font-size: 16px;
                margin-left: 10px;
            }
        """)
        file_list_layout.addWidget(self.image_position_label, alignment=Qt.AlignRight)
        file_list_inner_layout = QVBoxLayout()
        self.file_list = QListWidget()
        file_list_inner_layout.addWidget(self.file_list)
        file_list_layout.addLayout(file_list_inner_layout)
        file_list_group.setLayout(file_list_layout)
        right_panel.addWidget(file_list_group)

        right_panel_widget = QWidget()
        right_panel_widget.setLayout(right_panel)
        right_panel_widget.setFixedWidth(300)

        main_layout.addWidget(left_panel_widget)
        main_layout.addWidget(self.image_viewer)
        main_layout.addWidget(right_panel_widget)

        central_widget = QWidget()
        central_widget.setLayout(main_layout)
        self.setCentralWidget(central_widget)

        self.setup_signals()

        self.image_files = []
        self.current_image_index = -1
        self.bbox = None
        self.start_point = None
        self.end_point = None
        self.drawing = False
        self.selected_bbox = None
        self.bounding_boxes = {}  # {image_path: [(coords, label, shape_type)]}
        self.label_items = {}  # {image_path: {bbox_item: label}}
        self.save_directory = None
        self.last_bbox_per_image = {}

        self.vertical_line = QGraphicsLineItem()
        self.vertical_line.setPen(QPen(QColor(0, 0, 255), 2, Qt.DashLine))
        self.image_scene.addItem(self.vertical_line)

        self.horizontal_line = QGraphicsLineItem()
        self.horizontal_line.setPen(QPen(QColor(0, 0, 255), 2, Qt.DashLine))
        self.image_scene.addItem(self.horizontal_line)

        self.image_scene.addItem(self.preview_path)

        self.image_viewer.setMouseTracking(True)
        self.image_viewer.viewport().installEventFilter(self)
        self.setFocusPolicy(Qt.StrongFocus)
        self.setFocus()

    def distance_viewport(self, scene_pos1, scene_pos2):
        """Calculate distance between two scene points in viewport (screen) pixels"""
        v1 = self.image_viewer.mapFromScene(scene_pos1)
        v2 = self.image_viewer.mapFromScene(scene_pos2)
        return ((v1.x() - v2.x()) ** 2 + (v1.y() - v2.y()) ** 2) ** 0.5

    def is_near_first_point(self, scene_pos):
        if not self.is_drawing_polygon or not self.first_polygon_point:
            return False

        dx = scene_pos.x() - self.first_polygon_point.x()
        dy = scene_pos.y() - self.first_polygon_point.y()
        dist = (dx * dx + dy * dy) ** 0.5

        # threshold shrinks when zoomed in
        threshold = 12 / self.image_viewer.transform().m11()
        return dist <= threshold

    def is_click_on_first_handle(self, scene_pos):
        if not self.vertex_handles:
            return False
        first = self.vertex_handles[0]
        p = first.mapFromScene(scene_pos)
        return first.rect().contains(p)

    def setup_signals(self):
        try:
            self.back_button.clicked.connect(self.open_data_annotation_window)
            self.folder_selector_button.clicked.connect(self.select_folder)
            self.save_directory_button.clicked.connect(self.select_save_directory)
            self.shape_selector.currentTextChanged.connect(self.update_shape)
            self.bbox_submit_button.clicked.connect(self.set_bbox_limit)
            self.previous_image_button.clicked.connect(self.previous_image)
            self.next_image_button.clicked.connect(self.next_image)
            self.delete_bounding_box_button.clicked.connect(self.delete_selected_bounding_box)
            self.zoom_in_button.clicked.connect(self.zoom_in)
            self.zoom_out_button.clicked.connect(self.zoom_out)
            self.label_reflector.itemSelectionChanged.connect(self.on_label_reflector_selection_changed)
            self.file_list.itemSelectionChanged.connect(self.on_file_selected)
        except Exception as e:
            QMessageBox.critical(self, "Error", f"Failed to connect signals: {str(e)}")

    def open_data_annotation_window(self):
        try:
            from .fun_pyqt5 import DataAnnotationWindow
            from PyQt5.QtGui import QPixmap
            bg_image_path = os.path.join(IMAGES_DIR, 'bg_help_image.png')
            if not os.path.exists(bg_image_path):
                raise FileNotFoundError(f"Background image {bg_image_path} not found")
            self.data_annotation_window = DataAnnotationWindow(QPixmap(bg_image_path))
            self.data_annotation_window.showMaximized()
            self.hide()
        except FileNotFoundError as e:
            QMessageBox.critical(self, "Error", str(e))
        except Exception as e:
            QMessageBox.critical(self, "Error", f"Failed to open Data Annotation Options: {str(e)}")

    def closeEvent(self, event):
        try:
            from .fun_pyqt5 import DataAnnotationWindow
            self.data_annotation_window = DataAnnotationWindow()
            self.data_annotation_window.showMaximized()
            event.accept()
        except Exception as e:
            QMessageBox.critical(self, "Error", f"Failed to return to Data Annotation window: {str(e)}")
            event.accept()

    def update_shape(self, shape):
        self.current_shape = shape

    def update_image_position_label(self):
        total_images = len(self.image_files)
        current_position = self.current_image_index + 1 if self.current_image_index >= 0 else 0
        self.image_position_label.setText(f"{current_position}/{total_images}")

    def update_classes_file(self):
        if not self.save_directory:
            return
        classes_file = os.path.join(self.save_directory, "classes.txt")
        try:
            with open(classes_file, 'w') as f:
                for class_id in sorted(self.id_to_class.keys()):
                    f.write(f"{self.id_to_class[class_id]}\n")
        except Exception as e:
            QMessageBox.warning(self, "Error", f"Failed to update classes file: {str(e)}")

    def get_class_id(self, label):
        if label not in self.class_to_id:
            class_id = len(self.class_to_id)
            self.class_to_id[label] = class_id
            self.id_to_class[class_id] = label
            self.update_classes_file()
        return self.class_to_id[label]

    def store_bbox(self, start, end, label, shape_type, existing_item=None):
        if not self.save_directory:
            QMessageBox.warning(self, "Save Directory Not Set", "Please select a save directory before annotating.")
            return

        try:
            image_width = self.original_image_size.width()
            image_height = self.original_image_size.height()

            if shape_type == "Polygon":
                if not self.polygon_points or len(self.polygon_points) < 3:
                    return
                coords = [(p.x() / image_width, p.y() / image_height) for p in self.polygon_points]
            else:
                if existing_item:
                    if isinstance(existing_item, QGraphicsPolygonItem):
                        points = existing_item.polygon()
                        coords = [(p.x() / image_width, p.y() / image_height) for p in points]
                    else:
                        rect = existing_item.rect()
                        x1 = rect.x()
                        y1 = rect.y()
                        x2 = x1 + rect.width()
                        y2 = y1 + rect.height()
                        x_center = ((x1 + x2) / 2) / image_width
                        y_center = ((y1 + y2) / 2) / image_height
                        width = (x2 - x1) / image_width
                        height = (y2 - y1) / image_height
                        coords = (x_center, y_center, width, height)
                else:
                    x1 = start.x() if start else 0
                    y1 = start.y() if start else 0
                    x2 = end.x() if end else 0
                    y2 = end.y() if end else 0
                    x1, x2 = min(x1, x2), max(x1, x2)
                    y1, y2 = min(y1, y2), max(y1, y2)
                    if shape_type == "Square":
                        size = max(x2 - x1, y2 - y1)
                        x2 = x1 + size
                        y2 = y1 + size
                    elif shape_type == "Circle":
                        width = x2 - x1
                        height = y2 - y1
                        x2 = x1 + width
                        y2 = y1 + height
                    x_center = ((x1 + x2) / 2) / image_width
                    y_center = ((y1 + y2) / 2) / image_height
                    width = (x2 - x1) / image_width
                    height = (y2 - y1) / image_height
                    coords = (x_center, y_center, width, height)

            image_path = self.image_files[self.current_image_index]
            if image_path not in self.bounding_boxes:
                self.bounding_boxes[image_path] = []
            if image_path not in self.label_items:
                self.label_items[image_path] = {}

            class_id = self.get_class_id(label)
            bbox_data = (coords, label, shape_type)

            if existing_item:
                removed = False
                for i, (old_coords, old_label, old_shape_type) in enumerate(self.bounding_boxes[image_path][:]):
                    if old_shape_type == shape_type and old_label == label:
                        for item, item_label in self.label_items[image_path].items():
                            if item == existing_item and item_label == old_label:
                                self.bounding_boxes[image_path].pop(i)
                                removed = True
                                break
                        if removed:
                            break
                self.bounding_boxes[image_path].append(bbox_data)
                self.label_items[image_path][existing_item] = label
            else:
                self.bounding_boxes[image_path].append(bbox_data)
                self.label_items[image_path][self.bbox] = label
                self.current_bbox_count += 1
                self.bbox_count_label.setText(f"Bounding boxes: {self.current_bbox_count}/{self.max_bboxes}")

            self.save_annotation(image_path, self.bounding_boxes[image_path])

            if self.bbox and not existing_item:
                self.last_bbox_per_image[image_path] = self.bbox
                self.bbox.setPen(QPen(QColor(0, 255, 0), 2))
                self.bbox.setFlag(QGraphicsItem.ItemIsSelectable, True)
                if self.bbox not in self.image_scene.items():
                    self.image_scene.addItem(self.bbox)

            self.update_label_reflector()

            if self.current_bbox_count >= self.max_bboxes:
                QTimer.singleShot(100, self.next_image)

        except Exception as e:
            QMessageBox.warning(self, "Error", f"Failed to store bounding box: {str(e)}")

    def set_bbox_limit(self):
        try:
            if self.bbox_input.text().strip() == "":
                self.max_bboxes = float('inf')
            else:
                self.max_bboxes = int(self.bbox_input.text())
            self.current_bbox_count = 0
            self.bbox_count_label.setText(f"Bounding boxes: {self.current_bbox_count}/{self.max_bboxes}")
        except ValueError:
            QMessageBox.warning(self, "Invalid Input", "Please enter a valid number.")

    def update_label_reflector(self):
        self.label_reflector.clear()
        if self.current_image_index >= 0:
            image_path = self.image_files[self.current_image_index]
            if image_path in self.bounding_boxes:
                seen_labels = set()
                for _, label, _ in self.bounding_boxes[image_path]:
                    if label not in seen_labels:
                        self.label_reflector.addItem(label)
                        seen_labels.add(label)

    def highlight_label_in_reflector(self, label):
        items = self.label_reflector.findItems(label, Qt.MatchExactly)
        for i in range(self.label_reflector.count()):
            item = self.label_reflector.item(i)
            if item in items:
                item.setBackground(Qt.gray)
            else:
                item.setBackground(Qt.white)

    def clear_highlight_in_reflector(self):
        for i in range(self.label_reflector.count()):
            item = self.label_reflector.item(i)
            item.setBackground(Qt.white)

    def highlight_bounding_boxes(self, label):
        image_path = self.image_files[self.current_image_index]
        for item in self.image_scene.items():
            if isinstance(item, (QGraphicsRectItem, QGraphicsEllipseItem, QGraphicsPolygonItem)):
                if image_path in self.label_items and item in self.label_items[image_path]:
                    item_label = self.label_items[image_path][item]
                    if item_label == label:
                        item.setPen(QPen(QColor(255, 255, 0), 3))
                    else:
                        item.setPen(QPen(QColor(0, 255, 0), 2))

    def on_label_reflector_selection_changed(self):
        selected_items = self.label_reflector.selectedItems()
        if selected_items:
            label = selected_items[0].text()
            self.highlight_bounding_boxes(label)
        else:
            image_path = self.image_files[self.current_image_index]
            for item in self.image_scene.items():
                if isinstance(item, (QGraphicsRectItem, QGraphicsEllipseItem, QGraphicsPolygonItem)):
                    item.setPen(QPen(QColor(0, 255, 0), 2))

    def on_file_selected(self):
        selected_items = self.file_list.selectedItems()
        if selected_items:
            selected_file = selected_items[0].text()
            self.save_annotations()
            self.display_image(selected_file)
            self.current_image_index = self.image_files.index(selected_file)
            self.update_label_reflector()

    def select_folder(self):
        folder = QFileDialog.getExistingDirectory(self, "Select Folder")
        if folder:
            self.load_images_from_folder(folder)

    def select_save_directory(self):
        directory = QFileDialog.getExistingDirectory(self, "Select Save Directory")
        if directory:
            self.save_directory = directory
            self.load_classes()

    def load_classes(self):
        classes_file = os.path.join(self.save_directory, "classes.txt")
        self.class_to_id.clear()
        self.id_to_class.clear()
        try:
            if os.path.exists(classes_file):
                with open(classes_file, 'r') as f:
                    for class_id, label in enumerate(f.read().splitlines()):
                        self.class_to_id[label] = class_id
                        self.id_to_class[class_id] = label
        except Exception as e:
            QMessageBox.warning(self, "Error", f"Failed to load classes: {str(e)}")

    def load_images_from_folder(self, folder):
        try:
            self.image_files = [os.path.join(folder, f) for f in os.listdir(folder) if
                                f.lower().endswith(('.png', '.jpg', '.jpeg'))]
            self.file_list.clear()
            self.file_list.addItems(self.image_files)
            if self.image_files:
                self.current_image_index = 0
                self.current_bbox_count = 0
                self.file_list.setCurrentRow(self.current_image_index)
                self.display_image(self.image_files[0])
                self.update_image_position_label()
        except Exception as e:
            QMessageBox.warning(self, "Error", f"Failed to load images: {str(e)}")

    def display_image(self, image_path):
        try:
            image = QImage(image_path)
            if image.isNull():
                raise ValueError("Failed to load image")
            pixmap = QPixmap.fromImage(image)

            self.image_scene.clear()

            self.preview_path = QGraphicsPathItem()
            self.preview_path.setPen(QPen(QColor(255, 100, 100), 2, Qt.DashLine))
            self.preview_path.setZValue(100)
            self.image_scene.addItem(self.preview_path)
            self.preview_path.setPath(QPainterPath())

            self.image_scene.addPixmap(pixmap)
            self.original_image_size = pixmap.size()
            self.image_scene.setSceneRect(QRectF(pixmap.rect()))
            self.image_viewer.fitInView(self.image_scene.sceneRect(), Qt.KeepAspectRatio)

            self.vertical_line = QGraphicsLineItem()
            self.vertical_line.setPen(QPen(QColor(0, 0, 255), 2, Qt.DashLine))
            self.image_scene.addItem(self.vertical_line)

            self.horizontal_line = QGraphicsLineItem()
            self.horizontal_line.setPen(QPen(QColor(0, 0, 255), 2, Qt.DashLine))
            self.image_scene.addItem(self.horizontal_line)

            if image_path in self.label_items:
                del self.label_items[image_path]
            self.label_items[image_path] = {}

            if image_path in self.bounding_boxes:
                image_width = self.original_image_size.width()
                image_height = self.original_image_size.height()
                for coords, label, shape_type in self.bounding_boxes[image_path]:
                    if shape_type == "Polygon":
                        points = [QPointF(x * image_width, y * image_height) for x, y in coords]
                        rect = QGraphicsPolygonItem(QPolygonF(points))
                    else:
                        x_center, y_center, width, height = coords
                        x1 = (x_center - width / 2) * image_width
                        y1 = (y_center - height / 2) * image_height
                        rect_width = width * image_width
                        rect_height = height * image_height
                        if shape_type == "Circle":
                            rect = QGraphicsEllipseItem(x1, y1, rect_width, rect_height)
                        elif shape_type == "Square":
                            size = max(rect_width, rect_height)
                            rect = QGraphicsRectItem(x1, y1, size, size)
                        else:  # Rectangle
                            rect = QGraphicsRectItem(x1, y1, rect_width, rect_height)

                    rect.setPen(QPen(QColor(0, 255, 0), 2))
                    rect.setFlag(QGraphicsItem.ItemIsSelectable, True)
                    self.image_scene.addItem(rect)

                    if (image_path in self.last_bbox_per_image and
                            self.last_bbox_per_image[image_path] is None and
                            self.bounding_boxes[image_path][-1] == (coords, label, shape_type)):
                        self.last_bbox_per_image[image_path] = rect

                    self.label_items[image_path][rect] = label

            self.current_bbox_count = len(self.bounding_boxes.get(image_path, []))
            self.bbox_count_label.setText(f"Bounding boxes: {self.current_bbox_count}/{self.max_bboxes}")
            self.update_label_reflector()
            self.update_image_position_label()

        except Exception as e:
            QMessageBox.warning(self, "Error", f"Failed to display image: {str(e)}")

    def previous_image(self):
        if self.current_image_index > 0:
            try:
                self.save_annotations()
                self.current_image_index -= 1
                self.file_list.setCurrentRow(self.current_image_index)
                self.display_image(self.image_files[self.current_image_index])
                self.update_label_reflector()
                self.update_image_position_label()
            except Exception as e:
                QMessageBox.warning(self, "Error", f"Failed to load previous image: {str(e)}")

    def next_image(self):
        if self.current_image_index < len(self.image_files) - 1:
            try:
                self.save_annotations()
                self.current_image_index += 1
                self.current_bbox_count = 0
                self.file_list.setCurrentRow(self.current_image_index)
                self.display_image(self.image_files[self.current_image_index])
                self.update_label_reflector()
                self.update_image_position_label()
            except Exception as e:
                QMessageBox.warning(self, "Error", f"Failed to load next image: {str(e)}")

    def zoom(self, factor, mouse_pos=None):
        try:
            if not self.image_scene or self.image_scene.sceneRect().isEmpty():
                return

            if mouse_pos is not None:
                if isinstance(mouse_pos, tuple) and len(mouse_pos) == 2:
                    cursor_pos = QPointF(mouse_pos[0], mouse_pos[1])
                elif isinstance(mouse_pos, QPoint):
                    cursor_pos = QPointF(mouse_pos)
                else:
                    cursor_pos = self.image_viewer.viewport().mapFromGlobal(QCursor.pos())
                    cursor_pos = self.image_viewer.mapToScene(cursor_pos)
            else:
                cursor_pos = self.image_scene.sceneRect().center()

            if not self.image_scene.sceneRect().contains(cursor_pos):
                cursor_pos = self.image_scene.sceneRect().center()

            old_pos = self.image_viewer.mapToScene(self.image_viewer.mapFromScene(cursor_pos))
            if old_pos.isNull():
                return

            self.image_viewer.scale(factor, factor)

            new_pos = self.image_viewer.mapToScene(self.image_viewer.mapFromScene(cursor_pos))
            if new_pos.isNull():
                return

            delta = new_pos - old_pos
            self.image_viewer.translate(delta.x(), delta.y())

        except Exception as e:
            QMessageBox.warning(self, "Error", f"Zoom error: {str(e)}")

    def zoom_in(self):
        mouse_pos = self.image_viewer.viewport().mapFromGlobal(QCursor.pos())
        self.zoom(1.1, mouse_pos)

    def zoom_out(self):
        mouse_pos = self.image_viewer.viewport().mapFromGlobal(QCursor.pos())
        self.zoom(0.9, mouse_pos)

    def get_resize_handle(self, item, pos):
        try:
            transform = self.image_viewer.transform()
            scale = transform.m11()
            handle_size = 60 / max(scale, 0.1)

            if isinstance(item, QGraphicsPolygonItem):
                points = item.polygon()
                for i, point in enumerate(points):
                    scene_point = item.mapToScene(point)
                    handle_rect = QRectF(scene_point.x() - handle_size / 2, scene_point.y() - handle_size / 2,
                                         handle_size, handle_size)
                    if handle_rect.contains(pos):
                        return f"vertex_{i}"
                return None
            else:
                rect = item.rect()
                scene_rect = item.mapRectToScene(rect)
                handles = {
                    "top_left": QRectF(scene_rect.left() - handle_size / 2, scene_rect.top() - handle_size / 2,
                                       handle_size, handle_size),
                    "top_right": QRectF(scene_rect.right() - handle_size / 2, scene_rect.top() - handle_size / 2,
                                        handle_size, handle_size),
                    "bottom_left": QRectF(scene_rect.left() - handle_size / 2, scene_rect.bottom() - handle_size / 2,
                                          handle_size, handle_size),
                    "bottom_right": QRectF(scene_rect.right() - handle_size / 2, scene_rect.bottom() - handle_size / 2,
                                           handle_size, handle_size),
                    "top": QRectF(scene_rect.left(), scene_rect.top() - handle_size / 2, scene_rect.width(),
                                  handle_size),
                    "bottom": QRectF(scene_rect.left(), scene_rect.bottom() - handle_size / 2, scene_rect.width(),
                                     handle_size),
                    "left": QRectF(scene_rect.left() - handle_size / 2, scene_rect.top(), handle_size,
                                   scene_rect.height()),
                    "right": QRectF(scene_rect.right() - handle_size / 2, scene_rect.top(), handle_size,
                                    scene_rect.height())
                }
                for handle, handle_rect in handles.items():
                    if handle_rect.contains(pos):
                        return handle
                return None
        except Exception as e:
            QMessageBox.warning(self, "Error", f"Error getting resize handle: {str(e)}")
            return None

    def update_bbox_during_resize(self, item, handle, new_pos):
        try:
            image_rect = self.image_scene.sceneRect()
            image_path = self.image_files[self.current_image_index]
            shape_type = None
            label = None
            if image_path in self.label_items and item in self.label_items[image_path]:
                label = self.label_items[image_path][item]
                for coords, lbl, st in self.bounding_boxes.get(image_path, []):
                    if lbl == label:
                        shape_type = st
                        break
            if not shape_type or not label:
                return

            if isinstance(item, QGraphicsPolygonItem):
                points = item.polygon()
                index = int(handle.split('_')[1])
                new_x = max(image_rect.left(), min(new_pos.x(), image_rect.right()))
                new_y = max(image_rect.top(), min(new_pos.y(), image_rect.bottom()))
                points[index] = QPointF(new_x, new_y)
                item.setPolygon(QPolygonF(points))
                if index < len(self.vertex_handles):
                    self.vertex_handles[index].setRect(new_x - 10, new_y - 10, 20, 20)
            else:
                rect = item.rect()
                new_rect = QRectF(rect)
                new_pos_item = item.mapFromScene(new_pos)

                if handle == "top_left":
                    new_rect.setTopLeft(new_pos_item)
                elif handle == "top_right":
                    new_rect.setTopRight(new_pos_item)
                elif handle == "bottom_left":
                    new_rect.setBottomLeft(new_pos_item)
                elif handle == "bottom_right":
                    new_rect.setBottomRight(new_pos_item)
                elif handle == "top":
                    new_rect.setTop(new_pos_item.y())
                elif handle == "bottom":
                    new_rect.setBottom(new_pos_item.y())
                elif handle == "left":
                    new_rect.setLeft(new_pos_item.x())
                elif handle == "right":
                    new_rect.setRight(new_pos_item.x())

                if shape_type == "Square":
                    size = max(new_rect.width(), new_rect.height())
                    if handle in ["top_left", "bottom_right"]:
                        new_rect = QRectF(new_rect.topLeft(), QSizeF(size, size))
                    elif handle in ["top_right", "bottom_left"]:
                        new_rect = QRectF(QPointF(new_rect.right() - size, new_rect.top()), QSizeF(size, size))
                    elif handle == "top":
                        new_rect.setBottom(new_rect.top() + new_rect.width())
                    elif handle == "bottom":
                        new_rect.setTop(new_rect.bottom() - new_rect.width())
                    elif handle == "left":
                        new_rect.setRight(new_rect.left() + new_rect.height())
                    elif handle == "right":
                        new_rect.setLeft(new_rect.right() - new_rect.height())
                elif shape_type == "Circle":
                    size = max(new_rect.width(), new_rect.height())
                    new_rect = QRectF(new_rect.topLeft(), QSizeF(size, size))

                scene_rect = item.mapRectToScene(new_rect)
                scene_rect = QRectF(
                    max(image_rect.left(), scene_rect.left()),
                    max(image_rect.top(), scene_rect.top()),
                    min(scene_rect.width(), image_rect.width() - scene_rect.left() + image_rect.left()),
                    min(scene_rect.height(), image_rect.height() - scene_rect.top() + image_rect.top())
                )
                new_rect = item.mapRectFromScene(scene_rect)

                item.setRect(new_rect.normalized())

                self.store_bbox(None, None, label, shape_type, item)
        except Exception as e:
            QMessageBox.warning(self, "Error", f"Error updating bounding box during resize: {str(e)}")

    def update_moved_bbox_annotation(self, item):
        try:
            if not item or self.current_image_index < 0:
                return False

            image_path = self.image_files[self.current_image_index]
            if image_path not in self.bounding_boxes or image_path not in self.label_items:
                return False

            image_width = self.original_image_size.width()
            image_height = self.original_image_size.height()

            current_label = self.label_items[image_path].get(item)
            if not current_label:
                return False

            found = False
            for i, (coords, label, shape_type) in enumerate(self.bounding_boxes[image_path][:]):
                if label == current_label and self.label_items[image_path].get(item) == label:
                    if isinstance(item, QGraphicsPolygonItem) and shape_type == "Polygon":
                        new_points = [(p.x() / image_width, p.y() / image_height) for p in item.polygon()]
                        new_coords = new_points
                    else:
                        rect = item.rect()
                        scene_rect = item.mapRectToScene(rect)
                        image_rect = self.image_scene.sceneRect()
                        scene_rect = QRectF(
                            max(image_rect.left(), scene_rect.left()),
                            max(image_rect.top(), scene_rect.top()),
                            min(scene_rect.width(), image_rect.width() - scene_rect.left() + image_rect.left()),
                            min(scene_rect.height(), image_rect.height() - scene_rect.top() + image_rect.top())
                        )
                        new_rect = item.mapRectFromScene(scene_rect)
                        item.setRect(new_rect.normalized())

                        x_center = (new_rect.x() + new_rect.width() / 2) / image_width
                        y_center = (new_rect.y() + new_rect.height() / 2) / image_height
                        width = new_rect.width() / image_width
                        height = new_rect.height() / image_height
                        new_coords = (x_center, y_center, width, height)

                    self.bounding_boxes[image_path][i] = (new_coords, current_label, shape_type)
                    found = True
                    break

            if not found:
                return False

            self.save_annotation(image_path, self.bounding_boxes[image_path])
            return True
        except Exception as e:
            QMessageBox.warning(self, "Error", f"Error updating moved bounding box: {str(e)}")
            return False

    def add_vertex_handle(self, pos, is_first=False):
        radius = 18 if is_first else 16
        handle = QGraphicsEllipseItem(pos.x() - radius, pos.y() - radius, radius * 2, radius * 2)
        color = QColor(255, 80, 80, 220) if is_first else QColor(255, 0, 0, 200)
        handle.setBrush(color)
        handle.setPen(QPen(QColor('white'), 3))
        handle.setZValue(101)
        self.image_scene.addItem(handle)
        self.vertex_handles.append(handle)

    def clear_vertex_handles(self):
        for handle in self.vertex_handles:
            if handle.scene():
                self.image_scene.removeItem(handle)
        self.vertex_handles.clear()

    def eventFilter(self, source, event):
        try:
            if source is self.image_viewer.viewport() and isinstance(event, QMouseEvent):
                scene_pos = self.image_viewer.mapToScene(event.pos())
                scene_rect = self.image_scene.sceneRect()

                if event.type() == QEvent.MouseMove:
                    self.vertical_line.setPen(QPen(QColor(0, 0, 255), 2, Qt.DashLine))
                    self.horizontal_line.setPen(QPen(QColor(0, 0, 255), 2, Qt.DashLine))

                    if scene_rect.contains(scene_pos):
                        self.vertical_line.setLine(scene_pos.x(), scene_rect.top(), scene_pos.x(), scene_rect.bottom())
                        self.horizontal_line.setLine(scene_rect.left(), scene_pos.y(), scene_rect.right(),
                                                     scene_pos.y())
                        self.vertical_line.show()
                        self.horizontal_line.show()
                    else:
                        self.vertical_line.hide()
                        self.horizontal_line.hide()

                    cursor_set = False
                    items = self.image_scene.items(scene_pos)
                    for item in items:
                        if isinstance(item, (QGraphicsRectItem, QGraphicsEllipseItem, QGraphicsPolygonItem)):
                            handle = self.get_resize_handle(item, scene_pos)
                            if handle:
                                if handle in ["top_left", "bottom_right"]:
                                    self.image_viewer.setCursor(QCursor(Qt.SizeFDiagCursor))
                                elif handle in ["top_right", "bottom_left"]:
                                    self.image_viewer.setCursor(QCursor(Qt.SizeBDiagCursor))
                                elif handle in ["top", "bottom"]:
                                    self.image_viewer.setCursor(QCursor(Qt.SizeVerCursor))
                                elif handle in ["left", "right"]:
                                    self.image_viewer.setCursor(QCursor(Qt.SizeHorCursor))
                                elif handle.startswith("vertex_"):
                                    self.image_viewer.setCursor(QCursor(Qt.PointingHandCursor))
                                cursor_set = True
                                break

                    if self.is_drawing_polygon and self.first_polygon_point:
                        if self.is_near_first_point(scene_pos) or self.is_click_on_first_handle(scene_pos):
                            self.image_viewer.setCursor(QCursor(Qt.PointingHandCursor))
                            cursor_set = True

                    if not cursor_set:
                        self.image_viewer.setCursor(QCursor(Qt.CrossCursor))

                    self.clear_highlight_in_reflector()
                    for item in items:
                        if isinstance(item, (QGraphicsRectItem, QGraphicsEllipseItem, QGraphicsPolygonItem)):
                            image_path = self.image_files[self.current_image_index]
                            if image_path in self.label_items and item in self.label_items[image_path]:
                                label = self.label_items[image_path][item]
                                self.highlight_label_in_reflector(label)
                                break

                    if self.drawing and self.bbox and self.current_shape != "Polygon":
                        self.end_point = scene_pos
                        if self.current_shape == "Square":
                            size = max(abs(scene_pos.x() - self.start_point.x()),
                                       abs(scene_pos.y() - self.start_point.y()))
                            x1 = self.start_point.x()
                            y1 = self.start_point.y()
                            if scene_pos.x() < x1:
                                x1 = self.start_point.x() - size
                            if scene_pos.y() < y1:
                                y1 = self.start_point.y() - size
                            self.bbox.setRect(QRectF(x1, y1, size, size))
                        elif self.current_shape == "Circle":
                            self.bbox.setRect(QRectF(self.start_point, scene_pos).normalized())
                        else:
                            self.bbox.setRect(QRectF(self.start_point, scene_pos).normalized())
                        self.bbox.show()

                    elif self.is_drawing_polygon and self.polygon_points:
                        path = QPainterPath()
                        path.moveTo(self.polygon_points[0])
                        for p in self.polygon_points[1:]:
                            path.lineTo(p)
                        path.lineTo(scene_pos)
                        self.preview_path.setPath(path)
                        self.image_scene.update()

                    elif self.resizing and self.selected_bbox and self.resize_handle:
                        self.update_bbox_during_resize(self.selected_bbox, self.resize_handle, scene_pos)
                        self.image_scene.update()

                elif event.type() == QEvent.MouseButtonPress and event.button() == Qt.LeftButton:
                    if not self.save_directory:
                        QMessageBox.warning(self, "Save Directory Not Set",
                                            "Please select a save directory before annotating.")
                        return True

                    items = self.image_scene.items(scene_pos)

                    selected_item = None
                    resize_handle = None

                    for item in items:
                        if isinstance(item, (QGraphicsRectItem, QGraphicsEllipseItem, QGraphicsPolygonItem)) and item != self.bbox:
                            handle = self.get_resize_handle(item, scene_pos)
                            if handle:
                                selected_item = item
                                resize_handle = handle
                                break

                    if not selected_item:
                        for item in items:
                            if isinstance(item, (QGraphicsRectItem, QGraphicsEllipseItem, QGraphicsPolygonItem)) and item != self.bbox:
                                item_pos = item.mapFromScene(scene_pos)
                                if isinstance(item, QGraphicsPolygonItem):
                                    poly = item.polygon()
                                    if poly.containsPoint(item_pos, Qt.OddEvenFill):
                                        selected_item = item
                                        break
                                elif isinstance(item, QGraphicsEllipseItem):
                                    rect = item.rect()
                                    center = rect.center()
                                    dx = (item_pos.x() - center.x()) / (rect.width() / 2)
                                    dy = (item_pos.y() - center.y()) / (rect.height() / 2)
                                    if dx * dx + dy * dy <= 1:
                                        selected_item = item
                                        break
                                else:
                                    if item.rect().contains(item_pos):
                                        selected_item = item
                                        break

                    if resize_handle:
                        if self.selected_bbox and self.selected_bbox != selected_item:
                            self.selected_bbox.setSelected(False)
                            self.selected_bbox.setPen(QPen(QColor(0, 255, 0), 2))
                            self.selected_bbox.setFlag(QGraphicsItem.ItemIsSelectable, True)
                            if isinstance(self.selected_bbox, QGraphicsPolygonItem):
                                self.clear_vertex_handles()
                        self.selected_bbox = selected_item
                        self.selected_bbox.setSelected(True)
                        self.selected_bbox.setPen(QPen(QColor(255, 255, 0), 3))
                        self.selected_bbox.setFlag(QGraphicsItem.ItemIsSelectable, False)
                        self.resizing = True
                        self.resize_handle = resize_handle
                        self.resize_start_pos = scene_pos
                        self.drawing = False
                        self.image_scene.update()
                        return True

                    elif selected_item:
                        if self.selected_bbox and self.selected_bbox != selected_item:
                            self.selected_bbox.setSelected(False)
                            self.selected_bbox.setPen(QPen(QColor(0, 255, 0), 2))
                            self.selected_bbox.setFlag(QGraphicsItem.ItemIsSelectable, True)
                            if isinstance(self.selected_bbox, QGraphicsPolygonItem):
                                self.clear_vertex_handles()
                        self.selected_bbox = selected_item
                        self.selected_bbox.setSelected(True)
                        self.selected_bbox.setPen(QPen(QColor(255, 255, 0), 3))
                        self.selected_bbox.setFlag(QGraphicsItem.ItemIsSelectable, True)
                        if isinstance(self.selected_bbox, QGraphicsPolygonItem):
                            self.clear_vertex_handles()
                            for p in self.selected_bbox.polygon():
                                self.add_vertex_handle(p)
                        self.drawing = False
                        self.resizing = False
                        self.resize_handle = None
                        self.resize_start_pos = None
                        self.image_scene.update()
                        return True

                    else:
                        if self.selected_bbox:
                            self.selected_bbox.setSelected(False)
                            self.selected_bbox.setPen(QPen(QColor(0, 255, 0), 2))
                            self.selected_bbox.setFlag(QGraphicsItem.ItemIsSelectable, True)
                            if isinstance(self.selected_bbox, QGraphicsPolygonItem):
                                self.clear_vertex_handles()
                            self.selected_bbox = None
                            self.resizing = False
                            self.resize_handle = None
                            self.resize_start_pos = None
                            self.update_label_reflector()
                            self.image_scene.update()

                        if self.current_shape != "Polygon":
                            if not self.drawing:
                                self.start_point = scene_pos
                                if self.current_shape == "Circle":
                                    self.bbox = QGraphicsEllipseItem(QRectF(self.start_point, self.start_point))
                                else:
                                    self.bbox = QGraphicsRectItem(QRectF(self.start_point, self.start_point))
                                self.bbox.setPen(QPen(QColor(0, 255, 255), 2))
                                self.bbox.setFlag(QGraphicsItem.ItemIsSelectable, True)
                                self.image_scene.addItem(self.bbox)
                                self.drawing = True
                        else:  # Polygon mode
                            if not self.is_drawing_polygon:
                                # Start new polygon
                                self.is_drawing_polygon = True
                                self.polygon_points = [scene_pos]
                                self.first_polygon_point = scene_pos
                                self.clear_vertex_handles()
                                self.add_vertex_handle(scene_pos, is_first=True)
                                path = QPainterPath()
                                path.moveTo(scene_pos)
                                self.preview_path.setPath(path)
                            else:
                                if len(self.polygon_points) >= 3 and (
                                        self.is_near_first_point(scene_pos) or self.is_click_on_first_handle(scene_pos)
                                ):
                                    event.accept()


                                    label = self.prompt_label()
                                    if label:
                                        polygon = QGraphicsPolygonItem(QPolygonF(self.polygon_points))
                                        polygon.setPen(QPen(QColor(0, 255, 0), 2))
                                        polygon.setFlag(QGraphicsItem.ItemIsSelectable, True)
                                        self.image_scene.addItem(polygon)
                                        self.bbox = polygon
                                        self.store_bbox(None, None, label, "Polygon")
                                    self.preview_path.setPath(QPainterPath())
                                    self.clear_vertex_handles()
                                    self.polygon_points = []
                                    self.is_drawing_polygon = False
                                    self.first_polygon_point = None
                                else:
                                    # Add normal point
                                    self.polygon_points.append(scene_pos)
                                    self.add_vertex_handle(scene_pos)
                                    path = QPainterPath()
                                    path.moveTo(self.polygon_points[0])
                                    for p in self.polygon_points[1:]:
                                        path.lineTo(p)
                                    self.preview_path.setPath(path)

                        self.image_scene.update()
                        return True

                elif event.type() == QEvent.MouseButtonRelease and event.button() == Qt.LeftButton:
                    if self.drawing and self.bbox and self.current_shape != "Polygon":
                        self.end_point = scene_pos
                        if self.start_point and self.end_point and self.start_point != self.end_point:
                            if self.current_shape == "Square":
                                size = max(abs(self.end_point.x() - self.start_point.x()),
                                           abs(self.end_point.y() - self.start_point.y()))
                                x1 = self.start_point.x()
                                y1 = self.start_point.y()
                                if self.end_point.x() < x1:
                                    x1 = self.start_point.x() - size
                                if self.end_point.y() < y1:
                                    y1 = self.start_point.y() - size
                                self.bbox.setRect(QRectF(x1, y1, size, size))
                            else:
                                self.bbox.setRect(QRectF(self.start_point, self.end_point).normalized())
                            label = self.prompt_label()
                            if label:
                                self.bbox.setPen(QPen(QColor(0, 255, 0), 2))
                                self.store_bbox(self.start_point, self.end_point, label, self.current_shape)
                            else:
                                self.image_scene.removeItem(self.bbox)
                            self.bbox = None
                            self.drawing = False
                            self.start_point = None
                            self.end_point = None
                        else:
                            self.image_scene.removeItem(self.bbox)
                            self.bbox = None
                            self.drawing = False
                            self.start_point = None
                            self.end_point = None
                    elif self.resizing and self.selected_bbox and self.resize_handle:
                        self.selected_bbox.setFlag(QGraphicsItem.ItemIsSelectable, True)
                        image_path = self.image_files[self.current_image_index]
                        if image_path in self.label_items and self.selected_bbox in self.label_items[image_path]:
                            label = self.label_items[image_path][self.selected_bbox]
                            shape_type = None
                            for coords, lbl, st in self.bounding_boxes.get(image_path, []):
                                if lbl == label:
                                    shape_type = st
                                    break
                            if shape_type and label:
                                self.store_bbox(None, None, label, shape_type, self.selected_bbox)
                        self.selected_bbox.setSelected(True)
                        self.selected_bbox.setPen(QPen(QColor(255, 255, 0), 3))
                        self.resizing = False
                        self.resize_handle = None
                        self.resize_start_pos = None
                        self.image_scene.update()

                elif event.type() == QEvent.MouseButtonPress and event.button() == Qt.RightButton:
                    if self.is_drawing_polygon:
                        self.preview_path.setPath(QPainterPath())
                        self.clear_vertex_handles()
                        self.polygon_points = []
                        self.is_drawing_polygon = False
                        self.first_polygon_point = None
                        self.image_scene.update()
                        return True

                elif event.type() == QEvent.Wheel:
                    if event.angleDelta().y() > 0:
                        self.zoom_in()
                    else:
                        self.zoom_out()
                    return True

        except Exception as e:
            QMessageBox.warning(self, "Error", f"Mouse event error: {str(e)}")

        return super().eventFilter(source, event)

    def delete_selected_bounding_box(self):
        try:
            if not self.selected_bbox or not isinstance(self.selected_bbox,
                                                        (QGraphicsRectItem, QGraphicsEllipseItem, QGraphicsPolygonItem)):
                QMessageBox.warning(self, "Warning", "No bounding box selected.")
                return

            image_path = self.image_files[self.current_image_index]
            if image_path not in self.bounding_boxes:
                return

            image_width = self.original_image_size.width()
            image_height = self.original_image_size.height()
            removed = False
            label = self.label_items[image_path].get(self.selected_bbox, "Unknown")

            if isinstance(self.selected_bbox, QGraphicsPolygonItem):
                poly_points = [(p.x() / image_width, p.y() / image_height) for p in self.selected_bbox.polygon()]
                for i, (coords, lbl, shape_type) in enumerate(self.bounding_boxes[image_path][:]):
                    if shape_type == "Polygon" and len(coords) == len(poly_points):
                        match = all(abs(coords[j][0] - poly_points[j][0]) < 1e-5 and
                                    abs(coords[j][1] - poly_points[j][1]) < 1e-5
                                    for j in range(len(coords)))
                        if match:
                            self.bounding_boxes[image_path].pop(i)
                            removed = True
                            break
            else:
                rect = self.selected_bbox.rect()
                x_center = (rect.x() + rect.width() / 2) / image_width
                y_center = (rect.y() + rect.height() / 2) / image_height
                width = rect.width() / image_width
                height = rect.height() / image_height
                for i, (coords, lbl, shape_type) in enumerate(self.bounding_boxes[image_path][:]):
                    if shape_type != "Polygon" and len(coords) == 4:
                        if (abs(coords[0] - x_center) < 1e-5 and
                                abs(coords[1] - y_center) < 1e-5 and
                                abs(coords[2] - width) < 1e-5 and
                                abs(coords[3] - height) < 1e-5):
                            self.bounding_boxes[image_path].pop(i)
                            removed = True
                            break

            if not removed:
                return

            if self.selected_bbox in self.image_scene.items():
                self.image_scene.removeItem(self.selected_bbox)

            if image_path in self.label_items and self.selected_bbox in self.label_items[image_path]:
                del self.label_items[image_path][self.selected_bbox]

            if image_path in self.last_bbox_per_image and self.last_bbox_per_image[image_path] == self.selected_bbox:
                self.last_bbox_per_image[image_path] = None

            self.clear_vertex_handles()
            self.selected_bbox = None
            self.resizing = False
            self.resize_handle = None
            self.resize_start_pos = None

            self.current_bbox_count = max(0, self.current_bbox_count - 1)
            self.bbox_count_label.setText(f"Bounding boxes: {self.current_bbox_count}/{self.max_bboxes}")

            if self.bounding_boxes.get(image_path):
                self.save_annotation(image_path, self.bounding_boxes[image_path])
            else:
                image_name = os.path.splitext(os.path.basename(image_path))[0]
                annotation_file_path = os.path.join(self.save_directory, f"{image_name}.txt")
                if os.path.exists(annotation_file_path):
                    os.remove(annotation_file_path)
                if image_path in self.bounding_boxes:
                    del self.bounding_boxes[image_path]

            self.update_label_reflector()

        except Exception as e:
            QMessageBox.warning(self, "Error", f"Error deleting bounding box: {str(e)}")

    def prompt_label(self):
        try:
            if self.default_label_checkbox.isChecked() and self.default_label_textbox.text():
                return self.default_label_textbox.text()
            else:
                label, ok = QInputDialog.getText(self, "Label Input", "Enter label for the bounding box:")
                if not ok or not label:
                    return None
                return label
        except Exception as e:
            QMessageBox.warning(self, "Error", f"Error prompting label: {str(e)}")
            return None

    def save_annotation(self, image_path, annotations):
        try:
            if not self.save_directory:
                return
            if not annotations:
                image_name = os.path.splitext(os.path.basename(image_path))[0]
                annotation_file_path = os.path.join(self.save_directory, f"{image_name}.txt")
                if os.path.exists(annotation_file_path):
                    os.remove(annotation_file_path)
                return
            image_name = os.path.splitext(os.path.basename(image_path))[0]
            annotation_file_path = os.path.join(self.save_directory, f"{image_name}.txt")
            with open(annotation_file_path, 'w') as f:
                for coords, label, shape_type in annotations:
                    class_id = self.get_class_id(label)
                    if shape_type == "Polygon":
                        coords_str = " ".join(f"{x:.6f} {y:.6f}" for x, y in coords)
                        f.write(f"{class_id} {coords_str}\n")
                    else:
                        x_center, y_center, width, height = coords
                        f.write(f"{class_id} {x_center:.6f} {y_center:.6f} {width:.6f} {height:.6f}\n")
        except Exception as e:
            QMessageBox.warning(self, "Error", f"Error saving annotation: {str(e)}")

    def save_annotations(self):
        try:
            if self.save_directory and self.current_image_index >= 0:
                image_path = self.image_files[self.current_image_index]
                annotations = self.bounding_boxes.get(image_path, [])
                self.save_annotation(image_path, annotations)
        except Exception as e:
            QMessageBox.warning(self, "Error", f"Error saving annotations: {str(e)}")

    def load_annotations(self):
        try:
            if self.save_directory:
                self.bounding_boxes.clear()
                for image_path in self.image_files:
                    image_name = os.path.splitext(os.path.basename(image_path))[0]
                    annotation_file_path = os.path.join(self.save_directory, f"{image_name}.txt")
                    if os.path.exists(annotation_file_path):
                        annotations = []
                        with open(annotation_file_path, 'r') as f:
                            for line in f:
                                parts = line.strip().split()
                                if len(parts) >= 5:
                                    class_id = int(parts[0])
                                    coords = list(map(float, parts[1:]))
                                    if class_id in self.id_to_class:
                                        label = self.id_to_class[class_id]
                                        if len(coords) == 4:
                                            shape_type = "Rectangle"  # Default for bounding box
                                            annotations.append((coords, label, shape_type))
                                        elif len(coords) >= 6 and len(coords) % 2 == 0:
                                            shape_type = "Polygon"
                                            points = [(coords[i], coords[i + 1]) for i in range(0, len(coords), 2)]
                                            annotations.append((points, label, shape_type))
                        if annotations:
                            self.bounding_boxes[image_path] = annotations
        except Exception as e:
            QMessageBox.warning(self, "Error", f"Error loading annotations: {str(e)}")

    def keyPressEvent(self, event):
        try:
            if event.key() == Qt.Key_Delete:
                self.delete_selected_bounding_box()
            elif event.key() == Qt.Key_Escape:
                if self.drawing:
                    self.image_scene.removeItem(self.bbox)
                    self.bbox = None
                    self.drawing = False
                    self.start_point = None
                    self.end_point = None
                if self.is_drawing_polygon:
                    self.preview_path.setPath(QPainterPath())
                    self.clear_vertex_handles()
                    self.polygon_points = []
                    self.is_drawing_polygon = False
                    self.first_polygon_point = None
                self.image_scene.update()
            elif event.key() == Qt.Key_Right or event.key() == Qt.Key_D:
                if self.current_image_index < len(self.image_files) - 1:
                    self.save_annotations()
                    self.current_image_index += 1
                    self.current_bbox_count = 0
                    self.file_list.setCurrentRow(self.current_image_index)
                    self.display_image(self.image_files[self.current_image_index])
                    self.update_label_reflector()
                    self.update_image_position_label()
            elif event.key() == Qt.Key_Left or event.key() == Qt.Key_A:
                if self.current_image_index > 0:
                    self.save_annotations()
                    prev_image_index = self.current_image_index - 1
                    prev_image_path = self.image_files[prev_image_index]
                    if (prev_image_path in self.last_bbox_per_image and
                            self.last_bbox_per_image[prev_image_path] and
                            prev_image_path in self.bounding_boxes and
                            self.bounding_boxes[prev_image_path]):
                        last_bbox = self.last_bbox_per_image[prev_image_path]
                        if last_bbox in self.image_scene.items():
                            self.image_scene.removeItem(last_bbox)
                        self.bounding_boxes[prev_image_path].pop()
                        self.save_annotation(prev_image_path, self.bounding_boxes[prev_image_path])
                        self.last_bbox_per_image[prev_image_path] = None
                    self.current_image_index -= 1
                    self.file_list.setCurrentRow(self.current_image_index)
                    self.display_image(self.image_files[self.current_image_index])
                    self.update_label_reflector()
                    self.update_image_position_label()
                elif self.current_image_index == 0:
                    self.save_annotations()
                    image_path = self.image_files[0]
                    if (image_path in self.last_bbox_per_image and
                            self.last_bbox_per_image[image_path] and
                            image_path in self.bounding_boxes and
                            self.bounding_boxes[image_path]):
                        last_bbox = self.last_bbox_per_image[image_path]
                        if last_bbox in self.image_scene.items():
                            self.image_scene.removeItem(last_bbox)
                        self.bounding_boxes[image_path].pop()
                        self.save_annotation(image_path, self.bounding_boxes[image_path])
                        self.last_bbox_per_image[image_path] = None
                        self.image_scene.removeItem(self.vertical_line)
                        self.image_scene.removeItem(self.horizontal_line)
                        self.display_image(image_path)
                        self.update_label_reflector()
                        self.update_image_position_label()
            elif event.key() == Qt.Key_Return or event.key() == Qt.Key_Enter:
                if self.selected_bbox and not self.resizing:
                    self.save_annotations()
            super().keyPressEvent(event)
        except Exception as e:
            QMessageBox.warning(self, "Error", f"Key press error: {str(e)}")


def main():
    try:
        app = QApplication(sys.argv)
        annotation_tool = AnnotationTool()
        annotation_tool.show()
        annotation_tool.load_annotations()
        sys.exit(app.exec_())
    except Exception as e:
        print("Application failed to start:", str(e))
        sys.exit(1)


if __name__ == '__main__':
    main()