import numpy as np

from .condense import condense


def width_open(im_label, width):
    """Removes thin objects from label image using maximum of distance
    transform values within each object.

    Parameters
    ----------
    im_label : array_like
        A uint32 type label image generated by segmentation methods.
    width : int
        width threshold for objects. Objects with fewer than 'Area' pixels will
        be zeroed to merge with background.

    Notes
    -----
    Objects are assumed to have positive nonzero values. A binary mask is
    generated for each object setting all other objects to the background value
    (0). The maximum chamfered distance transform value of this mask is used
    to represent the object width.

    Returns
    -------
    im_thinned : array_like
        A uint32 label where objects with pixels < Area are removed.

    See Also
    --------
    histomicstk.segmentation.label.condense,
    histomicstk.segmentation.label.shuffle,
    histomicstk.segmentation.label.split,
    histomicstk.segmentation.label.area_open

    """
    import scipy.ndimage as ndi

    # copy input image
    im_thinned = im_label.copy()

    # condense label image
    if np.unique(im_thinned).size - 1 != im_thinned.max():
        im_thinned = condense(im_thinned)

    # get locations of objects in initial label image
    Locations = ndi.find_objects(im_thinned)

    # iterate through objects, calculating distances where needed
    for i in np.arange(1, len(Locations) + 1):

        # extract object from label image
        W = im_thinned[Locations[i - 1]]

        # embed into mask with boundary
        Mask = np.zeros((W.shape[0] + 2, W.shape[1] + 2), dtype=bool)
        Mask[1:-1, 1:-1] = W == i

        # calculate distance transform of mask
        D = ndi.distance_transform_cdt(Mask, metric='taxicab')

        # get max distance
        Max = D.max()

        # zero label mask of object 'i'
        if Max < width:
            W[W == i] = 0

    # condense to fill gaps
    im_thinned = condense(im_thinned)

    return im_thinned
