import os
from PyQt5.QtWidgets import QMainWindow, QPushButton, QLabel, QWidget, QVBoxLayout, QHBoxLayout, QSpacerItem, \
    QSizePolicy, QMessageBox, QScrollArea, QApplication
from PyQt5.QtGui import QFont, QPixmap, QCursor, QIcon
from PyQt5.QtCore import Qt, QSize
from .backup_annotation_pyqt5 import AnnotationTool
from .imageclassification_pyqt5 import ClassificationTool

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


def update_positions(main_window, background_label, overlay_widget, background_image):
    try:
        background_label.setGeometry(0, 0, main_window.width(), main_window.height())
        background_label.setPixmap(
            background_image.scaled(main_window.size(), Qt.KeepAspectRatioByExpanding, Qt.SmoothTransformation))
        overlay_widget.setGeometry(0, 0, main_window.width(), main_window.height())
    except Exception as e:
        print(f"Error in update_positions: {e}")


def on_button_click():
    try:
        msg = QMessageBox()
        msg.setWindowTitle("Button Clicked")
        msg.setText("You clicked the button!")
        msg.exec_()
    except Exception as e:
        print(f"Error in on_button_click: {e}")


class MainWindow(QMainWindow):
    def __init__(self):
        try:
            super().__init__()
            self.setWindowTitle('vivaa')

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

            # Set background image
            self.background_label = QLabel(central_widget)
            self.background_image = QPixmap(os.path.join(IMAGES_DIR, 'bg_image.png'))
            self.background_image_help = QPixmap(os.path.join(IMAGES_DIR, 'bg_help_image.png'))
            self.background_label.setPixmap(self.background_image)
            self.background_label.setScaledContents(True)
            self.background_label.setAlignment(Qt.AlignCenter)
            self.background_label.setGeometry(0, 0, self.width(), self.height())

            # Overlay widget for title and buttons
            self.overlay_widget = QWidget(central_widget)
            overlay_layout = QVBoxLayout()

            # Spacer to push content to bottom
            overlay_layout.addSpacerItem(QSpacerItem(20, 40, QSizePolicy.Minimum, QSizePolicy.Expanding))

            # Project title
            self.title_label = QLabel('VIVA □ ▭ ○', self.overlay_widget)
            title_font = QFont('Comic Sans MS', 50, QFont.Bold)
            self.title_label.setFont(title_font)
            self.title_label.setAlignment(Qt.AlignCenter)
            self.title_label.setStyleSheet('''
                color: white;
            ''')
            self.title_label.setCursor(QCursor(Qt.PointingHandCursor))
            self.title_label.mousePressEvent = self.open_help_window

            # Start button
            start_button = QPushButton('Hit to Start', self.overlay_widget)
            start_button.setFont(QFont('Impact', 18))
            start_button.setStyleSheet('''
                QPushButton {
                    background-color: lightgray;
                    color: Black;
                    border: 2px solid black;
                    border-radius: 15px;
                    padding: 15px 30px;
                }
                QPushButton:hover {
                    background-color: darkgray;
                }
                QPushButton:pressed {
                    background-color: darkgray;
                }
            ''')
            start_button.clicked.connect(self.open_options_window)

            # Exit button
            exit_button = QPushButton('Exit', self.overlay_widget)
            exit_button.setFont(QFont('Impact', 18))
            exit_button.setStyleSheet('''
                QPushButton {
                    background-color: lightgray;
                    color: Black;
                    border: 2px solid black;
                    border-radius: 15px;
                    padding: 15px 30px;
                }
                QPushButton:hover {
                    background-color: darkgray;
                }
                QPushButton:pressed {
                    background-color: darkgray;
                }
            ''')
            exit_button.clicked.connect(self.exit_application)

            # Add title and buttons to layout
            overlay_layout.addWidget(self.title_label, alignment=Qt.AlignCenter)
            overlay_layout.addWidget(start_button, alignment=Qt.AlignCenter)
            overlay_layout.addWidget(exit_button, alignment=Qt.AlignCenter)

            # Set layout to overlay widget
            self.overlay_widget.setLayout(overlay_layout)
            self.overlay_widget.setGeometry(0, 0, self.width(), self.height())

            self.resizeEvent = self.update_positions_event
        except Exception as e:
            print(f"Error in MainWindow: {e}")

    def update_positions_event(self, event):
        update_positions(self, self.background_label, self.overlay_widget, self.background_image)
        if hasattr(self, 'dark_overlay'):
            self.dark_overlay.setGeometry(0, 0, self.width(), self.height())

    def open_help_window(self, event):
        try:
            self.help_window = HelpWindow(self.background_image_help)
            self.help_window.showMaximized()
            self.close()
        except Exception as e:
            print(f"Error in open_help_window: {e}")

    def open_options_window(self, event=None):
        try:
            self.options_window = OptionsWindow(self.background_image_help)
            self.options_window.showMaximized()
            self.close()
        except Exception as e:
            print(f"Error in open_options_window: {e}")

    def exit_application(self):
        try:
            QApplication.instance().quit()
        except Exception as e:
            print(f"Error in exit_application: {e}")


class HelpWindow(QMainWindow):
    def __init__(self, background_image):
        try:
            super().__init__()
            self.setWindowTitle('Help')
            self.setGeometry(100, 100, 800, 600)

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

            # Set background image
            self.background_label = QLabel(central_widget)
            self.background_image = background_image
            self.background_label.setPixmap(self.background_image)
            self.background_label.setScaledContents(True)
            self.background_label.setAlignment(Qt.AlignCenter)
            self.background_label.setGeometry(0, 0, self.width(), self.height())
            self.background_label = QLabel(central_widget)
            self.background_label.setPixmap(self.background_image)
            self.background_label.setScaledContents(True)
            self.background_label.setAlignment(Qt.AlignCenter)
            self.background_label.setGeometry(0, 0, self.width(), self.height())

            # Add dark overlay for dimming effect
            self.dark_overlay = QWidget(central_widget)
            self.dark_overlay.setStyleSheet("background-color: rgba(0, 0, 0, 128);")
            self.dark_overlay.setGeometry(0, 0, self.width(), self.height())
            self.dark_overlay.lower()

            # Scroll area for help content
            scroll_area = QScrollArea(central_widget)
            scroll_area.setGeometry(0, 0, self.width(), self.height())
            scroll_area.setWidgetResizable(True)
            scroll_area.setStyleSheet("background: transparent;")

            # Overlay widget for help content
            self.overlay_widget = QWidget()
            overlay_layout = QVBoxLayout()

            # Help content
            content = QLabel(self.overlay_widget)
            content.setFont(QFont('Times New Roman', 14))
            content.setAlignment(Qt.AlignLeft)
            content.setStyleSheet('color: white; background: transparent;')
            content.setWordWrap(True)
            content.setText('''
            <div style="text-align: center; font-family: 'Times New Roman'; color: white;">

            <h1>vivaa - Help —> Data Annotation -> Object Detection Tool</h1>

            <h2>Keyboard Shortcuts</h2>
            <ul style="text-align: left; display: inline-block;">
                <li><b>D / d:</b> Go to the <u>next image</u>.</li>
                <li><b>A / a:</b> Go to the <u>previous image</u> and <u>delete</u> the last bounding box drawn on that image.</li>
                <li><b>Delete:</b> Remove the <u>selected bounding box</u> from the current image.</li>
            </ul>

            <h2>Mouse Actions</h2>
            <ul style="text-align: left; display: inline-block;">
                <li><b>Left Click + Drag:</b> Draw a new bounding box on the image.</li>
                <li><b>Click on Bounding Box:</b> Select an existing bounding box for editing or deletion.</li>
                <li><b>Resize or Adjust Bounding Box:</b> Place the cursor near the edges or corners of a drawn bounding box, click and drag to resize or reshape it, then release to automatically update the bounding box coordinates in the corresponding label file.</li>
            </ul>

            <h2>Automatic Saving</h2>
            <ul style="text-align: left; display: inline-block;">
                <li>Every time you draw a bounding box and label it with a class, a corresponding label file is <b>automatically created</b> in YOLO format.</li>
                <li>No need to manually save after each annotation — the data is stored instantly.</li>
            </ul>

            <h2>Menu / Toolbar Options - left</h2>
            <ul style="text-align: left; display: inline-block;">
                <li><b>Back:</b> Return to the previous window or main menu from the annotation tool.</li>
                <li><b>Select Folder:</b> Load a directory containing images to annotate.</li>
                <li><b>Save Dir Location:</b> Load a directory to save the label files.</li>
                <li><b>Dropdown Menu:</b> Allows the user to select the shape type for annotation. Available options include <i>Rectangle</i>, <i>Square</i>, <i>Circle</i>, and <i>Polygon</i>. After choosing a shape, click and drag on the image to draw the corresponding bounding region.</li>
                <li><b>Bounding Box Limit:</b> Define how many bounding boxes can be drawn on each image. When the specified count is reached (e.g., <i>1</i>), the tool automatically switches to the next image, ensuring consistent annotation per image.</li>
                <li><b>Previous Image:</b> Go back to the previous image.</li>
                <li><b>Next Image:</b> Move forward to the next image in the folder.</li>
                <li><b>Delete Bounding Box:</b> Removes the selected bounding box from the current image when clicked.</li>
                <li><b>Zoom In:</b> Magnify the image for detailed annotation work, improving accuracy on small or complex objects.</li>
                <li><b>Zoom Out:</b> Reduce the image scale to view the full context or multiple objects in the frame.</li>
            </ul>

            <h2>Annotation Workflow</h2>
            <ol style="text-align: left; display: inline-block;">
                <li>Click <b>Select Folder</b> and select the image dataset directory.</li>
                <li>Use the mouse to <b>draw bounding boxes</b> around desired objects.</li>
                <li>Each box will be <b>saved automatically</b> in YOLO format.</li>
                <li>Press <b>D</b> to move to the next image once annotation is complete.</li>
                <li>Press <b>A</b> to return to the previous image and delete the last bounding box drawn in previous image.</li>
                <li>Use <b>Delete</b> to remove any incorrect bounding box.</li>
                <li>Repeat until all images are annotated.</li>
            </ol>

            <h2>Smart Options</h2>
            <ol style="text-align: left; display: inline-block;">
                <li><p><b>Bounding Box Limit & Auto-Next:</b> Specify the number of bounding boxes to draw per image in the left toolbar. 
                For example, if set to <i>1</i>, the tool will automatically move to the next image once the annotator draws the specified number of boxes and assigns a label. 
                This saves time by eliminating the need to manually click “Next Image” after each annotation.</p></li>

                <li><p><b>Default Label:</b> Set a default label for the current annotation session. 
                When enabled, any bounding box drawn will automatically be assigned this label. 
                This is especially useful when annotating large datasets with a fixed number of classes (e.g., "awake" or "drowsy"), saving time by avoiding repeated manual labeling for each image.</p></li>

                <li><p><b>Smart Previous (Delete Last Box):</b> Pressing <i>A</i> goes to the previous image and automatically deletes the last bounding box drawn on that image. 
                This is useful if you notice an incorrect annotation in the previous image and want to correct it quickly. 
                If you only want to view the previous image without deleting any boxes, use the <i>Previous Image</i> button in the toolbar instead.</p></li>

                <li><p><b>Full Smart Annotation Workflow:</b> Combine all smart options to speed up annotation for large datasets. 
                For example, if you have 1,000 images (500 "Awake" and 500 "Drowsy") and only one bounding box per image: 
                <ol>
                    <li>Set the number of bounding boxes per image to <i>1</i>.</li>
                    <li>Set the default label to <i>"Awake"</i> for the first 500 images.</li>
                    <li>Change the default label to <i>"Drowsy"</i> for the next 500 images.</li>
                </ol>
                With these settings, when you draw a bounding box, it is automatically labeled with the default class. 
                The tool will also automatically move to the next image once the specified number of boxes is drawn, eliminating the need to manually click “Next Image.” 
                If you notice an incorrect bounding box in the previous image, press <i>A</i> to delete the last drawn box and quickly correct it. 
                This workflow allows fluent, fast annotation without repeated manual steps.</p></li>
            </ol>

            <h2>Tips & Best Practices</h2>
            <ul style="text-align: left; display: inline-block;">
                <li><b>Precision Matters:</b> Draw bounding boxes tightly around objects to ensure high detection accuracy and reduce model confusion.</li>
                <li><b>Consistency is Key:</b> Use uniform labeling conventions across your dataset to maintain model reliability and prevent class ambiguity.</li>
                <li><b>Clear Class Definitions:</b> Choose unambiguous, descriptive class names and periodically validate label files to catch errors early.</li>
                <li><b>Zoom for Accuracy:</b> Utilize zoom-in features for small or overlapping objects to annotate with pixel-level precision.</li>
                <li><b>Review and Refine:</b> Regularly revisit previous annotations to detect inconsistencies, correct mistakes, and maintain dataset quality.</li>
                <li><b>Leverage Defaults & Limits:</b> Use default labels and bounding box limits to accelerate workflow and minimize repetitive manual steps.</li>
                <li><b>Maintain Annotation Hygiene:</b> Delete incorrect boxes immediately, and avoid overlapping or duplicate annotations to keep the dataset clean.</li>
                <li><b>Version Control:</b> Backup label files periodically to prevent data loss and track annotation progress systematically.</li>
            </ul>

            </div>
            ''')

            # Back button
            back_button = QPushButton('Back', self.overlay_widget)
            back_button.setFont(QFont('Impact', 18))
            back_button.setStyleSheet('''
                QPushButton {
                    background-color: lightgray;
                    color: Black;
                    border: 2px solid black;
                    border-radius: 15px;
                    padding: 15px 30px;
                }
                QPushButton:hover {
                    background-color: darkgray;
                }
                QPushButton:pressed {
                    background-color: darkgray;
                }
            ''')
            back_button.clicked.connect(self.open_main_window)

            overlay_layout.addWidget(content)
            overlay_layout.addWidget(back_button, alignment=Qt.AlignCenter)
            self.overlay_widget.setLayout(overlay_layout)
            self.overlay_widget.setStyleSheet('background: transparent;')

            scroll_area.setWidget(self.overlay_widget)
            layout = QVBoxLayout(central_widget)
            layout.addWidget(scroll_area)

            self.resizeEvent = self.update_positions_event
        except Exception as e:
            print(f"Error in HelpWindow: {e}")

    def update_positions_event(self, event):
        update_positions(self, self.background_label, self.overlay_widget, self.background_image)
        if hasattr(self, 'dark_overlay'):
            self.dark_overlay.setGeometry(0, 0, self.width(), self.height())

    def open_main_window(self):
        try:
            self.main_window = MainWindow()
            self.main_window.showMaximized()
            self.close()
        except Exception as e:
            print(f"Error in open_main_window: {e}")


class OptionsWindow(QMainWindow):
    def __init__(self, background_image):
        try:
            super().__init__()
            self.setWindowTitle('Options')
            self.setGeometry(100, 100, 800, 600)

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

            # Set background image
            self.background_label = QLabel(central_widget)
            self.background_image = background_image
            self.background_label.setPixmap(self.background_image)
            self.background_label.setScaledContents(True)
            self.background_label.setAlignment(Qt.AlignCenter)
            self.background_label.setGeometry(0, 0, self.width(), self.height())
            self.background_label = QLabel(central_widget)
            self.background_label.setPixmap(self.background_image)
            self.background_label.setScaledContents(True)
            self.background_label.setAlignment(Qt.AlignCenter)
            self.background_label.setGeometry(0, 0, self.width(), self.height())

            # Add dark overlay for dimming effect
            self.dark_overlay = QWidget(central_widget)
            self.dark_overlay.setStyleSheet("background-color: rgba(0, 0, 0, 128);")
            self.dark_overlay.setGeometry(0, 0, self.width(), self.height())
            self.dark_overlay.lower()

            # Overlay widget for buttons
            self.overlay_widget = QWidget(central_widget)
            overlay_layout = QVBoxLayout()

            # Spacer to center content
            overlay_layout.addSpacerItem(QSpacerItem(20, 40, QSizePolicy.Minimum, QSizePolicy.Expanding))

            # Option buttons
            data_annotation_icon = os.path.join(IMAGES_DIR, 'ANNOTATE_ICON.png')
            data_augmentation_icon = os.path.join(IMAGES_DIR, 'Image_Transformation_icon.png')

            data_annotation_button = QPushButton()
            data_annotation_button.setIcon(QIcon(data_annotation_icon))
            data_annotation_button.setIconSize(QSize(300, 300))
            data_annotation_button.setStyleSheet('''
                QPushButton {
                    background-color: #9370DB;
                    color: white;
                    border: 2px solid #8A2BE2;
                    border-radius: 15px;
                    padding: 10px;
                    font: bold 16px "Arial";
                }
                QPushButton:hover {
                    background-color: #EE82EE;
                    border: 2px solid #8A2BE2;
                }
                QPushButton:pressed {
                    background-color: #388E3C;
                    border: 2px solid #2E7D32;
                }
            ''')
            data_annotation_button.clicked.connect(self.open_data_annotation)

            data_annotation_label = QLabel("Data Annotation")
            data_annotation_label.setFont(QFont('Arial', 18, QFont.Bold))
            data_annotation_label.setAlignment(Qt.AlignCenter)
            data_annotation_label.setStyleSheet('''
                QLabel {
                    font-family: Verdana;
                    font-size: 28px;
                    font-weight: bold;
                    font-style: italic;
                    color: white;
                }
            ''')

            data_augmentation_button = QPushButton()
            data_augmentation_button.setIcon(QIcon(data_augmentation_icon))
            data_augmentation_button.setIconSize(QSize(300, 300))
            data_augmentation_button.setStyleSheet('''
                QPushButton {
                    background-color: #FFA500;
                    color: white;
                    border: 2px solid #FF8C00;
                    border-radius: 15px;
                    padding: 10px;
                    font: bold 16px "Arial";
                }
                QPushButton:hover {
                    background-color: #FFD700;
                    border: 2px solid #FFA500;
                }
                QPushButton:pressed {
                    background-color: #1976D2;
                    border: 2px solid #1565C0;
                }
            ''')
            data_augmentation_button.clicked.connect(self.open_data_augmentation)

            data_augmentation_label = QLabel("Data Augmentation")
            data_augmentation_label.setFont(QFont('Arial', 18, QFont.Bold))
            data_augmentation_label.setAlignment(Qt.AlignCenter)
            data_augmentation_label.setStyleSheet('''
                QLabel {
                    font-family: Verdana;
                    font-size: 28px;
                    font-weight: bold;
                    font-style: italic;
                    color: white;
                }
            ''')

            # Layout for buttons
            buttons_layout = QHBoxLayout()
            data_annotation_layout = QVBoxLayout()
            data_annotation_layout.addWidget(data_annotation_button, alignment=Qt.AlignCenter)
            data_annotation_layout.addWidget(data_annotation_label, alignment=Qt.AlignCenter)
            data_augmentation_layout = QVBoxLayout()
            data_augmentation_layout.addWidget(data_augmentation_button, alignment=Qt.AlignCenter)
            data_augmentation_layout.addWidget(data_augmentation_label, alignment=Qt.AlignCenter)
            buttons_layout.addLayout(data_annotation_layout)
            buttons_layout.addLayout(data_augmentation_layout)

            overlay_layout.addLayout(buttons_layout)

            # Spacer to center content
            overlay_layout.addSpacerItem(QSpacerItem(20, 40, QSizePolicy.Minimum, QSizePolicy.Expanding))

            # Back button
            back_button = QPushButton('Back')
            back_button.setFont(QFont('Impact', 18))
            back_button.setStyleSheet('''
                QPushButton {
                    background-color: lightgray;
                    color: Black;
                    border: 2px solid black;
                    border-radius: 15px;
                    padding: 15px 30px;
                }
                QPushButton:hover {
                    background-color: darkgray;
                }
                QPushButton:pressed {
                    background-color: darkgray;
                }
            ''')
            back_button.clicked.connect(self.open_main_window)

            overlay_layout.addWidget(back_button, alignment=Qt.AlignCenter)

            # Set layout to overlay widget
            self.overlay_widget.setLayout(overlay_layout)
            self.overlay_widget.setGeometry(0, 0, self.width(), self.height())

            self.resizeEvent = self.update_positions_event
        except Exception as e:
            print(f"Error in OptionsWindow: {e}")

    def update_positions_event(self, event):
        update_positions(self, self.background_label, self.overlay_widget, self.background_image)
        if hasattr(self, 'dark_overlay'):
            self.dark_overlay.setGeometry(0, 0, self.width(), self.height())

    def open_main_window(self):
        try:
            self.main_window = MainWindow()
            self.main_window.showMaximized()
            self.close()
        except Exception as e:
            print(f"Error in open_main_window: {e}")

    def open_data_annotation(self):
        try:
            self.data_annotation_window = DataAnnotationWindow(self.background_image)
            self.data_annotation_window.showMaximized()
            self.close()
        except Exception as e:
            print(f"Error in open_data_annotation: {e}")

    def open_data_augmentation(self):
        try:
            QMessageBox.information(self, "Data Augmentation", "Data Augmentation clicked!")
        except Exception as e:
            print(f"Error in open_data_augmentation: {e}")


class DataAnnotationWindow(QMainWindow):
    def __init__(self, background_image):
        try:
            super().__init__()
            self.setWindowTitle('Data Annotation Options')
            self.setGeometry(100, 100, 800, 600)

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

            # Set background image
            self.background_label = QLabel(central_widget)
            self.background_image = background_image
            self.background_label.setPixmap(self.background_image)
            self.background_label.setScaledContents(True)
            self.background_label.setAlignment(Qt.AlignCenter)
            self.background_label.setGeometry(0, 0, self.width(), self.height())
            self.background_label = QLabel(central_widget)
            self.background_label.setPixmap(self.background_image)
            self.background_label.setScaledContents(True)
            self.background_label.setAlignment(Qt.AlignCenter)
            self.background_label.setGeometry(0, 0, self.width(), self.height())

            # Add dark overlay for dimming effect
            self.dark_overlay = QWidget(central_widget)
            self.dark_overlay.setStyleSheet("background-color: rgba(0, 0, 0, 128);")
            self.dark_overlay.setGeometry(0, 0, self.width(), self.height())
            self.dark_overlay.lower()

            # Overlay widget for buttons
            self.overlay_widget = QWidget(central_widget)
            overlay_layout = QVBoxLayout()

            # Spacer to center content
            overlay_layout.addSpacerItem(QSpacerItem(20, 40, QSizePolicy.Minimum, QSizePolicy.Expanding))

            # Buttons for annotation types
            annotation_options = [
                (os.path.join(IMAGES_DIR, 'object_detection_icon.png'), "Object Detection", "#9370DB", "#8A2BE2",
                 "#388E3C", "#2E7D32"),
                (
                    os.path.join(IMAGES_DIR, 'image_classification_icon.png'), "Image Classification", "#FFA500",
                    "#FF8C00",
                    "#1976D2", "#1565C0"),
                (os.path.join(IMAGES_DIR, 'instance_segmentation_icon.png'), "Instance Segmentation", "#6A5ACD",
                 "#483D8B", "#4CAF50", "#388E3C"),
                (os.path.join(IMAGES_DIR, 'keypoint_detection_icon.png'), "Keypoint Detection", "#20B2AA", "#008B8B",
                 "#2196F3", "#1976D2"),
            ]

            buttons_layout = QHBoxLayout()
            for i in range(0, len(annotation_options), 2):  # Two columns
                column_layout = QVBoxLayout()
                for icon_path, label_text, bg_color, border_color, pressed_bg, pressed_border in annotation_options[
                                                                                                 i:i + 2]:
                    button = QPushButton()
                    button.setIcon(QIcon(icon_path))
                    button.setIconSize(QSize(200, 200))
                    button.setStyleSheet(f'''
                        QPushButton {{
                            background-color: {bg_color};
                            color: white;
                            border: 2px solid {border_color};
                            border-radius: 15px;
                            padding: 10px;
                            font: bold 14px "Arial";
                        }}
                        QPushButton:hover {{
                            background-color: {pressed_bg};
                            border: 2px solid {pressed_border};
                        }}
                        QPushButton:pressed {{
                            background-color: {pressed_bg};
                            border: 2px solid {pressed_border};
                        }}
                    ''')
                    if label_text == "Object Detection":
                        button.clicked.connect(self.open_object_detection_tool)
                    elif label_text == "Image Classification":
                        button.clicked.connect(self.open_classification_tool)
                    elif label_text == "Instance Segmentation":
                        button.clicked.connect(self.open_instance_segmentation_tool)
                    else:
                        button.clicked.connect(lambda checked, lt=label_text: self.on_annotation_selected(lt))
                    label = QLabel(label_text)
                    label.setFont(QFont('Arial', 16, QFont.Bold))
                    label.setAlignment(Qt.AlignCenter)
                    label.setStyleSheet('''
                        QLabel {
                            font-family: Verdana;
                            font-size: 20px;
                            font-weight: bold;
                            font-style: italic;
                            color: white;
                        }
                    ''')
                    column_layout.addWidget(button, alignment=Qt.AlignCenter)
                    column_layout.addWidget(label, alignment=Qt.AlignCenter)
                buttons_layout.addLayout(column_layout)

            overlay_layout.addLayout(buttons_layout)

            # Spacer to center content
            overlay_layout.addSpacerItem(QSpacerItem(20, 40, QSizePolicy.Minimum, QSizePolicy.Expanding))

            # Back button
            back_button = QPushButton('Back')
            back_button.setFont(QFont('Impact', 18))
            back_button.setStyleSheet('''
                QPushButton {
                    background-color: lightgray;
                    color: Black;
                    border: 2px solid black;
                    border-radius: 15px;
                    padding: 15px 30px;
                }
                QPushButton:hover {
                    background-color: darkgray;
                }
                QPushButton:pressed {
                    background-color: darkgray;
                }
            ''')
            back_button.clicked.connect(self.open_options_window)

            overlay_layout.addWidget(back_button, alignment=Qt.AlignCenter)

            # Set layout to overlay widget
            self.overlay_widget.setLayout(overlay_layout)
            self.overlay_widget.setGeometry(0, 0, self.width(), self.height())

            self.resizeEvent = self.update_positions_event
        except Exception as e:
            print(f"Error in DataAnnotationWindow: {e}")

    def update_positions_event(self, event):
        update_positions(self, self.background_label, self.overlay_widget, self.background_image)
        if hasattr(self, 'dark_overlay'):
            self.dark_overlay.setGeometry(0, 0, self.width(), self.height())

    def open_options_window(self):
        try:
            self.options_window = OptionsWindow(self.background_image)
            self.options_window.showMaximized()
            self.close()
        except Exception as e:
            print(f"Error in open_options_window: {e}")

    def on_annotation_selected(self, option):
        try:
            if option == "Object Detection":
                self.annotation_tool = AnnotationTool()
                self.annotation_tool.show()
                self.hide()
            elif option == "Image Classification":
                self.classification_tool = ClassificationTool()
                self.classification_tool.show()
                self.hide()
            else:
                QMessageBox.information(self, "Info", f"{option} is not implemented yet.")
        except Exception as e:
            QMessageBox.critical(self, "Error", f"Failed to open {option} tool: {str(e)}")

    def open_object_detection_tool(self):
        try:
            self.annotation_tool = AnnotationTool()
            self.annotation_tool.showMaximized()
            self.close()
        except Exception as e:
            print(f"Error in open_object_detection_tool: {e}")

    def open_classification_tool(self):
        try:
            self.classification_tool = ClassificationTool()
            self.classification_tool.showMaximized()
            self.close()
        except Exception as e:
            print(f"Error in open_classification_tool: {e}")

    def open_instance_segmentation_tool(self):
        try:
            # TODO: Implement the actual InstanceSegmentationTool
            # Placeholder message for now
            QMessageBox.information(
                self,
                "Instance Segmentation Tool",
                "This feature is under development and will be implemented soon.",
            )
            # If you had the real tool, you would uncomment these lines:
            # from .instancesegmentation_pyqt5 import InstanceSegmentationTool
            # self.instance_segmentation_tool = InstanceSegmentationTool()
            # self.instance_segmentation_tool.showMaximized()
            # self.close()
        except Exception as e:
            print(f"Error in open_instance_segmentation_tool: {e}")


def main():
    try:
        app = QApplication([])
        main_window = MainWindow()
        main_window.showMaximized()
        app.exec_()
    except Exception as e:
        print(f"Error in main: {e}")
