import numpy as np
from skimage import measure, filters, morphology
from scipy import ndimage as ndi

def strictly_non_overlapping_segmentation_median_split(
    image,
    max_amplification=30,
    step=2,
    min_area=10,
    split_threshold_ratio=0.7  # Adjust based on your data
):
    """
    Segment light spots in an image, splitting merged blobs using median intensity.
    """
    final_regions = {}
    next_id = 0

    for amp in range(-15, max_amplification + 1, step):
        # Amplify and threshold
        image_amplified = amplify_signal(image, amp)
        threshold = filters.threshold_otsu(image_amplified)
        binary = image_amplified > threshold
        binary = morphology.binary_opening(binary, morphology.disk(1))

        if not np.any(binary):
            continue

        # Initial labeling
        labeled_image, num_regions = measure.label(binary, return_num=True, connectivity=2)
        regions = measure.regionprops(labeled_image)

        # Temporary storage for this amplification level
        current_regions = {}

        for region in regions:
            if region.area < min_area:
                continue

            region_mask = labeled_image == region.label
            centroid = region.centroid

            # Extract contour
            contours = measure.find_contours(region_mask, 0.5)
            if not contours:
                continue
            contour = max(contours, key=len)

            # --- MEDIAN-BASED SPLITTING ---
            # If the blob is large, check if it should be split
            if region.area > min_area * 2:  # Only split if larger than min_area
                blob_intensities = image_amplified[region_mask]
                median_intensity = np.median(blob_intensities)

                # Apply a secondary threshold at median_intensity * split_threshold_ratio
                split_threshold = median_intensity * split_threshold_ratio
                split_mask = (image_amplified > split_threshold) & region_mask

                # Re-label the split regions
                split_labeled, _ = measure.label(split_mask, return_num=True, connectivity=2)
                split_regions = measure.regionprops(split_labeled)

                # If splitting produced multiple regions, use them instead
                if len(split_regions) > 1:
                    for split_region in split_regions:
                        if split_region.area < min_area:
                            continue

                        split_mask = split_labeled == split_region.label
                        split_contours = measure.find_contours(split_mask, 0.5)
                        if not split_contours:
                            continue
                        split_contour = max(split_contours, key=len)

                        current_regions[f"{region.label}_{split_region.label}"] = {
                            'mask': split_mask,
                            'contour': split_contour,
                            'centroid': split_region.centroid,
                            'area': split_region.area,
                            'amp': amp
                        }
                    continue  # Skip the original merged blob

            # If no splitting was needed, store the original region
            current_regions[region.label] = {
                'mask': region_mask,
                'contour': contour,
                'centroid': centroid,
                'area': region.area,
                'amp': amp
            }

        # --- REST OF YOUR EXISTING LOGIC (matching, overlap prevention, etc.) ---
        matched_pairs = []

        for curr_label, curr_data in current_regions.items():
            curr_centroid = curr_data['centroid']
            best_match = None
            min_distance = float('inf')

            for existing_id, existing_data in final_regions.items():
                exist_centroid = existing_data['centroid']
                distance = np.sqrt(
                    (curr_centroid[0] - exist_centroid[0])**2 +
                    (curr_centroid[1] - exist_centroid[1])**2
                )
                if distance < 20:
                    overlap = np.logical_and(curr_data['mask'], existing_data['mask'])
                    union = np.logical_or(curr_data['mask'], existing_data['mask'])
                    overlap_ratio = np.sum(overlap) / np.sum(union) if np.sum(union) > 0 else 0
                    if overlap_ratio > 0.2 and distance < min_distance:
                        best_match = existing_id
                        min_distance = distance

            if best_match is not None:
                matched_pairs.append((curr_label, best_match))

        # Update regions while preventing overlap
        working_mask = np.zeros_like(image, dtype=bool)
        for existing_id, data in final_regions.items():
            working_mask = np.logical_or(working_mask, data['mask'])

        for curr_label, existing_id in matched_pairs:
            curr_data = current_regions[curr_label]
            exist_data = final_regions[existing_id]
            working_mask = np.logical_and(working_mask, ~exist_data['mask'])
            would_overlap = np.any(np.logical_and(curr_data['mask'], working_mask))
            if not would_overlap:
                final_regions[existing_id] = curr_data
                working_mask = np.logical_or(working_mask, curr_data['mask'])
            else:
                working_mask = np.logical_or(working_mask, exist_data['mask'])

        matched_current_labels = [pair[0] for pair in matched_pairs]
        for curr_label, curr_data in current_regions.items():
            if curr_label in matched_current_labels:
                continue
            would_overlap = np.any(np.logical_and(curr_data['mask'], working_mask))
            if not would_overlap:
                region_id = next_id
                next_id += 1
                final_regions[region_id] = curr_data
                working_mask = np.logical_or(working_mask, curr_data['mask'])

    # Prepare final output
    final_mask = np.zeros_like(image, dtype=bool)
    polygons = []
    amplification_levels = []

    for region_id, data in final_regions.items():
        final_mask = np.logical_or(final_mask, data['mask'])
        polygons.append(data['contour'])
        amplification_levels.append(data['amp'])

    return final_mask, polygons, amplification_levels

i = 1
image = get_image(selected_images.Id.iloc[i])
_, polygons, _ = strictly_non_overlapping_segmentation_median_split(
    image,
    max_amplification=5,
    step=3,
    split_threshold_ratio=0.9999  # Adjust based on your data
)
plot_given_polygons_on_image(selected_images["Id"][i], polygons, gain=15)




the results are not bad but some of the blobs i am trying to detect are still merged at the first step even with very high split threshold. And some blobs that were detected before using this new threshold are not detected anymore