import numpy as np
from scipy.ndimage import find_objects, binary_fill_holes


def diameters(masks: np.ndarray):
    _, counts = np.unique(np.int32(masks), return_counts=True)
    counts = counts[1:]
    md = np.median(counts**0.5)
    if np.isnan(md):
        md = 0
    md /= (np.pi**0.5) / 2
    return md, counts**0.5


def fill_holes_and_remove_small_masks(mask: np.ndarray, min_size: int = 15):
    """
    Args:
        mask: Input mask of (height, width) dimensions to process
        min_size: Minimum area threshold allowed for the objects contained in the mask.

    Returns:
        The input mask without any object of area lower than min_size, and with the objects filled.
    """
    j = 1
    slices = find_objects(mask)

    for i, obj_slice in enumerate(slices):
        if obj_slice is not None:
            obj_mask = mask[obj_slice] == (i + 1)
            area = obj_mask.sum()

            if min_size > 0 and area < min_size:
                mask[obj_slice][obj_mask] = 0
            elif area > 0:
                obj_mask = binary_fill_holes(obj_mask)
                mask[obj_slice][obj_mask] = j
                j += 1
    return mask