import warnings

import copy

import numpy as np

from PyQt5.QtCore import Qt, QPointF, QRectF
from PyQt5.QtGui import QMouseEvent, QKeyEvent, QPen, QColor, QBrush
from PyQt5.QtWidgets import QMessageBox, QGraphicsRectItem, QApplication

from coralnet_toolbox.Tools.QtTool import Tool

from coralnet_toolbox.Results import ResultsProcessor
from coralnet_toolbox.Results import CombineResults
from coralnet_toolbox.Results import MapResults

from coralnet_toolbox.Annotations.QtRectangleAnnotation import RectangleAnnotation
from coralnet_toolbox.Annotations.QtPolygonAnnotation import PolygonAnnotation

from coralnet_toolbox.QtProgressBar import ProgressBar
from coralnet_toolbox.QtWorkArea import WorkArea

from coralnet_toolbox.utilities import pixmap_to_numpy
from coralnet_toolbox.utilities import simplify_polygon

warnings.filterwarnings("ignore", category=DeprecationWarning)


# ----------------------------------------------------------------------------------------------------------------------
# Classes
# ----------------------------------------------------------------------------------------------------------------------


class SeeAnythingTool(Tool):
    def __init__(self, annotation_window):
        super().__init__(annotation_window)
        self.annotation_window = annotation_window
        self.main_window = annotation_window.main_window
        self.see_anything_dialog = None
        
        self.animation_manager = self.annotation_window.animation_manager

        self.top_left = None

        self.cursor = Qt.CrossCursor
        self.default_cursor = Qt.ArrowCursor  
        self.annotation_graphics = None

        self.work_area_image = None
        self.rectangles = []       
        self.rectangle_items = []  

        self.working_area = None

        self.image_path = None
        self.original_image = None
        self.original_width = None
        self.original_height = None

        # Rectangle drawing attributes
        self.start_point = None
        self.end_point = None
        self.top_left = None
        self.bottom_right = None
        self.drawing_rectangle = False
        self.current_rect_graphics = None  
        self.rectangles_processed = False  

        # Add state variables for custom working area creation
        self.creating_working_area = False
        self.working_area_start = None
        self.working_area_temp_graphics = None
        
        # Add hover position tracking
        self.hover_pos = None
        
        self.annotations = []
        self.results = None
    
    def activate(self):
        """
        Activates the tool.
        """
        self.active = True
        self.annotation_window.setCursor(self.cursor)
        self.see_anything_dialog = self.main_window.see_anything_deploy_predictor_dialog

    def deactivate(self):
        """
        Deactivates the tool and cleans up all resources.
        """
        self.active = False
        self.annotation_window.setCursor(self.default_cursor)

        # Clear annotations that haven't been confirmed
        if self.annotations:
            self.clear_annotations()

        # Clear rectangle data and graphics
        self.clear_all_rectangles()
        # Clean up working area and shadow
        self.cancel_working_area()
        # Cancel working area creation if in progress
        self.cancel_working_area_creation()
        # Clear detection data
        self.results = None

        # Update the viewport
        self.annotation_window.scene.update()
        
        # Call parent deactivate to ensure crosshair is properly cleared
        super().deactivate()

    def set_working_area(self):
        """
        Set the working area for the tool using the WorkArea class.
        """
        self.annotation_window.setCursor(Qt.WaitCursor)

        # Cancel the current working area if it exists
        self.cancel_working_area()

        # Original image (grab current from the annotation window)
        self.image_path = self.annotation_window.current_image_path
        self.original_image = pixmap_to_numpy(self.annotation_window.pixmap_image)
        self.original_width = self.annotation_window.pixmap_image.size().width()
        self.original_height = self.annotation_window.pixmap_image.size().height()

        # Current extent (view)
        extent = self.annotation_window.viewportToScene()

        top = max(0, round(extent.top()))
        left = max(0, round(extent.left()))
        width = round(extent.width())
        height = round(extent.height())
        bottom = min(self.original_height, top + height)
        right = min(self.original_width, left + width)

        # Create the WorkArea instance
        self.working_area = WorkArea(left, top, right - left, bottom - top, self.image_path)
        # Set animation manager
        self.working_area.set_animation_manager(self.animation_manager)

        # Create and add the working area graphics
        self.working_area.create_graphics(self.annotation_window.scene, include_shadow=True)
        self.working_area.set_remove_button_visibility(False)
        self.working_area.removed.connect(self.on_working_area_removed)

        # Crop the image based on the working area
        self.work_area_image = self.original_image[top:bottom, left:right]

        # Set the image in the SeeAnything dialog
        self.see_anything_dialog.set_image(self.work_area_image, self.image_path)
        # self.see_anything_dialog.reload_model()

        self.annotation_window.setCursor(Qt.CrossCursor)
        self.annotation_window.scene.update()
        
    def set_custom_working_area(self, start_point, end_point):
        """
        Create a working area from custom points selected by the user.
        
        Args:
            start_point (QPointF): First corner of the working area
            end_point (QPointF): Opposite corner of the working area
        """
        self.annotation_window.setCursor(Qt.WaitCursor)
        
        # Cancel any existing working area
        self.cancel_working_area()
        
        # Calculate the rectangle bounds
        left = max(0, int(min(start_point.x(), end_point.x())))
        top = max(0, int(min(start_point.y(), end_point.y())))
        right = min(int(self.annotation_window.pixmap_image.size().width()), 
                    int(max(start_point.x(), end_point.x())))
        bottom = min(int(self.annotation_window.pixmap_image.size().height()),
                     int(max(start_point.y(), end_point.y())))
        
        # Ensure minimum size (at least 10x10 pixels)
        if right - left < 10:
            right = min(left + 10, int(self.annotation_window.pixmap_image.size().width()))
        if bottom - top < 10:
            bottom = min(top + 10, int(self.annotation_window.pixmap_image.size().height()))
            
        # Original image information
        self.image_path = self.annotation_window.current_image_path
        self.original_image = pixmap_to_numpy(self.annotation_window.pixmap_image)
        self.original_width = self.annotation_window.pixmap_image.size().width()
        self.original_height = self.annotation_window.pixmap_image.size().height()
            
        # Create the WorkArea instance
        self.working_area = WorkArea(left, top, right - left, bottom - top, self.image_path)
        # Set animation manager
        self.working_area.set_animation_manager(self.animation_manager)
        
        # Create and add the working area graphics
        self.working_area.create_graphics(self.annotation_window.scene, include_shadow=True)
        self.working_area.set_remove_button_visibility(False)
        self.working_area.removed.connect(self.on_working_area_removed)
        
        # Crop the image based on the working area
        self.work_area_image = self.original_image[top:bottom, left:right]
        
        # Set the image in the SeeAnything dialog
        self.see_anything_dialog.set_image(self.work_area_image, self.image_path)
        
        self.annotation_window.setCursor(Qt.CrossCursor)
        self.annotation_window.scene.update()

    def display_working_area_preview(self, current_pos):
        """
        Display a preview rectangle for the working area being created.
        
        Args:
            current_pos (QPointF): Current mouse position
        """
        if not self.working_area_start:
            return
            
        # Remove previous preview if it exists
        if self.working_area_temp_graphics:
            self.annotation_window.scene.removeItem(self.working_area_temp_graphics)
            self.working_area_temp_graphics = None
            
        # Create preview rectangle
        rect = QRectF(
            min(self.working_area_start.x(), current_pos.x()),
            min(self.working_area_start.y(), current_pos.y()),
            abs(current_pos.x() - self.working_area_start.x()),
            abs(current_pos.y() - self.working_area_start.y())
        )
        
        # Create a dashed blue pen for the working area preview
        pen = QPen(QColor(0, 168, 230))
        pen.setCosmetic(True)
        pen.setStyle(Qt.DashLine)
        pen.setWidth(3)
        
        self.working_area_temp_graphics = QGraphicsRectItem(rect)
        self.working_area_temp_graphics.setPen(pen)
        self.working_area_temp_graphics.setBrush(QBrush(QColor(0, 168, 230, 30)))  # Light blue transparent fill
        self.annotation_window.scene.addItem(self.working_area_temp_graphics)

    def cancel_working_area_creation(self):
        """
        Cancel the process of creating a working area.
        """
        self.creating_working_area = False
        self.working_area_start = None
        
        if self.working_area_temp_graphics:
            self.annotation_window.scene.removeItem(self.working_area_temp_graphics)
            self.working_area_temp_graphics = None
            
        self.annotation_window.scene.update()

    def on_working_area_removed(self, work_area):
        """
        Handle when the work area is removed via its internal mechanism.
        """
        self.cancel_working_area()

    def create_rectangle_graphics(self):
        """
        Create a new rectangle graphics item for drawing with the selected label color.
        """
        if self.start_point and self.end_point:
            # Calculate the rectangle dimensions
            rect = QRectF(
                min(self.start_point.x(), self.end_point.x()),
                min(self.start_point.y(), self.end_point.y()),
                abs(self.end_point.x() - self.start_point.x()),
                abs(self.end_point.y() - self.start_point.y())
            )

            # Remove current rectangle being drawn if it exists
            if self.current_rect_graphics:
                if self.current_rect_graphics in self.annotation_window.scene.items():
                    self.annotation_window.scene.removeItem(self.current_rect_graphics)
                self.current_rect_graphics = None

            # Create a new rectangle graphics item
            self.current_rect_graphics = QGraphicsRectItem(rect)

            # Get color from the selected label
            color = self.annotation_window.selected_label.color

            # Style the rectangle
            pen = QPen(QColor(color))
            pen.setCosmetic(True)
            pen.setWidth(4)
            pen.setStyle(Qt.DashLine)
            self.current_rect_graphics.setPen(pen)

            # Add to scene
            self.annotation_window.scene.addItem(self.current_rect_graphics)

    def update_rectangle_graphics(self):
        """
        Update the current rectangle graphics item while drawing.
        """
        if self.start_point and self.end_point and self.drawing_rectangle:
            # If no graphics item exists yet, create one
            if not self.current_rect_graphics:
                self.create_rectangle_graphics()
            else:
                # Update the existing rectangle
                rect = QRectF(
                    min(self.start_point.x(), self.end_point.x()),
                    min(self.start_point.y(), self.end_point.y()),
                    abs(self.end_point.x() - self.start_point.x()),
                    abs(self.end_point.y() - self.start_point.y())
                )
                self.current_rect_graphics.setRect(rect)

    def add_completed_rectangle(self):
        """
        Add the completed rectangle to the list of rectangles and their graphics.
        """
        if self.current_rect_graphics:
            # Add the current rectangle graphics item to the list if it's in the scene
            if self.current_rect_graphics in self.annotation_window.scene.items():
                self.rectangle_items.append(self.current_rect_graphics)
            else:
                # If it's not in the scene anymore for some reason, don't track it
                self.current_rect_graphics = None
                return

            # Calculate rectangle coordinates relative to working area
            working_area_top_left = self.working_area.rect.topLeft()

            top_left = QPointF(
                min(self.start_point.x(), self.end_point.x()) - working_area_top_left.x(),
                min(self.start_point.y(), self.end_point.y()) - working_area_top_left.y()
            )

            bottom_right = QPointF(
                max(self.start_point.x(), self.end_point.x()) - working_area_top_left.x(),
                max(self.start_point.y(), self.end_point.y()) - working_area_top_left.y()
            )

            # Add the rectangle coordinates to the list
            rectangle = np.array([top_left.x(), top_left.y(), bottom_right.x(), bottom_right.y()])
            self.rectangles.append(rectangle)

            # Reset the current rectangle graphics item without removing from scene
            # It's now tracked in rectangle_items
            self.current_rect_graphics = None

            # Set rectangles_processed to False since we have new user rectangles
            self.rectangles_processed = False  # Indicate prediction is needed

    def mousePressEvent(self, event: QMouseEvent):
        """
        Handles the mouse press event.

        Args:
            event (QMouseEvent): The mouse press event.
        """
        if not self.annotation_window.selected_label:
            QMessageBox.warning(self.annotation_window,
                                "No Label Selected",
                                "A label must be selected before adding an annotation.")
            return None

        # Get position in scene coordinates
        scene_pos = self.annotation_window.mapToScene(event.pos())
        
        # Handle working area creation mode
        if not self.working_area and event.button() == Qt.LeftButton:
            if not self.creating_working_area:
                # Start working area creation
                self.creating_working_area = True
                self.working_area_start = scene_pos
                return
            elif self.creating_working_area and self.working_area_start:
                # Finish working area creation
                self.set_custom_working_area(self.working_area_start, scene_pos)
                self.cancel_working_area_creation()
                return

        if not self.working_area:
            return

        # Check if the position is within the working area
        if not self.working_area.contains_point(scene_pos):
            return

        if event.button() == Qt.LeftButton and not self.drawing_rectangle:
            # Get the start point
            self.start_point = scene_pos
            # Start drawing the rectangle
            self.drawing_rectangle = True
            self.end_point = self.start_point  # Initialize end_point
            self.update_rectangle_graphics()

        elif event.button() == Qt.LeftButton and self.drawing_rectangle:
            # Get the end point
            self.end_point = scene_pos
            # Finish drawing the rectangle
            self.drawing_rectangle = False
            # Update the rectangle graphics before finalizing
            self.update_rectangle_graphics()

            # Add the completed rectangle to our lists
            self.add_completed_rectangle()

            # Reset drawing state
            self.start_point = None
            self.end_point = None

        self.annotation_window.scene.update()

    def mouseMoveEvent(self, event: QMouseEvent):
        """
        Handles the mouse move event.

        Args:
            event (QMouseEvent): The mouse move event.
        """
        # Call parent implementation to handle crosshair
        super().mouseMoveEvent(event)
        
        # Continue with tool-specific behavior
        scene_pos = self.annotation_window.mapToScene(event.pos())
        self.hover_pos = scene_pos
        
        # Update working area preview during creation
        if self.creating_working_area and self.working_area_start:
            self.display_working_area_preview(scene_pos)
            return
            
        if self.working_area and self.drawing_rectangle:
            # Update the end point while drawing the rectangle
            self.end_point = self.annotation_window.mapToScene(event.pos())
            self.update_rectangle_graphics()

        self.annotation_window.scene.update()

    def keyPressEvent(self, event: QKeyEvent):
        """
        Handles the key press event.

        Args:
            event (QKeyEvent): The key press event
        """
        if event.key() == Qt.Key_Space:
            # If creating working area, confirm it
            if self.creating_working_area and self.working_area_start and self.hover_pos:
                self.set_custom_working_area(self.working_area_start, self.hover_pos)
                self.cancel_working_area_creation()
                return
        
            # If there is no working area, set it
            if not self.working_area:
                self.set_working_area()

            # If there are user-drawn rectangles ready for processing, run the predictor
            elif len(self.rectangles) > 0 and not self.rectangles_processed:
                # Create annotation based on the user-drawn rectangles
                self.create_annotations_from_rectangles()
                # Clear the user-drawn rectangles (graphics and data) as they've been used
                self.clear_all_rectangles()
                # Mark rectangles as processed for this cycle
                self.rectangles_processed = True
            else:
                # If there's a working area but no new user rectangles,
                # or if rectangles have been processed, confirm the accumulated annotations.
                if self.annotations:  # Check if there are any annotations to confirm/process
                    if self.see_anything_dialog.use_sam_dropdown.currentText() == "True":
                        self.apply_sam_model()
                    else:
                        # Confirm the annotations accumulated so far
                        self.confirm_annotations()
                # Cancel the working area if no annotations were generated or after confirmation/SAM
                self.cancel_working_area()

        elif event.key() == Qt.Key_Backspace:
            # If creating working area, cancel it
            if self.creating_working_area:
                self.cancel_working_area_creation()
                return
                
            # Cancel current rectangle being drawn
            if self.drawing_rectangle:
                self.drawing_rectangle = False
                if self.current_rect_graphics:
                    self.annotation_window.scene.removeItem(self.current_rect_graphics)
                    self.current_rect_graphics = None
                self.start_point = None
                self.end_point = None
            # If we have a working area and accumulated annotations, clear them
            elif self.working_area and len(self.annotations) > 0:
                self.clear_annotations()  # Clears unconfirmed annotations
            # If not drawing and no annotations to clear, clear any pending user-drawn rectangles
            else:
                self.clear_all_rectangles()  # Clears user input rectangles

        self.annotation_window.scene.update()

    def create_annotations_from_rectangles(self):
        """
        Create annotations based on the user-drawn rectangles.
        """
        if not self.annotation_window.active_image:
            return None

        if not self.annotation_window.pixmap_image:
            return None

        if not self.working_area:
            return None

        if len(self.rectangles) == 0:  # Check specifically for user-drawn rectangles
            return None

        # Make cursor busy
        QApplication.setOverrideCursor(Qt.WaitCursor)

        # Move the points back to the original image space
        working_area_top_left = self.working_area.rect.topLeft()

        masks = None
        # Create masks from the rectangles (these are not polygons)
        if self.see_anything_dialog.task_dropdown.currentText() == 'segment':
            masks = []
            for r in self.rectangles:
                x1, y1, x2, y2 = r
                masks.append(np.array([[x1, y1], [x2, y1], [x2, y2], [x1, y2]]))

        # Predict from prompts, providing masks if the task is segmentation
        results = self.see_anything_dialog.predict_from_prompts(self.rectangles, masks=masks)[0]

        if not results:
            # Make cursor normal
            QApplication.restoreOverrideCursor()
            return None

        # Create a results processor to merge and filter results
        results_processor = ResultsProcessor(self.main_window, {})
        # Merge
        if self.results:
            results = CombineResults().combine_results([self.results, results])

        # Filter
        self.results = results_processor.apply_filters_to_results(results)

        # Calculate the area of the image
        image_area = self.work_area_image.shape[0] * self.work_area_image.shape[1]

        # Clear previous annotations if any
        self.clear_annotations()

        # Process results based on the task type (creates polygons or rectangle annotations)
        if self.see_anything_dialog.task_dropdown.currentText() == "segment":
            if self.results.masks:
                for i, polygon in enumerate(self.results.masks.xyn):
                    confidence = self.results.boxes.conf[i].item()
                    if confidence < self.main_window.get_uncertainty_thresh():
                        continue

                    # Get absolute bounding box for area check (relative to work area)
                    box_work_area = self.results.boxes.xyxy[i].detach().cpu().numpy()
                    box_area = (box_work_area[2] - box_work_area[0]) * (box_work_area[3] - box_work_area[1])

                    # Area filtering
                    min_area = self.main_window.get_area_thresh_min() * image_area
                    max_area = self.main_window.get_area_thresh_max() * image_area
                    if not (min_area <= box_area <= max_area):
                        continue

                    # Convert normalized polygon points to absolute coordinates in the whole image
                    polygon_abs = polygon.copy()
                    polygon_abs[:, 0] = polygon_abs[:, 0] * self.work_area_image.shape[1] + working_area_top_left.x()
                    polygon_abs[:, 1] = polygon_abs[:, 1] * self.work_area_image.shape[0] + working_area_top_left.y()

                    polygon_abs = simplify_polygon(polygon_abs, 0.1)
                    self.create_polygon_annotation(polygon_abs, confidence)
                    
        else:  # Task is 'detect'
            if self.results.boxes:
                for i, box_norm in enumerate(self.results.boxes.xyxyn):
                    confidence = self.results.boxes.conf[i].item()
                    if confidence < self.main_window.get_uncertainty_thresh():
                        continue

                    # Convert normalized box to absolute coordinates in the work area
                    box_abs_work_area = box_norm.detach().cpu().numpy() * np.array(
                        [self.work_area_image.shape[1], self.work_area_image.shape[0],
                         self.work_area_image.shape[1], self.work_area_image.shape[0]])
                    # Calculate the area of the bounding box
                    box_area = (box_abs_work_area[2] - box_abs_work_area[0]) * \
                               (box_abs_work_area[3] - box_abs_work_area[1])

                    # Area filtering
                    min_area = self.main_window.get_area_thresh_min() * image_area
                    max_area = self.main_window.get_area_thresh_max() * image_area
                    if not (min_area <= box_area <= max_area):
                        continue

                    # Add working area offset to get coordinates in the whole image
                    box_abs_full = box_abs_work_area.copy()
                    box_abs_full[0] += working_area_top_left.x()
                    box_abs_full[1] += working_area_top_left.y()
                    box_abs_full[2] += working_area_top_left.x()
                    box_abs_full[3] += working_area_top_left.y()
                    self.create_rectangle_annotation(box_abs_full, confidence)

        self.annotation_window.scene.update()

        # Make cursor normal
        QApplication.restoreOverrideCursor()

    def create_rectangle_annotation(self, box, confidence):
        """
        Create rectangle annotations based on the given box coordinates.

        Args:
            box (np.ndarray): The bounding box coordinates.
            confidence (float): The confidence score for the annotation.
        """
        if len(box):
            # Convert to QPointF
            top_left = QPointF(box[0], box[1])
            bottom_right = QPointF(box[2], box[3])

            # Create the annotation
            annotation = RectangleAnnotation(top_left,
                                             bottom_right,
                                             self.annotation_window.selected_label.short_label_code,
                                             self.annotation_window.selected_label.long_label_code,
                                             self.annotation_window.selected_label.color,
                                             self.annotation_window.current_image_path,
                                             self.annotation_window.selected_label.id,
                                             self.annotation_window.main_window.get_transparency_value())

            # Update the confidence score of annotation
            annotation.update_machine_confidence({self.annotation_window.selected_label: confidence})
            # Set the animation manager
            annotation.set_animation_manager(self.animation_manager)
            # Ensure the annotation is added to the scene after creation (but not saved yet)
            annotation.create_graphics_item(self.annotation_window.scene)
            # Animate the annotation
            annotation.animate(force=True)
            
            self.annotations.append(annotation)
            
    def update_transparency(self, value):
        """
        Update the transparency of all unconfirmed annotations in this tool.
        """
        for annotation in self.annotations:
            annotation.update_transparency(value)
        self.annotation_window.scene.update()

    def create_polygon_annotation(self, points, confidence):
        """
        Create polygon annotations based on the given points.

        Args:
            points (np.ndarray): The polygon points.
            confidence (float): The confidence score for the annotation.
        """
        if len(points) > 3:
            # Convert to QPointF
            points = [QPointF(point[0], point[1]) for point in points]
            # Create the annotation
            annotation = PolygonAnnotation(points,
                                           self.annotation_window.selected_label.short_label_code,
                                           self.annotation_window.selected_label.long_label_code,
                                           self.annotation_window.selected_label.color,
                                           self.annotation_window.current_image_path,
                                           self.annotation_window.selected_label.id,
                                           self.annotation_window.main_window.get_transparency_value())

            # Update the confidence score of annotation
            annotation.update_machine_confidence({self.annotation_window.selected_label: confidence})
            # Set the animation manager
            annotation.set_animation_manager(self.animation_manager)
            # Ensure the annotation is added to the scene after creation (but not saved yet)
            annotation.create_graphics_item(self.annotation_window.scene)
            # Animate the annotation
            annotation.animate(force=True)
            
            self.annotations.append(annotation)

    def confirm_annotations(self):
        """
        Confirm the annotations and clear the working area.
        """
        # Make cursor busy
        QApplication.setOverrideCursor(Qt.WaitCursor)
        progress_bar = ProgressBar(self.annotation_window, "Confirming Annotations")
        progress_bar.show()
        progress_bar.start_progress(len(self.annotations))
            
        for annotation in self.annotations:
            # The add_annotation_from_tool will re-bind it
            # and it will animate normally if selected.
            annotation.deanimate()
            
            # Create cropped image if not already done
            if not annotation.cropped_image and self.annotation_window.rasterio_image:
                annotation.create_cropped_image(self.annotation_window.rasterio_image)
            
            # Add the annotation using the add_annotation_from_tool method
            self.annotation_window.add_annotation_from_tool(annotation)

            # Update progress bar
            progress_bar.update_progress()
            
        # Update the scene to reflect deanimation
        self.annotation_window.scene.update()

        # Make cursor normal
        QApplication.restoreOverrideCursor()
        progress_bar.finish_progress()
        progress_bar.stop_progress()
        progress_bar.close()
        progress_bar = None

        # Clear all rectangles explicitly before clearing the working area
        self.clear_all_rectangles()

        # Clear the working area
        self.cancel_working_area()

        # Clear the annotations list
        self.annotations = []
        self.results = None

    def apply_sam_model(self):
        """Uses the Results with SAM predictor to create polygons instead of confirming the
        ones created by the SeeAnything predictor."""
        # Make cursor busy
        QApplication.setOverrideCursor(Qt.WaitCursor)

        # Create a class mapping dictionary
        class_mapping = {0: self.annotation_window.selected_label}

        # Create a results processor
        results_processor = ResultsProcessor(
            self.main_window,
            class_mapping
        )

        # Make a copy of the results
        results = copy.deepcopy(self.results)

        # Update the class mapping for the results
        results.names = {0: class_mapping[0].short_label_code}

        # Process the results with the SAM predictor using the new
        results = self.see_anything_dialog.sam_dialog.predict_from_results([results], self.image_path)

        # Get SAM resizing dimensions
        original_h, original_w = self.work_area_image.shape[:2]
        resized_h, resized_w = self.see_anything_dialog.sam_dialog.resized_image.shape[:2]

        # Calculate scaling factors
        scale_x = original_w / resized_w
        scale_y = original_h / resized_h

        # Update mask coordinates to account for resizing
        for i, mask in enumerate(results[0].masks.xy):
            if len(mask) > 0:
                # Scale coordinates back to original size
                mask[:, 0] *= scale_x
                mask[:, 1] *= scale_y
                results[0].masks.xy[i] = mask

        # Get the raster
        raster = self.main_window.image_window.raster_manager.get_raster(self.image_path)

        # Map results from working area to the original image coordinates
        results = MapResults().map_results_from_work_area(
            results,
            raster,
            self.working_area,
            map_masks=True
        )

        # Process the results
        results_processor.process_segmentation_results(results)

        # Make cursor normal
        QApplication.restoreOverrideCursor()
        # Clear the previous, non-confirmed annotations
        self.clear_annotations()
        # Clear the working area
        self.cancel_working_area()

    def clear_annotations(self):
        """
        Clear all *unconfirmed* annotations created by this tool from the scene.
        """
        for annotation in self.annotations:
            # Stop animation first
            annotation.deanimate()  # Deanimate the annotation before removing
            annotation.delete()  # Let the annotation handle all graphics cleanup
            annotation = None

        self.annotations = []
        self.annotation_window.scene.update()

    def clear_rectangle_graphics(self):
        """
        Clear rectangle graphics from the scene but keep the data.
        """
        # Remove all rectangle graphics from scene
        for rect_item in self.rectangle_items:
            if rect_item in self.annotation_window.scene.items():
                # Ensure any child items are removed first (like borders or handles)
                child_items = rect_item.childItems()
                for child in child_items:
                    self.annotation_window.scene.removeItem(child)

                # Remove the rectangle item itself
                self.annotation_window.scene.removeItem(rect_item)
            rect_item = None  # Explicitly dereference

        # Clear the rectangle graphics if one is being drawn
        if self.current_rect_graphics:
            if self.current_rect_graphics in self.annotation_window.scene.items():
                # Remove any child items first
                child_items = self.current_rect_graphics.childItems()
                for child in child_items:
                    self.annotation_window.scene.removeItem(child)

                self.annotation_window.scene.removeItem(self.current_rect_graphics)
            self.current_rect_graphics = None

        # Reset the graphics list
        self.rectangle_items = []

        # Force a full scene update and repaint
        self.annotation_window.scene.update()
        self.annotation_window.viewport().update()

    def clear_rectangle_data(self):
        """
        Clear rectangle data structures but keep the graphics.
        """
        self.rectangles = []
        self.start_point = None
        self.end_point = None
        self.drawing_rectangle = False
        self.rectangles_processed = False

    def clear_all_rectangles(self):
        """
        Clear all *user-drawn* rectangle graphics and data.
        """
        self.clear_rectangle_graphics()  # Clears items in self.rectangle_items and self.current_rect_graphics
        self.clear_rectangle_data()      # Clears self.rectangles list and drawing state

    def cancel_working_area(self):
        """
        Cancel the working area and clean up all associated resources.
        """
        if self.working_area:
            # Properly remove the working area using its method
            self.working_area.remove_from_scene()
            self.working_area = None

        self.image_path = None
        self.original_image = None
        self.work_area_image = None

        # Clear all rectangles when canceling the working area
        self.clear_all_rectangles()
        self.rectangles_processed = False

        self.annotations = []
        self.results = None

        # Force update to ensure graphics are removed visually
        self.annotation_window.scene.update()
