"""Cell segmentation from force map for background exclusion."""
import numpy as np
from scipy.ndimage import gaussian_filter
from skimage.filters import threshold_otsu
from skimage.morphology import closing, opening, dilation, remove_small_objects, disk
from skimage.measure import label, regionprops


def estimate_cell_mask(heatmap, sigma=2, min_size=200, exclude_full_image=True,
                       threshold_relax=0.85, dilate_radius=4, min_area_ratio=0.2):
    """
    Estimate cell region from force map using Otsu thresholding and morphological cleanup.

    Supports multiple disconnected regions (e.g., two cells): components whose area is
    at least min_area_ratio of the largest are merged into the final mask.

    Args:
        heatmap: 2D float array [0, 1] - predicted force map
        sigma: Gaussian smoothing sigma to reduce noise. Default 2.
        min_size: Minimum object size in pixels; smaller objects removed. Default 200.
        exclude_full_image: If True, exclude the largest connected component when it
            covers most of the image (>70%) and use the second largest. Default True.
        threshold_relax: Multiply Otsu threshold by this (<1 = looser, include more pixels).
            Default 0.85.
        dilate_radius: Radius to dilate mask outward to include surrounding pixels.
            Default 4.
        min_area_ratio: Include components with area >= this fraction of the largest
            component (0–1). E.g. 0.2 = include regions at least 20% the size of the
            largest. Handles multiple disconnected force regions. Default 0.2.

    Returns:
        mask: Binary uint8 array, 1 = estimated cell, 0 = background
    """
    heatmap = np.clip(heatmap, 0, 1).astype(np.float64)
    if np.max(heatmap) <= 0:
        return np.zeros_like(heatmap, dtype=np.uint8)

    # Smooth to reduce noise
    smoothed = gaussian_filter(heatmap, sigma=sigma)

    # Otsu automatic threshold, relaxed to include more pixels
    thresh = threshold_otsu(smoothed) * threshold_relax
    mask = (smoothed > thresh).astype(np.uint8)

    # Morphological cleanup
    mask = closing(mask, disk(5)).astype(np.uint8)
    mask = opening(mask, disk(3)).astype(np.uint8)
    mask_bool = mask.astype(bool)
    min_size_int = max(int(min_size), 0)
    # skimage >=0.26 deprecates `min_size` in favor of `max_size` (inclusive threshold).
    # Use `min_size - 1` so behavior stays equivalent to the prior strict `< min_size` rule.
    try:
        mask_bool = remove_small_objects(mask_bool, max_size=max(min_size_int - 1, 0))
    except TypeError:
        mask_bool = remove_small_objects(mask_bool, min_size=min_size_int)
    mask = mask_bool.astype(np.uint8)

    # Select component(s): optionally exclude full-image background, then merge
    # all significant components (handles multiple disconnected force regions)
    labeled = label(mask)
    props = list(regionprops(labeled))

    if len(props) == 0:
        return np.zeros_like(heatmap, dtype=np.uint8)

    props_sorted = sorted(props, key=lambda x: x.area, reverse=True)
    total_px = heatmap.shape[0] * heatmap.shape[1]

    # Skip largest if it covers most of image (likely background)
    if exclude_full_image and len(props_sorted) >= 2 and props_sorted[0].area > 0.7 * total_px:
        props_sorted = props_sorted[1:]

    # Reference area for "significant" components
    ref_area = props_sorted[0].area
    # Include all components with area >= min_area_ratio * ref_area
    labels_to_keep = [p.label for p in props_sorted if p.area >= min_area_ratio * ref_area]

    mask = np.zeros_like(labeled, dtype=np.uint8)
    for lab in labels_to_keep:
        mask[labeled == lab] = 1

    # Dilate to include surrounding pixels
    if dilate_radius > 0:
        mask = dilation(mask, disk(dilate_radius)).astype(np.uint8)

    return mask