import numpy as np
import skimage.measure
import skimage.segmentation
import skimage.morphology
import skimage.filters
import scipy.ndimage
import utils_plots


def segment_HU_scan(x):
    mask = np.asarray(x < -350, dtype='int32')
    for iz in xrange(mask.shape[0]):
        skimage.segmentation.clear_border(mask[iz], in_place=True)
        skimage.morphology.binary_opening(mask[iz], selem=skimage.morphology.disk(5), out=mask[iz])
        if np.sum(mask[iz]):
            mask[iz] = skimage.morphology.convex_hull_image(mask[iz])
    return mask


def segment_HU_scan_frederic(x, threshold=-350):
    mask = np.copy(x)
    binary_part = mask > threshold
    selem1 = skimage.morphology.disk(8)
    selem2 = skimage.morphology.disk(2)
    selem3 = skimage.morphology.disk(13)

    for iz in xrange(mask.shape[0]):
        # fill the body part
        filled = scipy.ndimage.binary_fill_holes(binary_part[iz])  # fill body
        filled_borders_mask = skimage.morphology.binary_erosion(filled, selem1)
        mask[iz] *= filled_borders_mask


        mask[iz] = skimage.morphology.closing(mask[iz], selem2)
        mask[iz] = skimage.morphology.erosion(mask[iz], selem3)
        mask[iz] = mask[iz] < threshold

    return mask


def segment_HU_scan_elias(x, threshold=-350, pid='test', plot=False, verbose=False):
    mask = np.copy(x)
    binary_part = mask > threshold
    selem1 = skimage.morphology.disk(8)
    selem2 = skimage.morphology.disk(2)
    selem3 = skimage.morphology.disk(13)

    for iz in xrange(mask.shape[0]):
        # fill the body part
        filled = scipy.ndimage.binary_fill_holes(binary_part[iz])  # fill body
        filled_borders_mask = skimage.morphology.binary_erosion(filled, selem1)
        mask[iz] *= filled_borders_mask

        mask[iz] = skimage.morphology.closing(mask[iz], selem2)
        mask[iz] = mask[iz] < threshold

    # params
    overlap_treshold = 7
    ratio_overlap_treshold = 0.015

    #discard disconnected regions, start at the middle slice and go to the head
    for iz in range(mask.shape[0]/2, mask.shape[0]-1):
        overlap = mask[iz] * mask[iz+1] 
        label_image = skimage.measure.label(mask[iz+1])
        if verbose: 
            print 'iz', iz
        for idx, region in enumerate(skimage.measure.regionprops(label_image)):
            total_overlap = 0
            for coordinates in region.coords:                
                total_overlap += overlap[coordinates[0], coordinates[1]]
            ratio_overlap = 1.*total_overlap/region.area
            if verbose:
                print 'region', idx, ', t_overlap', total_overlap, ', r_overlap ', ratio_overlap, ', area ', region.area, ', center', np.round(region.centroid)
            if total_overlap < overlap_treshold or ratio_overlap < ratio_overlap_treshold:
                if verbose:
                    print 'region', idx, 'in slice z=', iz-1, 'has a low overlap (', total_overlap, ratio_overlap, ') and will be discarded'
                for coordinates in region.coords: 
                    mask[iz+1, coordinates[0], coordinates[1]] = 0

    #discard disconnected regions, start at the middle slice and go to the head
    for iz in range(mask.shape[0]/2,0,-1 ):
        overlap = mask[iz] * mask[iz-1] 
        label_image = skimage.measure.label(mask[iz-1])
        if verbose: 
            print 'iz', iz
        for idx, region in enumerate(skimage.measure.regionprops(label_image)):
            total_overlap = 0
            for coordinates in region.coords:                
                total_overlap += overlap[coordinates[0], coordinates[1]]
            ratio_overlap = 1.*total_overlap/region.area
            if verbose: 
                print 'region', idx, ', t_overlap', total_overlap, ', r_overlap ', ratio_overlap, ', area ', region.area, ', center', np.round(region.centroid)
            if total_overlap < overlap_treshold or ratio_overlap < ratio_overlap_treshold:
                if verbose: 
                    print 'region', idx, 'in slice z=', iz-1, 'has a low overlap (', total_overlap, ratio_overlap, ') and will be discarded'
                for coordinates in region.coords: 
                    mask[iz-1, coordinates[0], coordinates[1]] = 0


    #erode out the blood vessels and the borders of the lung for a cleaner mask
    for iz in xrange(mask.shape[0]):
        mask[iz] = skimage.morphology.binary_dilation(mask[iz], selem3)
        #mask[iz] = scipy.ndimage.binary_fill_holes(mask[iz])


    if plot:
        utils_plots.plot_all_slices(x, mask, pid, './plots/segment_HU_scan_elias/')

    return mask


def segment_HU_scan_ira(x, threshold=-350, min_area=300):
    mask = np.asarray(x < threshold, dtype='int8')

    for zi in xrange(mask.shape[0]):
        skimage.segmentation.clear_border(mask[zi, :, :], in_place=True)

    # noise reduction
    mask = skimage.morphology.binary_opening(mask, skimage.morphology.cube(2))
    mask = np.asarray(mask, dtype='int8')

    # label regions
    label_image = skimage.measure.label(mask)
    region_props = skimage.measure.regionprops(label_image)
    sorted_regions = sorted(region_props, key=lambda x: x.area, reverse=True)
    lung_label = sorted_regions[0].label
    lung_mask = np.asarray((label_image == lung_label), dtype='int8')

    # convex hull mask
    lung_mask_convex = np.zeros_like(lung_mask)
    for i in range(lung_mask.shape[2]):
        if np.any(lung_mask[:, :, i]):
            lung_mask_convex[:, :, i] = skimage.morphology.convex_hull_image(lung_mask[:, :, i])

    # old mask inside the convex hull
    mask *= lung_mask_convex
    label_image = skimage.measure.label(mask)
    region_props = skimage.measure.regionprops(label_image)
    sorted_regions = sorted(region_props, key=lambda x: x.area, reverse=True)

    for r in sorted_regions[1:]:
        if r.area > min_area:
            # make an image only containing that region
            label_image_r = label_image == r.label
            # grow the mask
            label_image_r = scipy.ndimage.binary_dilation(label_image_r,
                                                          structure=scipy.ndimage.generate_binary_structure(3, 2))
            # compute the overlap with true lungs
            overlap = label_image_r * lung_mask
            if not np.any(overlap):
                for i in range(label_image_r.shape[0]):
                    if np.any(label_image_r[i]):
                        label_image_r[i] = skimage.morphology.convex_hull_image(label_image_r[i])
                lung_mask_convex *= 1 - label_image_r

    return lung_mask_convex