"""Change detection module for remote sensing imagery using torchange."""

import logging
import os
import threading
from typing import Any, Dict, List, Optional, Tuple, Union

logger = logging.getLogger(__name__)

import matplotlib.pyplot as plt
import numpy as np
import rasterio
from rasterio.windows import from_bounds
from skimage.transform import resize

try:
    from torchange.models.segment_any_change import AnyChange, show_change_masks
except ImportError:
    AnyChange = None
    show_change_masks = None

from .utils import download_file


class ChangeDetection:
    """A class for change detection on geospatial imagery using torchange and SAM."""

    def __init__(self, sam_model_type="vit_h", sam_checkpoint=None) -> None:
        """
        Initialize the ChangeDetection class.

        Args:
            sam_model_type (str): SAM model type ('vit_h', 'vit_l', 'vit_b')
            sam_checkpoint (str): Path to SAM checkpoint file
        """
        self.sam_model_type = sam_model_type
        self.sam_checkpoint = sam_checkpoint
        self.model = None
        self._init_model()

    def _init_model(self):
        """Initialize the AnyChange model."""
        if AnyChange is None:
            raise ImportError(
                "The 'torchange' package is required for change detection. "
                "Please install it using: pip install torchange\n"
                "Note: torchange requires Python 3.11 or higher."
            )

        if self.sam_checkpoint is None:
            self.sam_checkpoint = download_checkpoint(self.sam_model_type)

        self.model = AnyChange(self.sam_model_type, sam_checkpoint=self.sam_checkpoint)

        # Set default hyperparameters
        self.model.make_mask_generator(
            points_per_side=32,
            stability_score_thresh=0.95,
        )
        self.model.set_hyperparameters(
            change_confidence_threshold=145,
            use_normalized_feature=True,
            bitemporal_match=True,
        )

    def set_hyperparameters(
        self,
        change_confidence_threshold: int = 155,
        auto_threshold: bool = False,
        use_normalized_feature: bool = True,
        area_thresh: float = 0.8,
        match_hist: bool = False,
        object_sim_thresh: int = 60,
        bitemporal_match: bool = True,
        **kwargs: Any,
    ) -> None:
        """
        Set hyperparameters for the change detection model.

        Args:
            change_confidence_threshold (int): Change confidence threshold for SAM
            auto_threshold (bool): Whether to use auto threshold for SAM
            use_normalized_feature (bool): Whether to use normalized feature for SAM
            area_thresh (float): Area threshold for SAM
            match_hist (bool): Whether to use match hist for SAM
            object_sim_thresh (int): Object similarity threshold for SAM
            bitemporal_match (bool): Whether to use bitemporal match for SAM
            **kwargs: Keyword arguments for model hyperparameters
        """
        if self.model:
            self.model.set_hyperparameters(
                change_confidence_threshold=change_confidence_threshold,
                auto_threshold=auto_threshold,
                use_normalized_feature=use_normalized_feature,
                area_thresh=area_thresh,
                match_hist=match_hist,
                object_sim_thresh=object_sim_thresh,
                bitemporal_match=bitemporal_match,
                **kwargs,
            )

    def set_mask_generator_params(
        self,
        points_per_side: int = 32,
        points_per_batch: int = 64,
        pred_iou_thresh: float = 0.5,
        stability_score_thresh: float = 0.95,
        stability_score_offset: float = 1.0,
        box_nms_thresh: float = 0.7,
        point_grids: Optional[List] = None,
        min_mask_region_area: int = 0,
        **kwargs: Any,
    ) -> None:
        """
        Set mask generator parameters.

        Args:
            points_per_side (int): Number of points per side for SAM
            points_per_batch (int): Number of points per batch for SAM
            pred_iou_thresh (float): IoU threshold for SAM
            stability_score_thresh (float): Stability score threshold for SAM
            stability_score_offset (float): Stability score offset for SAM
            box_nms_thresh (float): NMS threshold for SAM
            point_grids (list): Point grids for SAM
            min_mask_region_area (int): Minimum mask region area for SAM
            **kwargs: Keyword arguments for mask generator
        """
        if self.model:
            self.model.make_mask_generator(
                points_per_side=points_per_side,
                points_per_batch=points_per_batch,
                pred_iou_thresh=pred_iou_thresh,
                stability_score_thresh=stability_score_thresh,
                stability_score_offset=stability_score_offset,
                box_nms_thresh=box_nms_thresh,
                point_grids=point_grids,
                min_mask_region_area=min_mask_region_area,
                **kwargs,
            )

    def _read_and_align_images(self, image1_path, image2_path, target_size=1024):
        """
        Read and align two GeoTIFF images, handling different extents and projections.

        Args:
            image1_path (str): Path to first image
            image2_path (str): Path to second image
            target_size (int): Target size for processing (default 1024 for torchange)

        Returns:
            tuple: (aligned_img1, aligned_img2, transform, crs, bounds)
        """
        with rasterio.open(image1_path) as src1, rasterio.open(image2_path) as src2:
            # Get the intersection of bounds
            bounds1 = src1.bounds
            bounds2 = src2.bounds

            # Calculate intersection bounds
            left = max(bounds1.left, bounds2.left)
            bottom = max(bounds1.bottom, bounds2.bottom)
            right = min(bounds1.right, bounds2.right)
            top = min(bounds1.top, bounds2.top)

            if left >= right or bottom >= top:
                raise ValueError("Images do not overlap")

            intersection_bounds = (left, bottom, right, top)

            # Read the intersecting area from both images
            window1 = from_bounds(*intersection_bounds, src1.transform)
            window2 = from_bounds(*intersection_bounds, src2.transform)

            # Read data
            img1_data = src1.read(window=window1)
            img2_data = src2.read(window=window2)

            # Get transform for the intersecting area
            transform = src1.window_transform(window1)
            crs = src1.crs

            # Convert from (bands, height, width) to (height, width, bands)
            img1_data = np.transpose(img1_data, (1, 2, 0))
            img2_data = np.transpose(img2_data, (1, 2, 0))

            # Use only RGB bands (first 3 channels) for torchange
            if img1_data.shape[2] >= 3:
                img1_data = img1_data[:, :, :3]
            if img2_data.shape[2] >= 3:
                img2_data = img2_data[:, :, :3]

            # Normalize to 0-255 range if needed
            if img1_data.dtype != np.uint8:
                img1_data = (
                    (img1_data - img1_data.min())
                    / (img1_data.max() - img1_data.min())
                    * 255
                ).astype(np.uint8)
            if img2_data.dtype != np.uint8:
                img2_data = (
                    (img2_data - img2_data.min())
                    / (img2_data.max() - img2_data.min())
                    * 255
                ).astype(np.uint8)

            # Store original size for later use
            original_shape = img1_data.shape[:2]

            # Resize to target size for torchange processing
            if img1_data.shape[0] != target_size or img1_data.shape[1] != target_size:
                img1_resized = resize(
                    img1_data, (target_size, target_size), preserve_range=True
                ).astype(np.uint8)
                img2_resized = resize(
                    img2_data, (target_size, target_size), preserve_range=True
                ).astype(np.uint8)
            else:
                img1_resized = img1_data
                img2_resized = img2_data

            return (img1_resized, img2_resized, transform, crs, original_shape)

    def detect_changes(
        self,
        image1_path: str,
        image2_path: str,
        output_path: Optional[str] = None,
        target_size: int = 1024,
        return_results: bool = True,
        export_probability: bool = False,
        probability_output_path: Optional[str] = None,
        export_instance_masks: bool = False,
        instance_masks_output_path: Optional[str] = None,
        return_detailed_results: bool = False,
    ) -> Union[Tuple[Any, np.ndarray, np.ndarray], Dict[str, Any], None]:
        """
        Detect changes between two GeoTIFF images with instance segmentation.

        Args:
            image1_path (str): Path to first image
            image2_path (str): Path to second image
            output_path (str): Optional path to save binary change mask as GeoTIFF
            target_size (int): Target size for processing
            return_results (bool): Whether to return results
            export_probability (bool): Whether to export probability mask
            probability_output_path (str): Path to save probability mask (required if export_probability=True)
            export_instance_masks (bool): Whether to export instance segmentation masks
            instance_masks_output_path (str): Path to save instance masks (required if export_instance_masks=True)
            return_detailed_results (bool): Whether to return detailed mask information

        Returns:
            tuple: (change_masks, img1, img2) if return_results=True
            dict: Detailed results if return_detailed_results=True
        """
        # Read and align images
        img1, img2, transform, crs, original_shape = self._read_and_align_images(
            image1_path, image2_path, target_size
        )

        # Detect changes
        change_masks, _, _ = self.model.forward(img1, img2)

        # If output path specified, save binary mask as GeoTIFF
        if output_path:
            self._save_change_mask(
                change_masks, output_path, transform, crs, original_shape, target_size
            )

        # If probability export requested, save probability mask
        if export_probability:
            if probability_output_path is None:
                raise ValueError(
                    "probability_output_path must be specified when export_probability=True"
                )
            self._save_probability_mask(
                change_masks,
                probability_output_path,
                transform,
                crs,
                original_shape,
                target_size,
            )

        # If instance masks export requested, save instance segmentation masks
        if export_instance_masks:
            if instance_masks_output_path is None:
                raise ValueError(
                    "instance_masks_output_path must be specified when export_instance_masks=True"
                )
            num_instances = self._save_instance_segmentation_masks(
                change_masks,
                instance_masks_output_path,
                transform,
                crs,
                original_shape,
                target_size,
            )

            # Also save instance scores if requested
            scores_path = instance_masks_output_path.replace(".tif", "_scores.tif")
            self._save_instance_scores_mask(
                change_masks,
                scores_path,
                transform,
                crs,
                original_shape,
                target_size,
            )

        # Return detailed results if requested
        if return_detailed_results:
            return self._extract_detailed_results(
                change_masks, transform, crs, original_shape, target_size
            )

        if return_results:
            return change_masks, img1, img2

    def _save_change_mask(
        self, change_masks, output_path, transform, crs, original_shape, target_size
    ):
        """
        Save change masks as a GeoTIFF with proper georeference.

        Args:
            change_masks: Change detection masks (MaskData object)
            output_path (str): Output file path
            transform: Rasterio transform
            crs: Coordinate reference system
            original_shape (tuple): Original image shape
            target_size (int): Processing target size
        """
        # Convert MaskData to binary mask by decoding RLE masks
        combined_mask = np.zeros((target_size, target_size), dtype=bool)

        # Extract RLE masks from MaskData object
        mask_items = dict(change_masks.items())
        if "rles" in mask_items:
            rles = mask_items["rles"]
            for rle in rles:
                if isinstance(rle, dict) and "size" in rle and "counts" in rle:
                    try:
                        # Decode RLE to binary mask
                        size = rle["size"]
                        counts = rle["counts"]

                        # Create binary mask from RLE counts
                        mask = np.zeros(size[0] * size[1], dtype=np.uint8)
                        pos = 0
                        value = 0

                        for count in counts:
                            if pos + count <= len(mask):
                                if value == 1:
                                    mask[pos : pos + count] = 1
                                pos += count
                                value = 1 - value  # Toggle between 0 and 1
                            else:
                                break

                        # RLE is column-major, reshape and transpose
                        mask = mask.reshape(size).T
                        if mask.shape == (target_size, target_size):
                            combined_mask = np.logical_or(
                                combined_mask, mask.astype(bool)
                            )

                    except (ValueError, IndexError, KeyError) as e:
                        logger.warning(f"Failed to decode RLE mask: {e}")
                        continue

        # Convert to uint8 first, then resize if needed
        combined_mask_uint8 = combined_mask.astype(np.uint8) * 255

        # Resize back to original shape if needed
        if original_shape != (target_size, target_size):
            # Use precise resize
            combined_mask_resized = resize(
                combined_mask_uint8.astype(np.float32),
                original_shape,
                preserve_range=True,
                anti_aliasing=False,
                order=0,
            )
            combined_mask = (combined_mask_resized > 127).astype(np.uint8) * 255
        else:
            combined_mask = combined_mask_uint8

        # Save as GeoTIFF
        with rasterio.open(
            output_path,
            "w",
            driver="GTiff",
            height=combined_mask.shape[0],
            width=combined_mask.shape[1],
            count=1,
            dtype=combined_mask.dtype,
            crs=crs,
            transform=transform,
            compress="lzw",
        ) as dst:
            dst.write(combined_mask, 1)

    def _save_probability_mask(
        self, change_masks, output_path, transform, crs, original_shape, target_size
    ):
        """
        Save probability masks as a GeoTIFF with proper georeference.

        Args:
            change_masks: Change detection masks (MaskData object)
            output_path (str): Output file path
            transform: Rasterio transform
            crs: Coordinate reference system
            original_shape (tuple): Original image shape
            target_size (int): Processing target size
        """
        # Extract mask components for probability calculation
        mask_items = dict(change_masks.items())
        rles = mask_items.get("rles", [])
        iou_preds = mask_items.get("iou_preds", None)
        stability_scores = mask_items.get("stability_score", None)
        change_confidence = mask_items.get("change_confidence", None)
        areas = mask_items.get("areas", None)

        # Convert tensors to numpy if needed
        if iou_preds is not None:
            iou_preds = iou_preds.detach().cpu().numpy()
        if stability_scores is not None:
            stability_scores = stability_scores.detach().cpu().numpy()
        if change_confidence is not None:
            change_confidence = change_confidence.detach().cpu().numpy()
        if areas is not None:
            areas = areas.detach().cpu().numpy()

        # Create probability mask
        probability_mask = np.zeros((target_size, target_size), dtype=np.float32)

        # Process each mask with probability weighting
        for i, rle in enumerate(rles):
            if isinstance(rle, dict) and "size" in rle and "counts" in rle:
                try:
                    # Decode RLE to binary mask
                    size = rle["size"]
                    counts = rle["counts"]

                    mask = np.zeros(size[0] * size[1], dtype=np.uint8)
                    pos = 0
                    value = 0

                    for count in counts:
                        if pos + count <= len(mask):
                            if value == 1:
                                mask[pos : pos + count] = 1
                            pos += count
                            value = 1 - value
                        else:
                            break

                    mask = mask.reshape(size).T
                    if mask.shape != (target_size, target_size):
                        continue

                    mask_bool = mask.astype(bool)

                    # Calculate probability using multiple factors
                    prob_components = []

                    # IoU prediction (0-1, higher is better)
                    if iou_preds is not None and i < len(iou_preds):
                        iou_score = float(iou_preds[i])
                        prob_components.append(("iou", iou_score))
                    else:
                        prob_components.append(("iou", 0.8))

                    # Stability score (0-1, higher is better)
                    if stability_scores is not None and i < len(stability_scores):
                        stability = float(stability_scores[i])
                        prob_components.append(("stability", stability))
                    else:
                        prob_components.append(("stability", 0.8))

                    # Change confidence (normalize based on threshold)
                    if change_confidence is not None and i < len(change_confidence):
                        conf = float(change_confidence[i])
                        # Normalize confidence: threshold is 145, values above indicate higher confidence
                        if conf >= 145:
                            conf_normalized = 0.5 + min(0.5, (conf - 145) / 145)
                        else:
                            conf_normalized = max(0.0, conf / 145 * 0.5)
                        prob_components.append(("confidence", conf_normalized))
                    else:
                        prob_components.append(("confidence", 0.5))

                    # Area-based weight (normalize using log scale)
                    if areas is not None and i < len(areas):
                        area = float(areas[i])
                        area_normalized = 0.2 + 0.8 * min(1.0, np.log(area + 1) / 15.0)
                        prob_components.append(("area", area_normalized))
                    else:
                        prob_components.append(("area", 0.6))

                    # Calculate weighted probability
                    weights = {
                        "iou": 0.3,
                        "stability": 0.3,
                        "confidence": 0.35,
                        "area": 0.05,
                    }
                    prob_weight = sum(
                        weights[name] * value for name, value in prob_components
                    )
                    prob_weight = np.clip(prob_weight, 0.0, 1.0)

                    # Add to probability mask (take maximum where masks overlap)
                    current_prob = probability_mask[mask_bool]
                    new_prob = np.maximum(current_prob, prob_weight)
                    probability_mask[mask_bool] = new_prob

                except (ValueError, IndexError, KeyError) as e:
                    logger.warning(f"Failed to process probability mask {i}: {e}")
                    continue

        # Resize back to original shape if needed
        if original_shape != (target_size, target_size):
            prob_resized = resize(
                probability_mask,
                original_shape,
                preserve_range=True,
                anti_aliasing=True,
                order=1,
            )
            prob_final = np.clip(prob_resized, 0.0, 1.0)
        else:
            prob_final = probability_mask

        # Save as float32 GeoTIFF
        with rasterio.open(
            output_path,
            "w",
            driver="GTiff",
            height=prob_final.shape[0],
            width=prob_final.shape[1],
            count=1,
            dtype=rasterio.float32,
            crs=crs,
            transform=transform,
            compress="lzw",
        ) as dst:
            dst.write(prob_final.astype(np.float32), 1)

    def visualize_changes(
        self, image1_path: str, image2_path: str, figsize: Tuple[int, int] = (15, 5)
    ) -> plt.Figure:
        """
        Visualize change detection results.

        Args:
            image1_path (str): Path to first image
            image2_path (str): Path to second image
            figsize (tuple): Figure size

        Returns:
            matplotlib.figure.Figure: The figure object
        """
        if show_change_masks is None:
            raise ImportError(
                "The 'torchange' package is required for change detection visualization. "
                "Please install it using: pip install torchange\n"
                "Note: torchange requires Python 3.11 or higher."
            )

        change_masks, img1, img2 = self.detect_changes(
            image1_path, image2_path, return_results=True
        )

        # Use torchange's visualization function
        fig, _ = show_change_masks(img1, img2, change_masks)
        fig.set_size_inches(figsize)

        return fig

    def visualize_results(
        self,
        image1_path,
        image2_path,
        binary_path,
        prob_path,
        title1="Earlier Image",
        title2="Later Image",
    ) -> None:
        """Create enhanced visualization with probability analysis."""

        # Load data
        with rasterio.open(image1_path) as src:
            img1 = src.read([1, 2, 3])
            img1 = np.transpose(img1, (1, 2, 0))

        with rasterio.open(image2_path) as src:
            img2 = src.read([1, 2, 3])
            img2 = np.transpose(img2, (1, 2, 0))

        with rasterio.open(binary_path) as src:
            binary_mask = src.read(1)

        with rasterio.open(prob_path) as src:
            prob_mask = src.read(1)

        # Create comprehensive visualization
        fig, axes = plt.subplots(2, 4, figsize=(24, 12))

        # Crop for better visualization
        h, w = img1.shape[:2]
        y1, y2 = h // 4, 3 * h // 4
        x1, x2 = w // 4, 3 * w // 4

        img1_crop = img1[y1:y2, x1:x2]
        img2_crop = img2[y1:y2, x1:x2]
        binary_crop = binary_mask[y1:y2, x1:x2]
        prob_crop = prob_mask[y1:y2, x1:x2]

        # Row 1: Original and overlays
        axes[0, 0].imshow(img1_crop)
        axes[0, 0].set_title(title1, fontweight="bold")
        axes[0, 0].axis("off")

        axes[0, 1].imshow(img2_crop)
        axes[0, 1].set_title(title2, fontweight="bold")
        axes[0, 1].axis("off")

        # Binary overlay
        overlay_binary = img2_crop.copy()
        overlay_binary[binary_crop > 0] = [255, 0, 0]
        axes[0, 2].imshow(overlay_binary)
        axes[0, 2].set_title("Binary Changes\n(Red = Change)", fontweight="bold")
        axes[0, 2].axis("off")

        # Probability heatmap
        im1 = axes[0, 3].imshow(prob_crop, cmap="hot", vmin=0, vmax=1)
        axes[0, 3].set_title(
            "Probability Heatmap\n(White = High Confidence)", fontweight="bold"
        )
        axes[0, 3].axis("off")
        plt.colorbar(im1, ax=axes[0, 3], shrink=0.8)

        # Row 2: Detailed probability analysis
        # Confidence levels overlay
        overlay_conf = img2_crop.copy()
        high_conf = prob_crop > 0.7
        med_conf = (prob_crop > 0.4) & (prob_crop <= 0.7)
        low_conf = (prob_crop > 0.1) & (prob_crop <= 0.4)

        overlay_conf[high_conf] = [255, 0, 0]  # Red for high
        overlay_conf[med_conf] = [255, 165, 0]  # Orange for medium
        overlay_conf[low_conf] = [255, 255, 0]  # Yellow for low

        axes[1, 0].imshow(overlay_conf)
        axes[1, 0].set_title(
            "Confidence Levels\n(Red>0.7, Orange>0.4, Yellow>0.1)", fontweight="bold"
        )
        axes[1, 0].axis("off")

        # Thresholded probability (>0.5)
        overlay_thresh = img2_crop.copy()
        high_prob = prob_crop > 0.5
        overlay_thresh[high_prob] = [255, 0, 0]
        axes[1, 1].imshow(overlay_thresh)
        axes[1, 1].set_title(
            "High Confidence Only\n(Probability > 0.5)", fontweight="bold"
        )
        axes[1, 1].axis("off")

        # Probability histogram
        prob_values = prob_crop[prob_crop > 0]
        if len(prob_values) > 0:
            axes[1, 2].hist(
                prob_values, bins=50, alpha=0.7, color="red", edgecolor="black"
            )
            axes[1, 2].axvline(
                x=0.5, color="blue", linestyle="--", label="0.5 threshold"
            )
            axes[1, 2].axvline(
                x=0.7, color="green", linestyle="--", label="0.7 threshold"
            )
            axes[1, 2].set_xlabel("Change Probability")
            axes[1, 2].set_ylabel("Pixel Count")
            axes[1, 2].set_title(
                f"Probability Distribution\n({len(prob_values):,} pixels)"
            )
            axes[1, 2].legend()
            axes[1, 2].grid(True, alpha=0.3)

        # Statistics text
        stats_text = f"""Probability Statistics:
    Min: {np.min(prob_values):.3f}
    Max: {np.max(prob_values):.3f}
    Mean: {np.mean(prob_values):.3f}
    Median: {np.median(prob_values):.3f}

    Confidence Levels:
    High (>0.7): {np.sum(prob_crop > 0.7):,}
    Med (0.4-0.7): {np.sum((prob_crop > 0.4) & (prob_crop <= 0.7)):,}
    Low (0.1-0.4): {np.sum((prob_crop > 0.1) & (prob_crop <= 0.4)):,}"""

        axes[1, 3].text(
            0.05,
            0.95,
            stats_text,
            transform=axes[1, 3].transAxes,
            fontsize=11,
            verticalalignment="top",
            fontfamily="monospace",
        )
        axes[1, 3].set_xlim(0, 1)
        axes[1, 3].set_ylim(0, 1)
        axes[1, 3].axis("off")
        axes[1, 3].set_title("Statistics Summary", fontweight="bold")

        plt.tight_layout()
        plt.suptitle(
            "Enhanced Probability-Based Change Detection",
            fontsize=16,
            fontweight="bold",
            y=0.98,
        )

        plt.savefig("enhanced_probability_results.png", dpi=150, bbox_inches="tight")
        plt.show()

        logger.info(
            "Enhanced visualization saved as 'enhanced_probability_results.png'"
        )

    def create_split_comparison(
        self,
        image1_path,
        image2_path,
        binary_path,
        prob_path,
        title1="Earlier Image",
        title2="Later Image",
        output_path="split_comparison.png",
    ) -> None:
        """Create a split comparison visualization showing before/after with change overlay."""
        import cv2  # Lazy import to avoid QGIS opencv conflicts

        # Load data
        with rasterio.open(image1_path) as src:
            img1 = src.read([1, 2, 3])
            img1 = np.transpose(img1, (1, 2, 0))
            if img1.dtype != np.uint8:
                img1 = ((img1 - img1.min()) / (img1.max() - img1.min()) * 255).astype(
                    np.uint8
                )

        with rasterio.open(image2_path) as src:
            img2 = src.read([1, 2, 3])
            img2 = np.transpose(img2, (1, 2, 0))
            if img2.dtype != np.uint8:
                img2 = ((img2 - img2.min()) / (img2.max() - img2.min()) * 255).astype(
                    np.uint8
                )

        with rasterio.open(prob_path) as src:
            prob_mask = src.read(1)

        # Ensure all arrays have the same shape
        h, w = img1.shape[:2]
        if prob_mask.shape != (h, w):
            prob_mask = resize(
                prob_mask, (h, w), preserve_range=True, anti_aliasing=True, order=1
            )

        # Create split comparison
        fig, ax = plt.subplots(1, 1, figsize=(15, 10))

        # Create combined image - left half is earlier, right half is later
        combined_img = np.zeros_like(img1)
        combined_img[:, : w // 2] = img1[:, : w // 2]
        combined_img[:, w // 2 :] = img2[:, w // 2 :]

        # Create overlay with changes - ensure prob_mask is 2D and matches image dimensions
        overlay = combined_img.copy()
        high_conf_changes = prob_mask > 0.5

        # Apply overlay only where changes are detected
        if len(overlay.shape) == 3:  # RGB image
            overlay[high_conf_changes] = [255, 0, 0]  # Red for high confidence changes

        # Blend overlay with original
        blended = cv2.addWeighted(combined_img, 0.7, overlay, 0.3, 0)

        ax.imshow(blended)
        ax.axvline(x=w // 2, color="white", linewidth=3, linestyle="--", alpha=0.8)
        ax.text(
            w // 4,
            50,
            title1,
            fontsize=20,
            color="white",
            ha="center",
            bbox={"boxstyle": "round,pad=0.3", "facecolor": "black", "alpha": 0.8},
        )
        ax.text(
            3 * w // 4,
            50,
            title2,
            fontsize=20,
            color="white",
            ha="center",
            bbox={"boxstyle": "round,pad=0.3", "facecolor": "black", "alpha": 0.8},
        )

        ax.set_title(
            "Split Comparison with Change Detection\n(Red = High Confidence Changes)",
            fontsize=16,
            fontweight="bold",
            pad=20,
        )
        ax.axis("off")

        plt.tight_layout()
        plt.savefig(output_path, dpi=150, bbox_inches="tight")
        plt.show()

        logger.info(f"Split comparison saved as '{output_path}'")

    def analyze_instances(
        self, instance_mask_path, scores_path, output_path="instance_analysis.png"
    ) -> List[Dict[str, Any]]:
        """Analyze and visualize instance segmentation results."""

        # Load instance mask and scores
        with rasterio.open(instance_mask_path) as src:
            instance_mask = src.read(1)

        with rasterio.open(scores_path) as src:
            scores_mask = src.read(1)

        # Get unique instances (excluding background)
        unique_instances = np.unique(instance_mask)
        unique_instances = unique_instances[unique_instances > 0]

        # Calculate statistics for each instance
        instance_stats = []
        for instance_id in unique_instances:
            mask = instance_mask == instance_id
            area = np.sum(mask)
            score = np.mean(scores_mask[mask])
            instance_stats.append({"id": instance_id, "area": area, "score": score})

        # Sort by score
        instance_stats.sort(key=lambda x: x["score"], reverse=True)

        # Create visualization
        fig, axes = plt.subplots(2, 2, figsize=(16, 12))

        # 1. Instance segmentation visualization
        colored_mask = np.zeros((*instance_mask.shape, 3), dtype=np.uint8)
        colors = plt.cm.Set3(np.linspace(0, 1, len(unique_instances)))

        for i, instance_id in enumerate(unique_instances):
            mask = instance_mask == instance_id
            colored_mask[mask] = (colors[i][:3] * 255).astype(np.uint8)

        axes[0, 0].imshow(colored_mask)
        axes[0, 0].set_title(
            f"Instance Segmentation\n({len(unique_instances)} instances)",
            fontweight="bold",
        )
        axes[0, 0].axis("off")

        # 2. Scores heatmap
        im = axes[0, 1].imshow(scores_mask, cmap="viridis", vmin=0, vmax=1)
        axes[0, 1].set_title("Instance Confidence Scores", fontweight="bold")
        axes[0, 1].axis("off")
        plt.colorbar(im, ax=axes[0, 1], shrink=0.8)

        # 3. Score distribution
        all_scores = [stat["score"] for stat in instance_stats]
        axes[1, 0].hist(
            all_scores, bins=20, alpha=0.7, color="skyblue", edgecolor="black"
        )
        axes[1, 0].axvline(
            x=np.mean(all_scores),
            color="red",
            linestyle="--",
            label=f"Mean: {np.mean(all_scores):.3f}",
        )
        axes[1, 0].set_xlabel("Confidence Score")
        axes[1, 0].set_ylabel("Instance Count")
        axes[1, 0].set_title("Score Distribution", fontweight="bold")
        axes[1, 0].legend()
        axes[1, 0].grid(True, alpha=0.3)

        # 4. Top instances by score
        top_instances = instance_stats[:10]
        instance_ids = [stat["id"] for stat in top_instances]
        scores = [stat["score"] for stat in top_instances]
        areas = [stat["area"] for stat in top_instances]

        bars = axes[1, 1].bar(
            range(len(top_instances)), scores, color="coral", alpha=0.7
        )
        axes[1, 1].set_xlabel("Top 10 Instances")
        axes[1, 1].set_ylabel("Confidence Score")
        axes[1, 1].set_title("Top Instances by Confidence", fontweight="bold")
        axes[1, 1].set_xticks(range(len(top_instances)))
        axes[1, 1].set_xticklabels([f"#{id}" for id in instance_ids], rotation=45)

        # Add area info as text on bars
        for i, (bar, area) in enumerate(zip(bars, areas)):
            height = bar.get_height()
            axes[1, 1].text(
                bar.get_x() + bar.get_width() / 2.0,
                height,
                f"{area}px",
                ha="center",
                va="bottom",
                fontsize=8,
            )

        plt.tight_layout()
        plt.savefig(output_path, dpi=150, bbox_inches="tight")
        plt.show()

        # Print summary statistics
        logger.info(f"Instance Analysis Summary:")
        logger.info(f"   Total instances: {len(unique_instances)}")
        logger.info(f"   Average confidence: {np.mean(all_scores):.3f}")
        logger.info(
            f"   Score range: {np.min(all_scores):.3f} - {np.max(all_scores):.3f}"
        )
        logger.info(f"   Total change area: {sum(areas):,} pixels")

        logger.info(f"Instance analysis saved as '{output_path}'")

        return instance_stats

    def create_comprehensive_report(
        self, results_dict, output_path="comprehensive_report.png"
    ) -> None:
        """Create a comprehensive visualization report from detailed results."""

        if not results_dict or "masks" not in results_dict:
            logger.warning("No detailed results provided")
            return

        masks = results_dict["masks"]
        stats = results_dict["statistics"]

        # Create comprehensive visualization
        fig, axes = plt.subplots(2, 3, figsize=(18, 12))

        # 1. Score distributions
        if "iou_predictions" in stats:
            iou_scores = [
                mask["iou_pred"] for mask in masks if mask["iou_pred"] is not None
            ]
            axes[0, 0].hist(
                iou_scores, bins=20, alpha=0.7, color="lightblue", edgecolor="black"
            )
            axes[0, 0].axvline(
                x=stats["iou_predictions"]["mean"],
                color="red",
                linestyle="--",
                label=f"Mean: {stats['iou_predictions']['mean']:.3f}",
            )
            axes[0, 0].set_xlabel("IoU Score")
            axes[0, 0].set_ylabel("Count")
            axes[0, 0].set_title("IoU Predictions Distribution", fontweight="bold")
            axes[0, 0].legend()
            axes[0, 0].grid(True, alpha=0.3)

        # 2. Stability scores
        if "stability_scores" in stats:
            stability_scores = [
                mask["stability_score"]
                for mask in masks
                if mask["stability_score"] is not None
            ]
            axes[0, 1].hist(
                stability_scores,
                bins=20,
                alpha=0.7,
                color="lightgreen",
                edgecolor="black",
            )
            axes[0, 1].axvline(
                x=stats["stability_scores"]["mean"],
                color="red",
                linestyle="--",
                label=f"Mean: {stats['stability_scores']['mean']:.3f}",
            )
            axes[0, 1].set_xlabel("Stability Score")
            axes[0, 1].set_ylabel("Count")
            axes[0, 1].set_title("Stability Scores Distribution", fontweight="bold")
            axes[0, 1].legend()
            axes[0, 1].grid(True, alpha=0.3)

        # 3. Change confidence
        if "change_confidence" in stats:
            change_conf = [
                mask["change_confidence"]
                for mask in masks
                if mask["change_confidence"] is not None
            ]
            axes[0, 2].hist(
                change_conf, bins=20, alpha=0.7, color="lightyellow", edgecolor="black"
            )
            axes[0, 2].axvline(
                x=stats["change_confidence"]["mean"],
                color="red",
                linestyle="--",
                label=f"Mean: {stats['change_confidence']['mean']:.1f}",
            )
            axes[0, 2].set_xlabel("Change Confidence")
            axes[0, 2].set_ylabel("Count")
            axes[0, 2].set_title("Change Confidence Distribution", fontweight="bold")
            axes[0, 2].legend()
            axes[0, 2].grid(True, alpha=0.3)

        # 4. Area distribution
        if "areas" in stats:
            areas = [mask["area"] for mask in masks if mask["area"] is not None]
            axes[1, 0].hist(
                areas, bins=20, alpha=0.7, color="lightcoral", edgecolor="black"
            )
            axes[1, 0].axvline(
                x=stats["areas"]["mean"],
                color="red",
                linestyle="--",
                label=f"Mean: {stats['areas']['mean']:.1f}",
            )
            axes[1, 0].set_xlabel("Area (pixels)")
            axes[1, 0].set_ylabel("Count")
            axes[1, 0].set_title("Area Distribution", fontweight="bold")
            axes[1, 0].legend()
            axes[1, 0].grid(True, alpha=0.3)

        # 5. Combined confidence vs area scatter
        combined_conf = [
            mask["combined_confidence"]
            for mask in masks
            if "combined_confidence" in mask
        ]
        areas_for_scatter = [
            mask["area"]
            for mask in masks
            if "combined_confidence" in mask and mask["area"] is not None
        ]

        if combined_conf and areas_for_scatter:
            scatter = axes[1, 1].scatter(
                areas_for_scatter,
                combined_conf,
                alpha=0.6,
                c=combined_conf,
                cmap="viridis",
                s=50,
            )
            axes[1, 1].set_xlabel("Area (pixels)")
            axes[1, 1].set_ylabel("Combined Confidence")
            axes[1, 1].set_title("Confidence vs Area", fontweight="bold")
            axes[1, 1].grid(True, alpha=0.3)
            plt.colorbar(scatter, ax=axes[1, 1], shrink=0.8)

        # 6. Summary statistics text
        summary_text = f"""Detection Summary:
Total Instances: {len(masks)}
Processing Size: {results_dict['summary']['target_size']}
Original Shape: {results_dict['summary']['original_shape']}

Quality Metrics:"""

        if "iou_predictions" in stats:
            summary_text += f"""
IoU Predictions:
  Mean: {stats['iou_predictions']['mean']:.3f}
  Range: {stats['iou_predictions']['min']:.3f} - {stats['iou_predictions']['max']:.3f}"""

        if "stability_scores" in stats:
            summary_text += f"""
Stability Scores:
  Mean: {stats['stability_scores']['mean']:.3f}
  Range: {stats['stability_scores']['min']:.3f} - {stats['stability_scores']['max']:.3f}"""

        if "change_confidence" in stats:
            summary_text += f"""
Change Confidence:
  Mean: {stats['change_confidence']['mean']:.1f}
  Range: {stats['change_confidence']['min']:.1f} - {stats['change_confidence']['max']:.1f}"""

        if "areas" in stats:
            summary_text += f"""
Areas:
  Mean: {stats['areas']['mean']:.1f}
  Total: {stats['areas']['total']:,.0f} pixels"""

        axes[1, 2].text(
            0.05,
            0.95,
            summary_text,
            transform=axes[1, 2].transAxes,
            fontsize=10,
            verticalalignment="top",
            fontfamily="monospace",
        )
        axes[1, 2].set_xlim(0, 1)
        axes[1, 2].set_ylim(0, 1)
        axes[1, 2].axis("off")
        axes[1, 2].set_title("Summary Statistics", fontweight="bold")

        plt.tight_layout()
        plt.suptitle(
            "Comprehensive Change Detection Report",
            fontsize=16,
            fontweight="bold",
            y=0.98,
        )
        plt.savefig(output_path, dpi=150, bbox_inches="tight")
        plt.show()

        logger.info(f"Comprehensive report saved as '{output_path}'")

    def run_complete_analysis(
        self, image1_path, image2_path, output_dir="change_detection_results"
    ) -> Dict[str, Any]:
        """Run complete change detection analysis with all outputs and visualizations."""

        # Create output directory
        os.makedirs(output_dir, exist_ok=True)

        # Define output paths
        binary_path = os.path.join(output_dir, "binary_mask.tif")
        prob_path = os.path.join(output_dir, "probability_mask.tif")
        instance_path = os.path.join(output_dir, "instance_masks.tif")

        logger.info("Running complete change detection analysis...")

        # Run detection with all outputs
        results = self.detect_changes(
            image1_path,
            image2_path,
            output_path=binary_path,
            export_probability=True,
            probability_output_path=prob_path,
            export_instance_masks=True,
            instance_masks_output_path=instance_path,
            return_detailed_results=True,
            return_results=False,
        )

        logger.info("Creating visualizations...")

        # Create all visualizations
        self.visualize_results(image1_path, image2_path, binary_path, prob_path)

        self.create_split_comparison(
            image1_path,
            image2_path,
            binary_path,
            prob_path,
            output_path=os.path.join(output_dir, "split_comparison.png"),
        )

        scores_path = instance_path.replace(".tif", "_scores.tif")
        self.analyze_instances(
            instance_path,
            scores_path,
            os.path.join(output_dir, "instance_analysis.png"),
        )

        self.create_comprehensive_report(
            results, os.path.join(output_dir, "comprehensive_report.png")
        )

        logger.info(f"Complete analysis finished! Results saved to: {output_dir}")
        return results

    def _save_instance_segmentation_masks(
        self, change_masks, output_path, transform, crs, original_shape, target_size
    ):
        """
        Save instance segmentation masks as a single GeoTIFF where each instance has a unique ID.

        Args:
            change_masks: Change detection masks (MaskData object)
            output_path (str): Output path for instance segmentation GeoTIFF
            transform: Rasterio transform
            crs: Coordinate reference system
            original_shape (tuple): Original image shape
            target_size (int): Processing target size
        """
        # Extract mask components
        mask_items = dict(change_masks.items())
        rles = mask_items.get("rles", [])

        # Create instance segmentation mask (each instance gets unique ID)
        instance_mask = np.zeros((target_size, target_size), dtype=np.uint16)

        # Process each mask and assign unique instance ID
        for instance_id, rle in enumerate(rles, start=1):
            if isinstance(rle, dict) and "size" in rle and "counts" in rle:
                try:
                    # Decode RLE to binary mask
                    size = rle["size"]
                    counts = rle["counts"]

                    mask = np.zeros(size[0] * size[1], dtype=np.uint8)
                    pos = 0
                    value = 0

                    for count in counts:
                        if pos + count <= len(mask):
                            if value == 1:
                                mask[pos : pos + count] = 1
                            pos += count
                            value = 1 - value
                        else:
                            break

                    # RLE is column-major, reshape and transpose
                    mask = mask.reshape(size).T
                    if mask.shape != (target_size, target_size):
                        continue

                    # Assign instance ID to this mask
                    instance_mask[mask.astype(bool)] = instance_id

                except (ValueError, IndexError, KeyError) as e:
                    logger.warning(f"Failed to process mask {instance_id}: {e}")
                    continue

        # Resize back to original shape if needed
        if original_shape != (target_size, target_size):
            instance_mask_resized = resize(
                instance_mask.astype(np.float32),
                original_shape,
                preserve_range=True,
                anti_aliasing=False,
                order=0,
            )
            instance_mask_final = np.round(instance_mask_resized).astype(np.uint16)
        else:
            instance_mask_final = instance_mask

        # Save as GeoTIFF
        with rasterio.open(
            output_path,
            "w",
            driver="GTiff",
            height=instance_mask_final.shape[0],
            width=instance_mask_final.shape[1],
            count=1,
            dtype=instance_mask_final.dtype,
            crs=crs,
            transform=transform,
            compress="lzw",
        ) as dst:
            dst.write(instance_mask_final, 1)

            # Add metadata
            dst.update_tags(
                description="Instance segmentation mask with unique IDs for each change object",
                total_instances=str(len(rles)),
                background_value="0",
                instance_range=f"1-{len(rles)}",
            )

        logger.info(
            f"Saved instance segmentation mask with {len(rles)} instances to {output_path}"
        )
        return len(rles)

    def _save_instance_scores_mask(
        self, change_masks, output_path, transform, crs, original_shape, target_size
    ):
        """
        Save instance scores/probability mask as a GeoTIFF where each instance has its confidence score.

        Args:
            change_masks: Change detection masks (MaskData object)
            output_path (str): Output path for instance scores GeoTIFF
            transform: Rasterio transform
            crs: Coordinate reference system
            original_shape (tuple): Original image shape
            target_size (int): Processing target size
        """
        # Extract mask components
        mask_items = dict(change_masks.items())
        rles = mask_items.get("rles", [])
        iou_preds = mask_items.get("iou_preds", None)
        stability_scores = mask_items.get("stability_score", None)
        change_confidence = mask_items.get("change_confidence", None)

        # Convert tensors to numpy if needed
        if iou_preds is not None:
            iou_preds = iou_preds.detach().cpu().numpy()
        if stability_scores is not None:
            stability_scores = stability_scores.detach().cpu().numpy()
        if change_confidence is not None:
            change_confidence = change_confidence.detach().cpu().numpy()

        # Create instance scores mask
        scores_mask = np.zeros((target_size, target_size), dtype=np.float32)

        # Process each mask and assign confidence score
        for instance_id, rle in enumerate(rles):
            if isinstance(rle, dict) and "size" in rle and "counts" in rle:
                try:
                    # Decode RLE to binary mask
                    size = rle["size"]
                    counts = rle["counts"]

                    mask = np.zeros(size[0] * size[1], dtype=np.uint8)
                    pos = 0
                    value = 0

                    for count in counts:
                        if pos + count <= len(mask):
                            if value == 1:
                                mask[pos : pos + count] = 1
                            pos += count
                            value = 1 - value
                        else:
                            break

                    # RLE is column-major, reshape and transpose
                    mask = mask.reshape(size).T
                    if mask.shape != (target_size, target_size):
                        continue

                    # Calculate combined confidence score
                    confidence_score = 0.5  # Default
                    if iou_preds is not None and instance_id < len(iou_preds):
                        iou_score = float(iou_preds[instance_id])

                        if stability_scores is not None and instance_id < len(
                            stability_scores
                        ):
                            stability_score = float(stability_scores[instance_id])

                            if change_confidence is not None and instance_id < len(
                                change_confidence
                            ):
                                change_conf = float(change_confidence[instance_id])
                                # Normalize change confidence (typically around 145 threshold)
                                change_conf_norm = max(
                                    0.0, min(1.0, abs(change_conf) / 200.0)
                                )

                                # Weighted combination of scores
                                confidence_score = (
                                    0.35 * iou_score
                                    + 0.35 * stability_score
                                    + 0.3 * change_conf_norm
                                )
                            else:
                                confidence_score = 0.5 * (iou_score + stability_score)
                        else:
                            confidence_score = iou_score

                    # Assign confidence score to this mask
                    scores_mask[mask.astype(bool)] = confidence_score

                except (ValueError, IndexError, KeyError) as e:
                    logger.warning(
                        f"Failed to process scores for mask {instance_id}: {e}"
                    )
                    continue

        # Resize back to original shape if needed
        if original_shape != (target_size, target_size):
            scores_mask_resized = resize(
                scores_mask,
                original_shape,
                preserve_range=True,
                anti_aliasing=True,
                order=1,
            )
            scores_mask_final = np.clip(scores_mask_resized, 0.0, 1.0).astype(
                np.float32
            )
        else:
            scores_mask_final = scores_mask

        # Save as GeoTIFF
        with rasterio.open(
            output_path,
            "w",
            driver="GTiff",
            height=scores_mask_final.shape[0],
            width=scores_mask_final.shape[1],
            count=1,
            dtype=scores_mask_final.dtype,
            crs=crs,
            transform=transform,
            compress="lzw",
        ) as dst:
            dst.write(scores_mask_final, 1)

            # Add metadata
            dst.update_tags(
                description="Instance scores mask with confidence values for each change object",
                total_instances=str(len(rles)),
                background_value="0.0",
                score_range="0.0-1.0",
            )

        logger.info(
            f"Saved instance scores mask with {len(rles)} instances to {output_path}"
        )
        return len(rles)

    def _extract_detailed_results(
        self, change_masks, transform, crs, original_shape, target_size
    ):
        """
        Extract detailed results from change masks.

        Args:
            change_masks: Change detection masks (MaskData object)
            transform: Rasterio transform
            crs: Coordinate reference system
            original_shape (tuple): Original image shape
            target_size (int): Processing target size

        Returns:
            dict: Detailed results with mask information and statistics
        """
        # Extract mask components
        mask_items = dict(change_masks.items())
        rles = mask_items.get("rles", [])
        iou_preds = mask_items.get("iou_preds", None)
        stability_scores = mask_items.get("stability_score", None)
        change_confidence = mask_items.get("change_confidence", None)
        areas = mask_items.get("areas", None)
        boxes = mask_items.get("boxes", None)
        points = mask_items.get("points", None)

        # Convert tensors to numpy if needed
        if iou_preds is not None:
            iou_preds = iou_preds.detach().cpu().numpy()
        if stability_scores is not None:
            stability_scores = stability_scores.detach().cpu().numpy()
        if change_confidence is not None:
            change_confidence = change_confidence.detach().cpu().numpy()
        if areas is not None:
            areas = areas.detach().cpu().numpy()
        if boxes is not None:
            boxes = boxes.detach().cpu().numpy()
        if points is not None:
            points = points.detach().cpu().numpy()

        # Calculate statistics
        results = {
            "summary": {
                "total_masks": len(rles),
                "target_size": target_size,
                "original_shape": original_shape,
                "crs": str(crs),
                "transform": transform.to_gdal(),
            },
            "statistics": {},
            "masks": [],
        }

        # Calculate statistics for each metric
        if iou_preds is not None and len(iou_preds) > 0:
            results["statistics"]["iou_predictions"] = {
                "mean": float(np.mean(iou_preds)),
                "std": float(np.std(iou_preds)),
                "min": float(np.min(iou_preds)),
                "max": float(np.max(iou_preds)),
                "median": float(np.median(iou_preds)),
            }

        if stability_scores is not None and len(stability_scores) > 0:
            results["statistics"]["stability_scores"] = {
                "mean": float(np.mean(stability_scores)),
                "std": float(np.std(stability_scores)),
                "min": float(np.min(stability_scores)),
                "max": float(np.max(stability_scores)),
                "median": float(np.median(stability_scores)),
            }

        if change_confidence is not None and len(change_confidence) > 0:
            results["statistics"]["change_confidence"] = {
                "mean": float(np.mean(change_confidence)),
                "std": float(np.std(change_confidence)),
                "min": float(np.min(change_confidence)),
                "max": float(np.max(change_confidence)),
                "median": float(np.median(change_confidence)),
            }

        if areas is not None and len(areas) > 0:
            results["statistics"]["areas"] = {
                "mean": float(np.mean(areas)),
                "std": float(np.std(areas)),
                "min": float(np.min(areas)),
                "max": float(np.max(areas)),
                "median": float(np.median(areas)),
                "total": float(np.sum(areas)),
            }

        # Extract individual mask details
        for i in range(len(rles)):
            mask_info = {
                "mask_id": i,
                "iou_pred": (
                    float(iou_preds[i])
                    if iou_preds is not None and i < len(iou_preds)
                    else None
                ),
                "stability_score": (
                    float(stability_scores[i])
                    if stability_scores is not None and i < len(stability_scores)
                    else None
                ),
                "change_confidence": (
                    float(change_confidence[i])
                    if change_confidence is not None and i < len(change_confidence)
                    else None
                ),
                "area": int(areas[i]) if areas is not None and i < len(areas) else None,
                "bbox": (
                    boxes[i].tolist() if boxes is not None and i < len(boxes) else None
                ),
                "center_point": (
                    points[i].tolist()
                    if points is not None and i < len(points)
                    else None
                ),
            }

            # Calculate combined confidence score
            if all(
                v is not None
                for v in [
                    mask_info["iou_pred"],
                    mask_info["stability_score"],
                    mask_info["change_confidence"],
                ]
            ):
                # Normalize change confidence (145 is typical threshold)
                conf_norm = max(0.0, min(1.0, mask_info["change_confidence"] / 145.0))
                combined_score = (
                    0.3 * mask_info["iou_pred"]
                    + 0.3 * mask_info["stability_score"]
                    + 0.4 * conf_norm
                )
                mask_info["combined_confidence"] = float(combined_score)

            results["masks"].append(mask_info)

        # Sort masks by combined confidence if available
        if results["masks"] and "combined_confidence" in results["masks"][0]:
            results["masks"].sort(key=lambda x: x["combined_confidence"], reverse=True)

        return results


def download_checkpoint(
    model_type: str = "vit_h", checkpoint_dir: Optional[str] = None
) -> str:
    """Download the SAM model checkpoint.

    Args:
        model_type (str, optional): The model type. Can be one of ['vit_h', 'vit_l', 'vit_b'].
            Defaults to 'vit_h'. See https://bit.ly/3VrpxUh for more details.
        checkpoint_dir (str, optional): The checkpoint_dir directory. Defaults to None,
            which uses "~/.cache/torch/hub/checkpoints".
    """

    model_types = {
        "vit_h": {
            "name": "sam_vit_h_4b8939.pth",
            "url": "https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth",
        },
        "vit_l": {
            "name": "sam_vit_l_0b3195.pth",
            "url": "https://dl.fbaipublicfiles.com/segment_anything/sam_vit_l_0b3195.pth",
        },
        "vit_b": {
            "name": "sam_vit_b_01ec64.pth",
            "url": "https://dl.fbaipublicfiles.com/segment_anything/sam_vit_b_01ec64.pth",
        },
    }

    if model_type not in model_types:
        raise ValueError(
            f"Invalid model_type: {model_type}. It must be one of {', '.join(model_types)}"
        )

    if checkpoint_dir is None:
        checkpoint_dir = os.environ.get(
            "TORCH_HOME", os.path.expanduser("~/.cache/torch/hub/checkpoints")
        )

    checkpoint = os.path.join(checkpoint_dir, model_types[model_type]["name"])
    if not os.path.exists(checkpoint):
        logger.info(f"Model checkpoint for {model_type} not found.")
        url = model_types[model_type]["url"]
        if isinstance(url, str):
            download_file(url, checkpoint)

    return checkpoint


# Available ChangeStar model variants
_import_config_lock = threading.Lock()

CHANGESTAR_MODELS = {
    "s0_s1c1_vitb": "s0_init_s1c1_changestar_vitb_1x256",
    "s0_s1c5_vitb": "s0_init_s1c5_changestar_vitb_1x256",
    "s0_s9c1_vitb": "s0_init_s9c1_changestar_vitb_1x256",
    "s0_xview2_s1c5_vitb": "s0_init_xView2_ft_s1c5_changestar_vitb_1x256",
    "s1_s1c1_vitb": "s1_init_s1c1_changestar_vitb_1x256",
    "s1_s1c1_vitl": "s1_init_s1c1_changestar_vitl_1x256",
    "s9_s9c1_vitb": "s9_init_s9c1_changestar_vitb_1x256",
}


def list_changestar_models() -> Dict[str, str]:
    """List available ChangeStar model variants.

    Returns:
        Dict[str, str]: A dictionary mapping short model names to their
            full function names in torchange.

    Example:
        >>> from geoai.change_detection import list_changestar_models
        >>> models = list_changestar_models()
        >>> for short_name, full_name in models.items():
        ...     print(f"{short_name}: {full_name}")
    """
    return dict(CHANGESTAR_MODELS)


class ChangeStarDetection:
    """Change detection using ChangeStar models from the torchange package.

    ChangeStar is a building change detection model that uses Changen2
    pre-trained weights for zero-shot or fine-tuned building change detection.
    It takes two images (before/after) and outputs change maps and semantic
    segmentation maps.

    This class wraps the torchange ChangeStar models with GeoTIFF I/O
    support, tiled processing for large rasters, and vector output
    capabilities.

    Args:
        model_name (str): Short name of the ChangeStar model variant.
            Use :func:`list_changestar_models` to see available options.
            Defaults to ``"s1_s1c1_vitb"``.
        device (str or None): Device to use for inference (``"cuda"``,
            ``"cpu"``, or ``None`` for auto-detection). Defaults to ``None``.

    Raises:
        ImportError: If ``torchange`` or ``albumentations`` is not installed.
        ValueError: If an invalid ``model_name`` is provided.

    Example:
        >>> from geoai.change_detection import ChangeStarDetection
        >>> detector = ChangeStarDetection(model_name="s1_s1c1_vitb")
        >>> result = detector.predict("before.tif", "after.tif")
    """

    def __init__(
        self,
        model_name: str = "s1_s1c1_vitb",
        device: Optional[str] = None,
    ) -> None:
        self.model_name = model_name
        self.device = device
        self.model = None
        self._preprocess = None
        self._init_model()

    def _init_model(self):
        """Initialize the ChangeStar model and preprocessing pipeline."""
        try:
            import torch
            from torchange.models import changen2
            import albumentations as A
            import albumentations.pytorch
        except ImportError as e:
            raise ImportError(
                "The 'torchange' and 'albumentations' packages are required "
                "for ChangeStar change detection. "
                "Please install them using: pip install torchange albumentations\n"
                f"Original error: {e}"
            ) from e

        if self.model_name not in CHANGESTAR_MODELS:
            raise ValueError(
                f"Invalid model_name: '{self.model_name}'. "
                f"Available models: {list(CHANGESTAR_MODELS.keys())}"
            )

        # Get the model constructor function
        func_name = CHANGESTAR_MODELS[self.model_name]
        model_fn = getattr(changen2, func_name)

        # Workaround for upstream bug in ever's import_config: when given an
        # absolute path without a .py extension, it splits on dots, which
        # breaks paths containing "python3.12" (split into "python3/12").
        # Fix by appending .py so import_config takes the direct-path branch.
        import ever as er

        _original_import_config = er.config.import_config

        def _patched_import_config(config_name_or_path, prefix="configs"):
            if (
                os.path.isabs(config_name_or_path)
                and not config_name_or_path.endswith((".py", ".pkl"))
                and os.path.isfile(config_name_or_path + ".py")
            ):
                config_name_or_path = config_name_or_path + ".py"
            return _original_import_config(config_name_or_path, prefix)

        with _import_config_lock:
            er.config.import_config = _patched_import_config
            try:
                self.model = model_fn()
            finally:
                er.config.import_config = _original_import_config
        self.model.eval()

        # Set device
        if self.device is None:
            try:
                import ever as er

                self._device = er.auto_device()
            except ImportError:
                self._device = torch.device(
                    "cuda" if torch.cuda.is_available() else "cpu"
                )
        else:
            self._device = torch.device(self.device)

        self.model = self.model.to(self._device)

        # Set up preprocessing
        self._preprocess = A.Compose(
            [A.Normalize(), A.pytorch.ToTensorV2()],
            additional_targets={"image2": "image"},
        )

    def _read_geotiff(
        self, path: str, window=None
    ) -> Tuple[np.ndarray, Any, Any, Tuple[int, int]]:
        """Read a GeoTIFF file and return image data with georeference info.

        Args:
            path (str): Path to the GeoTIFF file.
            window: Optional rasterio window for reading a subset.

        Returns:
            Tuple of (image_array, transform, crs, shape):
                - image_array: numpy array with shape (H, W, 3) in uint8.
                - transform: rasterio affine transform.
                - crs: coordinate reference system.
                - shape: original (height, width) of the read region.
        """
        with rasterio.open(path) as src:
            if window is not None:
                img_data = src.read(window=window)
                transform = src.window_transform(window)
            else:
                img_data = src.read()
                transform = src.transform
            crs = src.crs

            # Convert from (bands, height, width) to (height, width, bands)
            img_data = np.transpose(img_data, (1, 2, 0))

            # Use only RGB bands (first 3 channels)
            if img_data.shape[2] >= 3:
                img_data = img_data[:, :, :3]

            # Normalize to 0-255 range if needed
            if img_data.dtype != np.uint8:
                img_min = img_data.min()
                img_max = img_data.max()
                if img_max > img_min:
                    img_data = (
                        (img_data - img_min) / (img_max - img_min) * 255
                    ).astype(np.uint8)
                else:
                    img_data = np.zeros_like(img_data, dtype=np.uint8)

            shape = img_data.shape[:2]
            return img_data, transform, crs, shape

    def _read_and_align_images(
        self, image1_path: str, image2_path: str
    ) -> Tuple[np.ndarray, np.ndarray, Any, Any, Tuple[int, int]]:
        """Read and align two GeoTIFF images by their spatial overlap.

        Args:
            image1_path (str): Path to the first (before) image.
            image2_path (str): Path to the second (after) image.

        Returns:
            Tuple of (img1, img2, transform, crs, shape):
                - img1: first image as (H, W, 3) uint8 array.
                - img2: second image as (H, W, 3) uint8 array.
                - transform: affine transform for the overlap area.
                - crs: coordinate reference system.
                - shape: (height, width) of the overlap area.

        Raises:
            ValueError: If the images do not overlap.
        """
        with rasterio.open(image1_path) as src1, rasterio.open(image2_path) as src2:
            bounds1 = src1.bounds
            bounds2 = src2.bounds

            left = max(bounds1.left, bounds2.left)
            bottom = max(bounds1.bottom, bounds2.bottom)
            right = min(bounds1.right, bounds2.right)
            top = min(bounds1.top, bounds2.top)

            if left >= right or bottom >= top:
                raise ValueError("Images do not overlap")

            intersection_bounds = (left, bottom, right, top)

            window1 = from_bounds(*intersection_bounds, src1.transform)
            window2 = from_bounds(*intersection_bounds, src2.transform)

            img1_data = src1.read(window=window1)
            img2_data = src2.read(window=window2)

            transform = src1.window_transform(window1)
            crs = src1.crs

            img1_data = np.transpose(img1_data, (1, 2, 0))
            img2_data = np.transpose(img2_data, (1, 2, 0))

            # Use only RGB bands
            if img1_data.shape[2] >= 3:
                img1_data = img1_data[:, :, :3]
            if img2_data.shape[2] >= 3:
                img2_data = img2_data[:, :, :3]

            # Normalize to uint8
            for arr_name in ["img1_data", "img2_data"]:
                arr = locals()[arr_name]
                if arr.dtype != np.uint8:
                    arr_min = arr.min()
                    arr_max = arr.max()
                    if arr_max > arr_min:
                        arr = ((arr - arr_min) / (arr_max - arr_min) * 255).astype(
                            np.uint8
                        )
                    else:
                        arr = np.zeros_like(arr, dtype=np.uint8)
                    if arr_name == "img1_data":
                        img1_data = arr
                    else:
                        img2_data = arr

            # Ensure same size by resizing img2 to match img1
            if img1_data.shape[:2] != img2_data.shape[:2]:
                img2_data = resize(
                    img2_data,
                    img1_data.shape[:2],
                    preserve_range=True,
                    anti_aliasing=True,
                ).astype(np.uint8)

            shape = img1_data.shape[:2]
            return img1_data, img2_data, transform, crs, shape

    def _preprocess_pair(self, img1: np.ndarray, img2: np.ndarray):
        """Preprocess a pair of images for model input.

        Args:
            img1 (np.ndarray): First image as (H, W, 3) uint8 array.
            img2 (np.ndarray): Second image as (H, W, 3) uint8 array.

        Returns:
            torch.Tensor: Concatenated and normalized tensor of shape
                (1, 6, H, W).
        """
        import torch

        data = self._preprocess(image=img1, image2=img2)
        img = torch.cat([data["image"], data["image2"]], dim=0)  # (6, H, W)
        return img.unsqueeze(0)  # (1, 6, H, W)

    def _predict_tile(self, img1: np.ndarray, img2: np.ndarray) -> Dict[str, Any]:
        """Run model inference on a single tile pair.

        Args:
            img1 (np.ndarray): First image tile as (H, W, 3) uint8 array.
            img2 (np.ndarray): Second image tile as (H, W, 3) uint8 array.

        Returns:
            Dict[str, Any]: Dictionary with keys:
                - ``"change_map"``: binary change map as (H, W) uint8 array.
                - ``"change_prob"``: change probability as (H, W) float32 array.
                - ``"t1_semantic"``: semantic segmentation for t1 as
                  (H, W) uint8 array.
                - ``"t2_semantic"``: semantic segmentation for t2 as
                  (H, W) uint8 array.
                - ``"t1_semantic_prob"``: semantic probability for t1 as
                  (H, W) float32 array.
                - ``"t2_semantic_prob"``: semantic probability for t2 as
                  (H, W) float32 array.
        """
        import torch

        input_tensor = self._preprocess_pair(img1, img2)
        input_tensor = input_tensor.to(self._device)

        with torch.no_grad():
            prediction = self.model(input_tensor)

        # Move all tensors to CPU
        prediction = {
            k: v.detach().cpu() if isinstance(v, torch.Tensor) else v
            for k, v in prediction.items()
        }

        # The model already applies logit_to_prob_() during inference,
        # so outputs are probabilities — do NOT apply sigmoid/softmax again.

        # Extract change prediction (always single-channel)
        change_prob = prediction["change_prediction"].squeeze(0)  # (1, H, W)
        change_prob = change_prob.squeeze(0).numpy()  # (H, W)
        change_map = (change_prob > 0.5).astype(np.uint8)

        # Extract semantic predictions (may be single or multi-class)
        t1_probs = prediction["t1_semantic_prediction"].squeeze(0)  # (C, H, W)
        t2_probs = prediction["t2_semantic_prediction"].squeeze(0)  # (C, H, W)

        num_classes = t1_probs.shape[0]

        if num_classes == 1:
            # Binary: squeeze class dim → (H, W)
            t1_sem_prob = t1_probs.squeeze(0).numpy()
            t2_sem_prob = t2_probs.squeeze(0).numpy()
            t1_semantic = (t1_sem_prob > 0.5).astype(np.uint8)
            t2_semantic = (t2_sem_prob > 0.5).astype(np.uint8)
        else:
            # Multi-class: argmax for labels
            t1_sem_prob = t1_probs.numpy()  # (C, H, W)
            t2_sem_prob = t2_probs.numpy()  # (C, H, W)
            t1_semantic = np.argmax(t1_sem_prob, axis=0).astype(np.uint8)  # (H, W)
            t2_semantic = np.argmax(t2_sem_prob, axis=0).astype(np.uint8)  # (H, W)

        return {
            "change_map": change_map,
            "change_prob": change_prob.astype(np.float32),
            "t1_semantic": t1_semantic,
            "t2_semantic": t2_semantic,
            "t1_semantic_prob": t1_sem_prob.astype(np.float32),
            "t2_semantic_prob": t2_sem_prob.astype(np.float32),
            "num_semantic_classes": num_classes,
        }

    def predict(
        self,
        image1_path: str,
        image2_path: str,
        output_change: Optional[str] = None,
        output_t1_semantic: Optional[str] = None,
        output_t2_semantic: Optional[str] = None,
        output_vector: Optional[str] = None,
        tile_size: int = 1024,
        overlap: int = 64,
        threshold: float = 0.5,
        **kwargs: Any,
    ) -> Dict[str, np.ndarray]:
        """Run change detection on two GeoTIFF images.

        Supports tiled processing for large rasters. When the image is
        larger than ``tile_size``, it is split into overlapping tiles,
        each tile is processed independently, and results are stitched
        back together.

        Args:
            image1_path (str): Path to the first (before) GeoTIFF image.
            image2_path (str): Path to the second (after) GeoTIFF image.
            output_change (str or None): Path to save the binary change
                mask as a GeoTIFF. Defaults to ``None``.
            output_t1_semantic (str or None): Path to save the t1 semantic
                segmentation as a GeoTIFF. Defaults to ``None``.
            output_t2_semantic (str or None): Path to save the t2 semantic
                segmentation as a GeoTIFF. Defaults to ``None``.
            output_vector (str or None): Path to save change polygons as
                a vector file (e.g., GeoJSON, GPKG). Defaults to ``None``.
            tile_size (int): Size of tiles for processing large images.
                Defaults to ``1024``.
            overlap (int): Overlap between tiles in pixels. Defaults to
                ``64``.
            threshold (float): Probability threshold for binary change map.
                Defaults to ``0.5``.
            **kwargs: Additional keyword arguments (currently unused).

        Returns:
            Dict[str, np.ndarray]: Dictionary with keys:
                - ``"change_map"``: binary change map as (H, W) uint8 array.
                - ``"change_prob"``: change probability as (H, W)
                  float32 array.
                - ``"t1_semantic"``: semantic segmentation for t1 as
                  (H, W) uint8 array.
                - ``"t2_semantic"``: semantic segmentation for t2 as
                  (H, W) uint8 array.
                - ``"t1_semantic_prob"``: semantic probability for t1 as
                  (H, W) float32 array.
                - ``"t2_semantic_prob"``: semantic probability for t2 as
                  (H, W) float32 array.

        Example:
            >>> detector = ChangeStarDetection(model_name="s1_s1c1_vitb")
            >>> result = detector.predict(
            ...     "before.tif",
            ...     "after.tif",
            ...     output_change="change_map.tif",
            ...     output_vector="changes.gpkg",
            ... )
        """
        # Read and align images
        img1, img2, transform, crs, shape = self._read_and_align_images(
            image1_path, image2_path
        )
        h, w = shape

        # Decide whether to use tiled processing
        if h <= tile_size and w <= tile_size:
            result = self._predict_tile(img1, img2)
        else:
            result = self._predict_tiled(
                img1, img2, tile_size=tile_size, overlap=overlap
            )

        # Apply threshold for change map (always binary)
        result["change_map"] = (result["change_prob"] > threshold).astype(np.uint8)

        # Semantic maps: multi-class uses argmax (already computed), binary uses threshold
        num_classes = result.get("num_semantic_classes", 1)
        if num_classes == 1:
            result["t1_semantic"] = (result["t1_semantic_prob"] > threshold).astype(
                np.uint8
            )
            result["t2_semantic"] = (result["t2_semantic_prob"] > threshold).astype(
                np.uint8
            )

        # Save outputs
        if output_change:
            self._save_raster(
                result["change_map"], output_change, transform, crs, dtype=np.uint8
            )

        if output_t1_semantic:
            self._save_raster(
                result["t1_semantic"],
                output_t1_semantic,
                transform,
                crs,
                dtype=np.uint8,
            )

        if output_t2_semantic:
            self._save_raster(
                result["t2_semantic"],
                output_t2_semantic,
                transform,
                crs,
                dtype=np.uint8,
            )

        if output_vector:
            self._save_vector(result["change_map"], output_vector, transform, crs)

        return result

    def _predict_tiled(
        self,
        img1: np.ndarray,
        img2: np.ndarray,
        tile_size: int = 1024,
        overlap: int = 64,
    ) -> Dict[str, np.ndarray]:
        """Process large images in overlapping tiles and merge results.

        Args:
            img1 (np.ndarray): First image as (H, W, 3) uint8 array.
            img2 (np.ndarray): Second image as (H, W, 3) uint8 array.
            tile_size (int): Size of each tile. Defaults to ``1024``.
            overlap (int): Number of pixels to overlap between tiles.
                Defaults to ``64``.

        Returns:
            Dict[str, np.ndarray]: Merged prediction results.
        """
        h, w = img1.shape[:2]
        stride = tile_size - overlap

        # Initialize output arrays (semantic arrays may be multi-class)
        change_prob = np.zeros((h, w), dtype=np.float32)
        count = np.zeros((h, w), dtype=np.float32)
        t1_sem_prob = None  # initialized on first tile
        t2_sem_prob = None
        num_classes = 1

        # Process tiles
        for y in range(0, h, stride):
            for x in range(0, w, stride):
                # Calculate tile bounds
                y_end = min(y + tile_size, h)
                x_end = min(x + tile_size, w)
                y_start = max(0, y_end - tile_size)
                x_start = max(0, x_end - tile_size)

                tile1 = img1[y_start:y_end, x_start:x_end]
                tile2 = img2[y_start:y_end, x_start:x_end]

                # Skip empty tiles
                if tile1.max() == 0 and tile2.max() == 0:
                    continue

                # Pad if needed
                th, tw = tile1.shape[:2]
                if th < tile_size or tw < tile_size:
                    pad_h = tile_size - th
                    pad_w = tile_size - tw
                    tile1 = np.pad(
                        tile1,
                        ((0, pad_h), (0, pad_w), (0, 0)),
                        mode="reflect",
                    )
                    tile2 = np.pad(
                        tile2,
                        ((0, pad_h), (0, pad_w), (0, 0)),
                        mode="reflect",
                    )

                tile_result = self._predict_tile(tile1, tile2)
                num_classes = tile_result.get("num_semantic_classes", 1)

                # Lazy-init semantic accumulators on first tile
                if t1_sem_prob is None:
                    if num_classes > 1:
                        t1_sem_prob = np.zeros((num_classes, h, w), dtype=np.float32)
                        t2_sem_prob = np.zeros((num_classes, h, w), dtype=np.float32)
                    else:
                        t1_sem_prob = np.zeros((h, w), dtype=np.float32)
                        t2_sem_prob = np.zeros((h, w), dtype=np.float32)

                # Crop padding if applied
                tile_change = tile_result["change_prob"][:th, :tw]

                # Accumulate change
                change_prob[y_start:y_end, x_start:x_end] += tile_change
                count[y_start:y_end, x_start:x_end] += 1

                # Accumulate semantic probabilities
                if num_classes > 1:
                    tile_t1 = tile_result["t1_semantic_prob"][:, :th, :tw]
                    tile_t2 = tile_result["t2_semantic_prob"][:, :th, :tw]
                    t1_sem_prob[:, y_start:y_end, x_start:x_end] += tile_t1
                    t2_sem_prob[:, y_start:y_end, x_start:x_end] += tile_t2
                else:
                    tile_t1 = tile_result["t1_semantic_prob"][:th, :tw]
                    tile_t2 = tile_result["t2_semantic_prob"][:th, :tw]
                    t1_sem_prob[y_start:y_end, x_start:x_end] += tile_t1
                    t2_sem_prob[y_start:y_end, x_start:x_end] += tile_t2

        # Handle case where no tiles were processed
        if t1_sem_prob is None:
            t1_sem_prob = np.zeros((h, w), dtype=np.float32)
            t2_sem_prob = np.zeros((h, w), dtype=np.float32)

        # Average overlapping regions
        mask = count > 0
        change_prob[mask] /= count[mask]
        if num_classes > 1:
            for c in range(num_classes):
                t1_sem_prob[c][mask] /= count[mask]
                t2_sem_prob[c][mask] /= count[mask]
            t1_semantic = np.argmax(t1_sem_prob, axis=0).astype(np.uint8)
            t2_semantic = np.argmax(t2_sem_prob, axis=0).astype(np.uint8)
        else:
            t1_sem_prob[mask] /= count[mask]
            t2_sem_prob[mask] /= count[mask]
            t1_semantic = (t1_sem_prob > 0.5).astype(np.uint8)
            t2_semantic = (t2_sem_prob > 0.5).astype(np.uint8)

        return {
            "change_map": (change_prob > 0.5).astype(np.uint8),
            "change_prob": change_prob,
            "t1_semantic": t1_semantic,
            "t2_semantic": t2_semantic,
            "t1_semantic_prob": t1_sem_prob,
            "t2_semantic_prob": t2_sem_prob,
            "num_semantic_classes": num_classes,
        }

    def _save_raster(
        self,
        data: np.ndarray,
        output_path: str,
        transform: Any,
        crs: Any,
        dtype=np.uint8,
    ) -> None:
        """Save a 2D array as a GeoTIFF.

        Args:
            data (np.ndarray): 2D array to save.
            output_path (str): Path for the output GeoTIFF.
            transform: Rasterio affine transform.
            crs: Coordinate reference system.
            dtype: Output data type. Defaults to ``np.uint8``.
        """
        os.makedirs(os.path.dirname(os.path.abspath(output_path)), exist_ok=True)

        with rasterio.open(
            output_path,
            "w",
            driver="GTiff",
            height=data.shape[0],
            width=data.shape[1],
            count=1,
            dtype=dtype,
            crs=crs,
            transform=transform,
            compress="lzw",
        ) as dst:
            dst.write(data.astype(dtype), 1)

    def _save_vector(
        self,
        change_map: np.ndarray,
        output_path: str,
        transform: Any,
        crs: Any,
    ) -> None:
        """Convert a binary change map to vector polygons and save.

        Args:
            change_map (np.ndarray): Binary change map as (H, W) uint8.
            output_path (str): Path for the output vector file
                (GeoJSON, GPKG, etc.).
            transform: Rasterio affine transform.
            crs: Coordinate reference system.
        """
        try:
            import geopandas as gpd
            from rasterio.features import shapes
            from shapely.geometry import shape as shapely_shape
        except ImportError as e:
            logger.warning(
                f"Cannot save vector output. "
                f"Missing dependency: {e}. "
                f"Install geopandas and shapely."
            )
            return

        os.makedirs(os.path.dirname(os.path.abspath(output_path)), exist_ok=True)

        # Extract polygons from the binary mask
        mask = change_map.astype(np.uint8)
        geometries = []
        values = []

        for geom, val in shapes(mask, transform=transform):
            if val == 1:  # Only keep change regions
                geometries.append(shapely_shape(geom))
                values.append(int(val))

        if not geometries:
            logger.info("No change polygons found.")
            return

        gdf = gpd.GeoDataFrame(
            {"change": values},
            geometry=geometries,
            crs=crs,
        )

        # Determine driver from extension
        ext = os.path.splitext(output_path)[1].lower()
        driver_map = {
            ".geojson": "GeoJSON",
            ".json": "GeoJSON",
            ".gpkg": "GPKG",
            ".shp": "ESRI Shapefile",
        }
        driver = driver_map.get(ext, "GPKG")

        gdf.to_file(output_path, driver=driver)

    def visualize(
        self,
        image1_path: str,
        image2_path: str,
        result: Optional[Dict[str, np.ndarray]] = None,
        figsize: Tuple[int, int] = (20, 5),
        title1: str = "Before (T1)",
        title2: str = "After (T2)",
        **kwargs: Any,
    ) -> plt.Figure:
        """Visualize change detection results.

        Shows the before image, after image, change map, and optionally
        the semantic segmentation maps side by side.

        Args:
            image1_path (str): Path to the first (before) image.
            image2_path (str): Path to the second (after) image.
            result (Dict or None): Pre-computed prediction result from
                :meth:`predict`. If ``None``, prediction is run first.
                Defaults to ``None``.
            figsize (Tuple[int, int]): Figure size. Defaults to
                ``(20, 5)``.
            title1 (str): Title for the before image. Defaults to
                ``"Before (T1)"``.
            title2 (str): Title for the after image. Defaults to
                ``"After (T2)"``.
            **kwargs: Additional keyword arguments passed to
                :meth:`predict`.

        Returns:
            matplotlib.figure.Figure: The figure object.

        Example:
            >>> detector = ChangeStarDetection()
            >>> fig = detector.visualize("before.tif", "after.tif")
        """
        if result is None:
            result = self.predict(image1_path, image2_path, **kwargs)

        img1, img2, _, _, _ = self._read_and_align_images(image1_path, image2_path)

        num_classes = result.get("num_semantic_classes", 1)
        sem_cmap = "gray" if num_classes == 1 else "tab10"
        sem_label = "Buildings" if num_classes == 1 else "Semantic"

        fig, axes = plt.subplots(2, 3, figsize=figsize)

        # Row 1: T1 image, T2 image
        axes[0, 0].imshow(img1)
        axes[0, 0].set_title(title1)
        axes[0, 0].axis("off")

        axes[0, 1].imshow(img2)
        axes[0, 1].set_title(title2)
        axes[0, 1].axis("off")

        fig.delaxes(axes[0, 2])

        # Row 2: T1 semantic, T2 semantic, Change Map
        axes[1, 0].imshow(result["t1_semantic"], cmap=sem_cmap)
        axes[1, 0].set_title(f"T1 {sem_label}")
        axes[1, 0].axis("off")

        axes[1, 1].imshow(result["t2_semantic"], cmap=sem_cmap)
        axes[1, 1].set_title(f"T2 {sem_label}")
        axes[1, 1].axis("off")

        axes[1, 2].imshow(result["change_map"], cmap="gray")
        axes[1, 2].set_title("Change Map")
        axes[1, 2].axis("off")

        plt.tight_layout()
        return fig

    def visualize_overlay(
        self,
        image1_path: str,
        image2_path: str,
        result: Optional[Dict[str, np.ndarray]] = None,
        figsize: Tuple[int, int] = (15, 5),
        alpha: float = 0.4,
        title1: str = "Before (T1)",
        title2: str = "After (T2)",
        **kwargs: Any,
    ) -> plt.Figure:
        """Visualize results with change map overlaid on the images.

        Args:
            image1_path (str): Path to the first (before) image.
            image2_path (str): Path to the second (after) image.
            result (Dict or None): Pre-computed prediction result from
                :meth:`predict`. If ``None``, prediction is run first.
                Defaults to ``None``.
            figsize (Tuple[int, int]): Figure size. Defaults to
                ``(15, 5)``.
            alpha (float): Transparency for the overlay. Defaults to
                ``0.4``.
            title1 (str): Title for the before image. Defaults to
                ``"Before (T1)"``.
            title2 (str): Title for the after image. Defaults to
                ``"After (T2)"``.
            **kwargs: Additional keyword arguments passed to
                :meth:`predict`.

        Returns:
            matplotlib.figure.Figure: The figure object.

        Example:
            >>> detector = ChangeStarDetection()
            >>> fig = detector.visualize_overlay("before.tif", "after.tif")
        """
        if result is None:
            result = self.predict(image1_path, image2_path, **kwargs)

        img1, img2, _, _, _ = self._read_and_align_images(image1_path, image2_path)

        fig, axes = plt.subplots(1, 3, figsize=figsize)

        # Before image with T1 semantic overlay
        num_classes = result.get("num_semantic_classes", 1)
        axes[0].imshow(img1)
        change_overlay_t1 = np.zeros((*result["t1_semantic"].shape, 4))
        change_overlay_t1[result["t1_semantic"] > 0] = [0, 0, 1, alpha]
        axes[0].imshow(change_overlay_t1)
        sem_label = "Buildings" if num_classes == 1 else "Semantic"
        axes[0].set_title(f"{title1}\n(Blue = {sem_label})")
        axes[0].axis("off")

        # After image with T2 semantic overlay
        axes[1].imshow(img2)
        change_overlay_t2 = np.zeros((*result["t2_semantic"].shape, 4))
        change_overlay_t2[result["t2_semantic"] > 0] = [0, 0, 1, alpha]
        axes[1].imshow(change_overlay_t2)
        axes[1].set_title(f"{title2}\n(Blue = {sem_label})")
        axes[1].axis("off")

        # After image with change overlay
        axes[2].imshow(img2)
        change_overlay = np.zeros((*result["change_map"].shape, 4))
        change_overlay[result["change_map"] == 1] = [1, 0, 0, alpha]
        axes[2].imshow(change_overlay)
        axes[2].set_title("Change Detection\n(Red = Change)")
        axes[2].axis("off")

        plt.tight_layout()
        return fig


def changestar_detect(
    image1_path: str,
    image2_path: str,
    model_name: str = "s1_s1c1_vitb",
    output_change: Optional[str] = None,
    output_t1_semantic: Optional[str] = None,
    output_t2_semantic: Optional[str] = None,
    output_vector: Optional[str] = None,
    tile_size: int = 1024,
    overlap: int = 64,
    threshold: float = 0.5,
    device: Optional[str] = None,
    **kwargs: Any,
) -> Dict[str, np.ndarray]:
    """Convenience function for ChangeStar building change detection.

    Creates a :class:`ChangeStarDetection` instance and runs prediction
    in one step. For repeated use, instantiate :class:`ChangeStarDetection`
    directly to avoid reloading the model.

    Args:
        image1_path (str): Path to the first (before) GeoTIFF image.
        image2_path (str): Path to the second (after) GeoTIFF image.
        model_name (str): Short name of the ChangeStar model variant.
            Use :func:`list_changestar_models` to see options.
            Defaults to ``"s1_s1c1_vitb"``.
        output_change (str or None): Path to save the binary change mask
            as a GeoTIFF. Defaults to ``None``.
        output_t1_semantic (str or None): Path to save the t1 semantic
            segmentation as a GeoTIFF. Defaults to ``None``.
        output_t2_semantic (str or None): Path to save the t2 semantic
            segmentation as a GeoTIFF. Defaults to ``None``.
        output_vector (str or None): Path to save change polygons as a
            vector file. Defaults to ``None``.
        tile_size (int): Size of tiles for processing large images.
            Defaults to ``1024``.
        overlap (int): Overlap between tiles in pixels. Defaults to
            ``64``.
        threshold (float): Probability threshold for binary change map.
            Defaults to ``0.5``.
        device (str or None): Device for inference. Defaults to ``None``
            (auto-detect).
        **kwargs: Additional keyword arguments passed to
            :meth:`ChangeStarDetection.predict`.

    Returns:
        Dict[str, np.ndarray]: Dictionary with prediction results.

    Example:
        >>> from geoai import changestar_detect
        >>> result = changestar_detect(
        ...     "before.tif",
        ...     "after.tif",
        ...     output_change="change_map.tif",
        ...     output_vector="changes.gpkg",
        ... )
    """
    detector = ChangeStarDetection(model_name=model_name, device=device)
    return detector.predict(
        image1_path,
        image2_path,
        output_change=output_change,
        output_t1_semantic=output_t1_semantic,
        output_t2_semantic=output_t2_semantic,
        output_vector=output_vector,
        tile_size=tile_size,
        overlap=overlap,
        threshold=threshold,
        **kwargs,
    )
