import warnings

import base64
import rasterio

import numpy as np

from scipy.ndimage import label as ndimage_label
from skimage.measure import find_contours

from shapely.geometry import Polygon, Point

from rasterio.features import rasterize

from pycocotools import mask

from PyQt5.QtCore import Qt, QPointF, QRectF
from PyQt5.QtWidgets import QGraphicsScene, QGraphicsPixmapItem, QGraphicsItem, QApplication
from PyQt5.QtGui import QPixmap, QColor, QImage, QPainter, QBrush, QPolygonF

from coralnet_toolbox.Annotations.QtAnnotation import Annotation
from coralnet_toolbox.Annotations.QtPolygonAnnotation import PolygonAnnotation

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


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


class MaskGraphicsItem(QGraphicsItem):
    def __init__(self, mask_annotation):
        super().__init__()
        self.mask_annotation = mask_annotation
        self.setFlag(QGraphicsItem.ItemUsesExtendedStyleOption, True)

    def boundingRect(self):
        height, width = self.mask_annotation.mask_data.shape
        return QRectF(0, 0, width, height)

    def paint(self, painter, option, widget):
        if self.mask_annotation.qimage:
            # Apply transparency at render time for instant updates
            transparency = self.mask_annotation.get_current_transparency()
            if transparency < 255:
                painter.setOpacity(transparency / 255.0)
            painter.drawImage(0, 0, self.mask_annotation.qimage)
            if transparency < 255:
                painter.setOpacity(1.0)  # Reset opacity for other drawing operations


class MaskAnnotation(Annotation):
    LOCK_BIT = 2**7  # For uint8, this is 128

    def __init__(self,
                 image_path: str,
                 mask_data: np.ndarray,
                 initial_labels: list,
                 transparency: int = 128,
                 show_msg: bool = False,
                 rasterio_src=None):
        """
        Initialize a full-image semantic segmentation annotation.
        There should only be one MaskAnnotation per image.
        """
        # --- This block is for context ---
        if not initial_labels:
            raise ValueError("initial_labels cannot be empty.")
        placeholder_label = initial_labels[0]

        super().__init__(
            short_label_code=placeholder_label.short_label_code,
            long_label_code=placeholder_label.long_label_code,
            color=placeholder_label.color,
            image_path=image_path,
            label_id=placeholder_label.id,
            transparency=transparency,
            show_msg=show_msg
        )
        
        self.mask_data = mask_data.astype(np.uint8)
        
        self.class_id_to_label_map = {}
        self.label_id_to_class_id_map = {}
        self.visible_label_ids = set()
        
        self.next_class_id = 1
        self.sync_label_map(initial_labels)
        # Update visible labels to include all labels currently in the map
        self.visible_label_ids = set(self.label_id_to_class_id_map.keys())

        self.offset = QPointF(0, 0)
        self.rasterio_src = rasterio_src
        
        # Initialize the colored canvas and QImage once at creation time.
        # This moves the expensive, one-time image generation from the first
        # brush stroke to the object's creation.
        self.colored_mask = None
        self.qimage = None
        self._initialize_canvas()
        
        self.set_centroid()
        self.set_cropped_bbox()
        
        # Cache for class statistics, initialized to None
        self._stats_cache = None

    def sync_label_map(self, current_labels_in_project: list):
        """Ensures the internal maps are synced with the project's labels,
           assigning new, stable IDs to any new labels without changing existing ones."""
        
        # Add any new labels from the project list to our maps
        for label in current_labels_in_project:
            if label.id not in self.label_id_to_class_id_map:
                new_id = self.next_class_id
                self.class_id_to_label_map[new_id] = label
                self.label_id_to_class_id_map[label.id] = new_id
                self.next_class_id += 1
                
                # Also add the new label's ID to the set of visible labels.
                self.visible_label_ids.add(label.id)
                
    def _build_color_map(self):
        """Builds a numpy array mapping class IDs to RGBA colors."""
        # Ensure the map is large enough to handle locked class IDs
        max_id = max(self.class_id_to_label_map.keys()) if self.class_id_to_label_map else 0
        map_size = max(256, max_id + self.LOCK_BIT + 1)
        color_map = np.zeros((map_size, 4), dtype=np.uint8)

        for class_id, label in self.class_id_to_label_map.items():
            color = label.color
            # Always use full alpha in the array - transparency applied at render time
            alpha = 255 if label.id in self.visible_label_ids else 0
            
            # Set the color for both the normal and the locked version of the class ID
            rgba = [color.red(), color.green(), color.blue(), alpha]
            color_map[class_id] = rgba
            if class_id + self.LOCK_BIT < map_size:
                color_map[class_id + self.LOCK_BIT] = rgba

        return color_map
    
    def _initialize_canvas(self):
        """
        Creates the initial color canvas and QImage. This is the one-time expensive
        operation that happens when the mask is first loaded.
        """
        height, width = self.mask_data.shape
        color_map = self._build_color_map()
        
        # Create the full-size 4-channel RGBA numpy array
        self.colored_mask = color_map[self.mask_data]
        
        # Create a QImage that is a VIEW on the numpy array's data buffer.
        # Modifying the numpy array will now automatically update the QImage.
        self.qimage = QImage(self.colored_mask.data, width, height, QImage.Format_RGBA8888)

    def _update_full_canvas(self):
        """Regenerates the entire color canvas. Used for global changes like label color edits."""
        color_map = self._build_color_map()
        # Update the existing colored_mask array in-place
        np.copyto(self.colored_mask, color_map[self.mask_data])

    def _update_canvas_slice(self, update_rect):
        """
        Efficiently updates only a small rectangular slice of the color canvas.
        Used for tools like the brush.
        """
        x1, y1, x2, y2 = update_rect
        
        # Get the slice of the data that has changed
        data_slice = self.mask_data[y1:y2, x1:x2]
        
        # Rebuild the color map in case label properties have changed
        color_map = self._build_color_map()
        
        # Perform the color lookup ONLY for the small data slice
        color_slice = color_map[data_slice]
        
        # Update the corresponding slice in our persistent color canvas
        self.colored_mask[y1:y2, x1:x2] = color_slice

    def set_centroid(self):
        """Set the centroid to the center of the image."""
        height, width = self.mask_data.shape
        self.center_xy = QPointF(width / 2.0, height / 2.0)

    def set_cropped_bbox(self):
        """Set the bounding box to the full dimensions of the image."""
        height, width = self.mask_data.shape
        self.cropped_bbox = (0, 0, width, height)
        self.annotation_size = int(max(width, height))
                
    def contains_point(self, point: QPointF) -> bool:
        """Check if a point is within the mask's classified area."""
        x, y = int(point.x()), int(point.y())
        height, width = self.mask_data.shape
        if 0 <= y < height and 0 <= x < width:
            return self.mask_data[y, x] > 0
        return False

    def get_area(self):
        """Return the total number of non-background pixels."""
        return np.count_nonzero(self.mask_data)

    def get_bounding_box_top_left(self):
        """Get the top-left corner of the annotation's bounding box (always 0,0)."""
        return QPointF(0, 0)

    def get_bounding_box_bottom_right(self):
        """Get the bottom-right corner of the annotation's bounding box."""
        height, width = self.mask_data.shape
        return QPointF(width, height)

    def create_graphics_item(self, scene: QGraphicsScene):
        """Create a QGraphicsPixmapItem to display the mask."""
        self.graphics_item = MaskGraphicsItem(self)
        scene.addItem(self.graphics_item)

    def update_graphics_item(self, update_rect=None):
        """Update the graphics item when mask data has changed."""
        if self.graphics_item is None:
            return
        
        if update_rect:
            # Localized update for brush strokes - update only the changed area
            self._update_canvas_slice(update_rect)
            qt_rect = QRectF(update_rect[0], 
                             update_rect[1], 
                             update_rect[2] - update_rect[0], 
                             update_rect[3] - update_rect[1])
            self.graphics_item.update(qt_rect)
        else:
            # Full update for global changes (e.g., label color changes)
            self._update_full_canvas()
            self.graphics_item.update()

    def update_mask(self, brush_location: QPointF, brush_mask: np.ndarray, new_class_id: int):
        """
        Modify the mask data based on a brush stroke, respecting pre-locked pixels.
        """
        x_start, y_start = int(brush_location.x()), int(brush_location.y())
        brush_h, brush_w = brush_mask.shape
        mask_h, mask_w = self.mask_data.shape

        # Define the update area and clip it to the mask's bounds
        x_end = min(x_start + brush_w, mask_w)
        y_end = min(y_start + brush_h, mask_h)
        clipped_x_start = max(x_start, 0)
        clipped_y_start = max(y_start, 0)

        if clipped_x_start >= x_end or clipped_y_start >= y_end:
            return
            
        # Get the slice of the main mask data we will be updating
        target_slice = self.mask_data[clipped_y_start:y_end, clipped_x_start:x_end]
        
        # Clip the user's brush mask to match the on-screen portion
        brush_x_offset = clipped_x_start - x_start
        brush_y_offset = clipped_y_start - y_start
        clipped_brush_mask = brush_mask[brush_y_offset:brush_y_offset + target_slice.shape[0],
                                        brush_x_offset:brush_x_offset + target_slice.shape[1]]
        
        # Identify pixels within the target slice that are NOT locked
        unlocked_pixels_mask = target_slice < self.LOCK_BIT
        
        # The final brush can only be applied where the brush shape is true AND the target pixel is not locked
        final_brush_mask = clipped_brush_mask & unlocked_pixels_mask

        # Apply the final, protected brush mask to the data
        target_slice[final_brush_mask] = new_class_id
        
        # Trigger a visual update for the changed rectangle
        changed_rect_coords = (clipped_x_start, clipped_y_start, x_end, y_end)
        self.update_graphics_item(update_rect=changed_rect_coords)
        
        self._invalidate_stats_cache()
        self.annotationUpdated.emit(self)
        
    def update_mask_with_mask(self, subset_mask: np.ndarray, top_left: tuple[int, int]):
        """
        Updates a subset area of the mask with a provided mask containing multiple labels.
        
        Args:
            subset_mask: A 2D numpy array with class IDs for the subset area.
            top_left: A tuple (x, y) specifying the top-left corner where to apply the subset.
        """
        x, y = top_left
        h, w = subset_mask.shape
        
        # Calculate the region in the full mask where the subset will be applied
        x_end = min(x + w, self.mask_data.shape[1])
        y_end = min(y + h, self.mask_data.shape[0])
        x_start = max(x, 0)
        y_start = max(y, 0)
        
        # Calculate the corresponding region in the subset mask
        sx_start = x_start - x
        sy_start = y_start - y
        sx_end = sx_start + (x_end - x_start)
        sy_end = sy_start + (y_end - y_start)
        
        # Get the target slice
        target_slice = self.mask_data[y_start:y_end, x_start:x_end]
        subset_slice = subset_mask[sy_start:sy_end, sx_start:sx_end]
        
        # Identify pixels within the target slice that are NOT locked
        unlocked_pixels_mask = target_slice < self.LOCK_BIT
        
        # Apply the subset mask only to unlocked pixels
        target_slice[unlocked_pixels_mask] = subset_slice[unlocked_pixels_mask]
        
        # Trigger a visual update for the changed rectangle
        update_rect = (x_start, y_start, x_end, y_end)
        self.update_graphics_item(update_rect=update_rect)
        
        self._invalidate_stats_cache()
        self.annotationUpdated.emit(self)
        
    def update_mask_with_prediction_mask(self, prediction_mask, top_left=(0, 0)):
        """
        Updates a full-size prediction mask with the current mask data.

        This method is non-destructive; it only updates pixels where the
        prediction_mask has a valid (non-background) prediction.
        The final update respects any locked pixels.

        Args:
            prediction_mask (np.ndarray): A full-size (H, W) mask containing
                                          the new predictions.
        """
        if prediction_mask.shape != self.mask_data.shape:
            print(f"Error: Prediction mask shape {prediction_mask.shape} "
                  f"does not match annotation shape {self.mask_data.shape}.")
            return

        # 1. Get (row, col) indices of all valid predictions (class > 0)
        #    np.nonzero is highly efficient for sparse arrays.
        pred_rows, pred_cols = np.nonzero(prediction_mask)
        
        if pred_rows.size == 0:
            return  # No predictions, nothing to do.

        # 2. Find the bounding box (Region of Interest) of the predictions
        min_row, max_row = np.min(pred_rows), np.max(pred_rows)
        min_col, max_col = np.min(pred_cols), np.max(pred_cols)
        
        # 3. Define the top-left (x, y) corner for pasting
        #    Note: col = x, row = y
        paste_top_left = (min_col, min_row)
        
        # 4. Extract the *small tile* from the prediction mask
        prediction_tile = prediction_mask[min_row: max_row + 1, min_col: max_col + 1]

        # 5. Extract the corresponding *small tile* from the *current* mask data
        original_tile = self.mask_data[min_row: max_row + 1, min_col: max_col + 1]

        # 6. Create the *merged tile*
        #    Use np.where: if pred_tile > 0, use its value, else use original value.
        merged_tile = np.where(
            prediction_tile > 0,
            prediction_tile,
            original_tile
        )
        
        # 7. Call the existing update method with the *small merged tile*
        #    This reuses all your existing logic for locked pixels and
        #    graphics updates, but now only on a small, efficient region.
        self.update_mask_with_mask(merged_tile, paste_top_left)
        
        # Proactively recalculate stats, as this signals the end of an edit session
        self.recalculate_class_statistics()
        
    def get_current_transparency(self):
        """Get the current transparency value for rendering."""
        # Try to get the active label's transparency from the application
        try:
            from PyQt5.QtWidgets import QApplication
            app = QApplication.instance()
            if app:
                main_window = app.activeWindow()
                if hasattr(main_window, 'label_window') and hasattr(main_window.label_window, 'active_label'):
                    active_label = main_window.label_window.active_label
                    if active_label:
                        return active_label.transparency
        except Exception:
            pass
        return self.transparency
    
    def update_transparency(self, transparency):
        """Update transparency instantly using render-time application."""
        transparency = max(0, min(255, transparency))  # Clamp to valid range
        if self.transparency != transparency:
            self.transparency = transparency
            # Update transparency on all visible labels
            for label_id in self.visible_label_ids:
                label = next((lbl for lbl in self.class_id_to_label_map.values() if lbl.id == label_id), None)
                if label:
                    label.transparency = transparency
            
            # Trigger repaint - transparency applied during rendering
            if self.graphics_item:
                self.graphics_item.update()
            
    def update_visible_labels(self, visible_ids: set):
        """
        Updates the set of visible label IDs and triggers a redraw of the mask.

        Args:
            visible_ids (set): A set containing the UUIDs of labels that should be visible.
        """
        self.visible_label_ids = visible_ids
        self.update_graphics_item()
            
    def remove_from_scene(self):
        """Removes the graphics item from its scene, if it exists."""
        if self.graphics_item and self.graphics_item.scene():
            self.graphics_item.scene().removeItem(self.graphics_item)
            
        # Remove the graphics item reference
        self.graphics_item = None

    # --- Data Manipulation & Editing Methods ---

    def fill_region(self, point: QPointF, new_class_id: int):
        """
        Fills a contiguous region with a new class ID, respecting pre-locked pixels.
        """
        x, y = int(point.x()), int(point.y())
        height, width = self.mask_data.shape
        if not (0 <= y < height and 0 <= x < width):
            return

        # Check if the starting pixel is locked. If so, we cannot fill from it.
        if self.mask_data[y, x] >= self.LOCK_BIT:
            return

        old_class_id = self.mask_data[y, x]
        if old_class_id == new_class_id:
            return
        
        QApplication.setOverrideCursor(Qt.WaitCursor)

        # Find all contiguous pixels with the same starting class ID. This automatically respects
        # locked areas because they will have a different ID (e.g., 3 vs 131) and act as a border.
        labeled_array, num_features = ndimage_label(self.mask_data == old_class_id)
        region_label = labeled_array[y, x]
        fill_mask = labeled_array == region_label
        
        # Apply the fill. No further checks are needed because locked pixels were excluded by the ndimage_label step.
        self.mask_data[fill_mask] = new_class_id
        
        # Use localized update for efficiency
        coords = np.where(fill_mask)
        y_min, y_max = coords[0].min(), coords[0].max()
        x_min, x_max = coords[1].min(), coords[1].max()
        update_rect = (x_min, y_min, x_max + 1, y_max + 1)
        self.update_graphics_item(update_rect=update_rect)

        self._invalidate_stats_cache()
        QApplication.restoreOverrideCursor()
        self.annotationUpdated.emit(self)

    def _fast_rasterize(self, geometries, width, height, mode="rasterio"):
        """
        Fast vectorized rasterization using either Shapely or Rasterio.
        
        Args:
            geometries: List of Shapely Polygon geometries
            width, height: Dimensions of the output mask
            mode: "shapely" for Shapely vectorized method, "rasterio" for rasterio.features.rasterize
        
        Returns:
            Boolean numpy array where True indicates pixels covered by geometries
        """
        if mode == "shapely":
            # Use Shapely vectorized method (very fast for many points)
            from shapely.vectorized import contains
            
            # Create coordinate grids - much faster than manual iteration
            y_coords, x_coords = np.mgrid[0:height, 0:width]
            
            # Create empty mask
            raster_mask = np.zeros((height, width), dtype=bool)
            
            # Process each geometry
            for geom in geometries:
                if geom.is_valid:
                    # Vectorized point-in-polygon check - VERY fast
                    mask = contains(geom, x_coords, y_coords)
                    raster_mask = raster_mask | mask
                    
            return raster_mask
            
        elif mode == "rasterio":
            # Use rasterio.features.rasterize (more robust for complex geometries)
            raster_mask = rasterize(
                geometries,
                out_shape=(height, width),
                fill=0,
                default_value=1,
                dtype=np.uint8
            ).astype(bool)
            
            return raster_mask
        
        else:
            raise ValueError(f"Unknown rasterization mode: {mode}. Use 'shapely' or 'rasterio'.")
        
    def rasterize_annotations(self, all_annotations: list):
        """
        Unified method to sync vector annotations with the mask.
        
        Args:
            all_annotations: List of vector annotations to process
        """        
        if not all_annotations:
            return  # Nothing to do if no annotations
        
        # Make cursor busy
        QApplication.setOverrideCursor(Qt.WaitCursor)
            
        height, width = self.mask_data.shape
        
        # Section 1: Build geometries list AND filter annotations that actually need processing
        geometries = []
        annotations_to_process = []
        all_bounds = []
        
        for annotation in all_annotations:
            if not hasattr(annotation, 'get_polygon'):
                continue
            try:
                polygon = annotation.get_polygon()
                if polygon is None or polygon.isEmpty():
                    continue
                points = [(polygon.at(i).x(), polygon.at(i).y()) for i in range(polygon.count())]
                if len(points) < 3:
                    continue
                shapely_poly = Polygon(points)
                
                # Get the class ID for this annotation's label
                class_id = self.label_id_to_class_id_map.get(annotation.label.id)
                if class_id is None:
                    continue  # Skip if label not in mask
                    
                geometries.append(shapely_poly)
                annotations_to_process.append(annotation)
                all_bounds.append(shapely_poly.bounds)
                
            except Exception as e:
                print(f"Warning: Could not process annotation {annotation.id}: {e}")
                continue
        
        if not geometries:
            # No valid geometries to rasterize
            QApplication.restoreOverrideCursor()
            return
        
        # No need to update if the mask is already empty, check before rasterization
        update_canvas = True if np.any(self.mask_data) else False

        # Section 2: Rasterize geometries
        annotation_mask = self._fast_rasterize(geometries, width, height, mode="rasterio")
        
        # Section 3: Clear mask pixels under annotations FIRST
        if np.any(annotation_mask):
            # Clear mask pixels under annotations (set to 0)
            self.mask_data[annotation_mask] = 0
        
        # Section 4: Apply locking to the cleared areas
        # Only lock pixels that are NOT already locked
        to_lock = annotation_mask & (self.mask_data < self.LOCK_BIT)
        if np.any(to_lock):
            self.mask_data[to_lock] += self.LOCK_BIT
            
        # Invalidate the cache since we modified the mask data
        if np.any(annotation_mask) or np.any(to_lock):
            self._invalidate_stats_cache()
        
        if not update_canvas:
            QApplication.restoreOverrideCursor()
            return  # No visual update needed
        
        # Calculate smart combined bounding box
        if all_bounds:
            min_x = min(b[0] for b in all_bounds)
            min_y = min(b[1] for b in all_bounds)
            max_x = max(b[2] for b in all_bounds)
            max_y = max(b[3] for b in all_bounds)
            
            # Add padding (5 pixels) to handle anti-aliasing edges
            x_min = max(0, int(min_x) - 5)
            y_min = max(0, int(min_y) - 5)
            x_max = min(width, int(max_x) + 6)  # +1 for inclusive, +5 for padding
            y_max = min(height, int(max_y) + 6)
            
            update_rect = (x_min, y_min, x_max, y_max)
        else:
            update_rect = None
            
        # Section 5: Smart localized repaint
        if np.any(to_lock) or np.any(annotation_mask):
            if update_rect and (x_max - x_min) * (y_max - y_min) < width * height * 0.3:  # If < 30% of image
                # Use localized update - much faster
                self.update_graphics_item(update_rect=update_rect)
            else:
                # Fall back to full update if annotations cover too much area
                self.update_graphics_item()

        # Restore cursor
        QApplication.restoreOverrideCursor()

    def unrasterize_annotations(self):
        """
        Remove lock protection from all pixels that were marked as locked.
        This allows mask editing over previously protected vector annotation areas.
        """
        # Find all pixels that have the lock bit set and remove it
        locked_pixels = self.mask_data >= self.LOCK_BIT
        
        # Remove the lock bit from these pixels, keeping their original class
        if np.any(locked_pixels):
            self.mask_data[locked_pixels] = self.mask_data[locked_pixels] - self.LOCK_BIT

        # Proactively recalculate stats, as this signals the end of an edit session
        self.recalculate_class_statistics()

    def clear_pixels_for_class(self, class_id: int):
        """Finds all pixels matching a class ID (both locked and unlocked) and resets them to 0."""
        if class_id == 0:  # Cannot clear background class
            return

        # Create a boolean mask of all pixels whose real class ID matches.
        pixels_to_clear = (self.mask_data % self.LOCK_BIT) == class_id
        
        # Set these pixels back to 0 (unclassified).
        self.mask_data[pixels_to_clear] = 0
        self._invalidate_stats_cache()
        
    def clear_pixels_for_annotations(self, annotations_to_clear: list):
        """
        Rasterizes a list of vector annotations and sets the corresponding
        pixels in the mask_data to 0 (unclassified).
        """
        if not annotations_to_clear:
            return

        # 1. Convert all annotation polygons to Shapely Polygons.
        geometries = []
        for anno in annotations_to_clear:
            if hasattr(anno, 'get_polygon'):
                qt_polygon = anno.get_polygon()
                points = [(p.x(), p.y()) for p in qt_polygon]
                if len(points) >= 3:
                    geometries.append(Polygon(points))

        if not geometries:
            return

        # 2. Rasterize all shapes at once into a boolean mask.
        # This creates a numpy array where 'True' indicates a pixel is covered
        # by at least one of the vector annotations.
        height, width = self.mask_data.shape
        clear_mask = rasterize(
            geometries,
            out_shape=(height, width),
            fill=0,
            default_value=1,
            dtype=np.uint8
        ).astype(bool)

        # 3. Apply the mask to the data, setting pixels to 0.
        self.mask_data[clear_mask] = 0
        self._invalidate_stats_cache()

        # 4. Trigger a full repaint of the mask to show the changes.
        self.update_graphics_item()

    # --- Analysis & Information Retrieval Methods ---

    def _invalidate_stats_cache(self):
        """Invalidates the statistics cache.
        This should be called by any method that modifies self.mask_data.
        """
        self._stats_cache = None

    def recalculate_class_statistics(self) -> dict:
        """
        Performs the expensive calculation of class statistics and updates the cache.
        This is called proactively when mask editing is finished.
        """
        stats = {}
        
        # 1. Use modulus to get the "unlocked" class ID for every pixel.
        # This is a single, fast, vectorized operation.
        unlocked_mask_data = self.mask_data % self.LOCK_BIT

        # 2. Get total pixels in the entire image.
        total_image_pixels = self.mask_data.size

        # 3. Handle division by zero for an empty/invalid image
        if total_image_pixels == 0:
            self._stats_cache = stats
            return self._stats_cache

        # 4. Determine the minimum length for the bincount array
        # to ensure it's large enough for all known class IDs.
        max_id = max(self.class_id_to_label_map.keys()) if self.class_id_to_label_map else 0
        min_len = max_id + 1

        # 5. Get counts for ALL class IDs in one O(N) pass.
        bin_counts = np.bincount(unlocked_mask_data.ravel(), minlength=min_len)

        # 6. Populate stats by iterating only through the labels we care about.
        for class_id, label in self.class_id_to_label_map.items():
            # class_id is already guaranteed to be < min_len from step 4
            count = bin_counts[class_id]
            if count > 0:
                stats[label.short_label_code] = {
                    "pixel_count": int(count),
                    # FIXED: Percentage is (count / total_image_pixels)
                    "percentage": (count / total_image_pixels) * 100
                }
        
        self._stats_cache = stats
        return self._stats_cache

    def get_class_statistics(self) -> dict:
        """
        Gets class statistics from the cache.
        If the cache is empty, it triggers a one-time recalculation.
        """
        # If cache is empty, calculate it
        if self._stats_cache is None:
            # This is the "lazy" part, acting as a safety net
            self.recalculate_class_statistics()
        
        return self._stats_cache
    
    @property
    def cached_statistics(self) -> dict | None:
        """
        Safely returns the cached statistics dictionary, or None if not cached.
        This method will NOT trigger a recalculation.
        """
        return self._stats_cache

    def get_class_at_point(self, point: QPointF) -> int:
        """Returns the class ID at a specific point."""
        x, y = int(point.x()), int(point.y())
        height, width = self.mask_data.shape
        if 0 <= y < height and 0 <= x < width:
            return self.mask_data[y, x]
        return 0  # Return background class if outside bounds

    # --- Conversion & Exporting Methods ---
    
    def get_binary_mask(self, class_id: int) -> np.ndarray:
        """Returns a boolean numpy array where True corresponds to the given class ID."""
        return self.mask_data == class_id

    def to_instance_polygons(self, class_id: int) -> list:
        """Converts all contiguous regions of a class ID into PolygonAnnotations."""
        binary_mask = self.get_binary_mask(class_id)
        # Add padding to handle contours touching the border
        padded_mask = np.pad(binary_mask, pad_width=1, mode='constant', constant_values=0)
        
        # Level is 0.5 to find contours between 0 and 1 values
        contours = find_contours(padded_mask, level=0.5)
        
        annotations = []
        for contour in contours:
            # Remove padding offset and swap (row, col) to (x, y)
            points = [QPointF(p[1] - 1, p[0] - 1) for p in contour]
            if len(points) > 2:  # Must have at least 3 points for a valid polygon
                # Use the label associated with the class_id for the new annotation
                label = self.class_id_to_label_map[class_id]
                anno = PolygonAnnotation(
                    points=points,
                    short_label_code=label.short_label_code,
                    long_label_code=label.long_label_code,
                    color=label.color,
                    image_path=self.image_path,
                    label_id=label.id
                )
                annotations.append(anno)
        return annotations

    def export_as_png(self, path: str, use_label_colors: bool = True):
        """Saves the mask to a PNG file."""
        if use_label_colors:
            # Use current colored image for export
            self.qimage.save(path)
        else:
            # Save the raw class IDs as a grayscale image
            height, width = self.mask_data.shape
            # Ensure data is in a format QImage can handle (e.g., 8-bit grayscale)
            if self.mask_data.max() < 256:
                img_data = self.mask_data.astype(np.uint8)
                q_image = QImage(img_data.data, width, height, QImage.Format_Grayscale8)
                q_image.save(path)
            else:
                warnings.warn("Mask contains class IDs > 255; cannot save as 8-bit grayscale PNG.")

    def export_as_raster(self, path: str):
        """Saves the mask data to a raster file (e.g., GeoTIFF) using rasterio."""
        profile = {
            'driver': 'GTiff',
            'height': self.mask_data.shape[0],
            'width': self.mask_data.shape[1],
            'count': 1,
            'dtype': self.mask_data.dtype
        }
        
        # If the original image was opened with rasterio, copy its spatial metadata
        if self.rasterio_src:
            profile['crs'] = self.rasterio_src.crs
            profile['transform'] = self.rasterio_src.transform

        with rasterio.open(path, 'w', **profile) as dst:
            dst.write(self.mask_data, 1)

    # --- Serialization & Deserialization ---

    def to_dict(self):
        """Serialize the annotation to a dictionary, with RLE for the mask."""
        base_dict = super().to_dict()
        
        # Encode each class's binary mask using pycocotools
        rle_list = []
        unique_classes = np.unique(self.mask_data)
        for class_id in unique_classes:
            if class_id == 0:
                continue
            binary_mask = (self.mask_data == class_id).astype(np.uint8)
            rle = mask.encode(np.asfortranarray(binary_mask))
            rle['counts'] = base64.b64encode(rle['counts']).decode('ascii')
            rle_list.append({'class_id': int(class_id), 'rle': rle})
        
        # Convert the label map to a serializable format
        serializable_label_map = {}
        for cid, label in self.class_id_to_label_map.items():
            serializable_label_map[cid] = label.short_label_code

        base_dict.update({
            'shape': self.mask_data.shape,
            'rle_masks': rle_list,
            'label_map': serializable_label_map
        })
        return base_dict

    @classmethod
    def from_dict(cls, data, label_window):
        """Instantiate a MaskAnnotation from a dictionary."""
        # Get all labels currently in the project. This is needed for the constructor.
        all_project_labels = list(label_window.labels)
        if not all_project_labels:
            raise ValueError("Cannot import a MaskAnnotation without any labels loaded in the project.")

        # Decode the RLE mask data
        shape = tuple(data['shape'])
        mask_data = np.zeros(shape, dtype=np.uint8)
        for item in data['rle_masks']:
            class_id = item['class_id']
            rle = item['rle']
            try:
                rle['counts'] = base64.b64decode(rle['counts'])
                binary_mask = mask.decode(rle).astype(bool)
                if binary_mask.shape != shape:
                    print(f"Warning: RLE decoded shape {binary_mask.shape} does not match expected shape {shape}")
                    continue
                mask_data[binary_mask] = class_id
            except Exception as e:
                print(f"Error decoding RLE for class {class_id}: {e}")
                continue

        # Create the base annotation instance. It will have a generic label map initially.
        annotation = cls(
            image_path=data['image_path'],
            mask_data=mask_data,
            initial_labels=all_project_labels
        )
        
        # Clear the generic maps created by the constructor.
        annotation.class_id_to_label_map.clear()
        annotation.label_id_to_class_id_map.clear()
        
        max_id_found = 0
        if 'label_map' in data and data['label_map']:
            # Iterate through the saved map {class_id: short_code}
            for cid_str, short_code in data['label_map'].items():
                class_id = int(cid_str)
                label = label_window.get_label_by_short_code(short_code)
                
                if label:
                    # Rebuild the maps with the correct associations
                    annotation.class_id_to_label_map[class_id] = label
                    annotation.label_id_to_class_id_map[label.id] = class_id
                    if class_id > max_id_found:
                        max_id_found = class_id
                else:
                    print(f"Warning: Label with short code '{short_code}' not found in project during mask import.")

        # Ensure the next class ID is set correctly to avoid future conflicts.
        annotation.next_class_id = max_id_found + 1

        # Restore other annotation properties
        annotation.id = data.get('id', annotation.id)
        annotation.data = data.get('data', {})
        
        return annotation

    @classmethod
    def from_rasterio(cls, file_path: str, image_path: str, all_labels: list):
        """Creates a MaskAnnotation instance by loading data from a raster file."""
        with rasterio.open(file_path) as src:
            mask_data = src.read(1)
            return cls(
                image_path=image_path,
                mask_data=mask_data,
                initial_labels=all_labels,
                rasterio_src=src
            )

    # --- Compatibility Methods ---
    def get_perimeter(self):
        height, width = self.mask_data.shape
        return 2 * (width + height)
    
    def get_polygon(self):
        height, width = self.mask_data.shape
        return QPolygonF(QRectF(0, 0, width, height))

    def __repr__(self):
        return (f"MaskAnnotation(id={self.id}, image_path={self.image_path}, "
                f"shape={self.mask_data.shape})")