import numpy as np
import os

from skimage import measure, morphology, segmentation
from skimage.morphology import ball, disk, binary_erosion, binary_closing
from skimage.measure import label, regionprops
from skimage.filters import roberts
from skimage.segmentation import clear_border

from scipy import ndimage as ndi
import matplotlib.pyplot as plt
import scipy.misc

import config


# Optimal threshold, found in an article about segmentation algorithm using
# morphological operations. This is in HU units.
threshold = -420
nodule_threshold = -180


class SegmentationAlgorithm(object):
    def __init__(self, threshold):
        self._threshold = threshold

    def get_segmented_lungs(self, plot=False):
        pass

    def apply_threshold(self, scan):
        scan[scan < nodule_threshold] = 0
        return scan

    def get_lung_nodules_candidates(self, patient_imgs):
        nodules = [self.apply_threshold(scan) for scan in patient_imgs]
        return np.stack([nodule for nodule in nodules if nodule.any()])

    def get_slices_with_nodules(self, patient_imgs):
        return np.stack([slice_ for slice_ in patient_imgs if 
            self._has_nodule(slice_)])

    def _has_nodule(self, scan):
        scan_copy = scan.copy()
        scan_copy[scan_copy < nodule_threshold] = 0
        return scan_copy.any()


class MorphologicalSegmentation(SegmentationAlgorithm):
    def __init__(self, threshold=-420):
        super(MorphologicalSegmentation, self).__init__(threshold)

    def get_segmented_lungs(self, im, plot=False):
        '''
        This funtion segments the lungs from the given 2D slice.
        '''
        '''
        Step 1: Convert into a binary image. 
        '''
        binary = im < self._threshold
        if plot == True:
            plt.imshow(binary, cmap=plt.cm.bone)
            plt.show()
        '''
        Step 2: Remove the blobs connected to the border of the image.
        '''
        cleared = clear_border(binary)
        if plot == True:
            plt.imshow(cleared, cmap=plt.cm.bone)
            plt.show()
        '''
        Step 3: Closure operation with a disk of radius 2. This operation is 
        to keep nodules attached to the lung wall.
        '''
        selem = disk(2)
        binary = binary_closing(cleared, selem)
        if plot == True:
            plt.imshow(binary, cmap=plt.cm.bone)
            plt.show()
        '''    
        Step 4: Label the image.
        '''
        label_image = label(binary)
        if plot == True:
            plt.imshow(label_image, cmap=plt.cm.bone)
            plt.show() 
        '''
        Step 5: Keep the labels with 2 largest areas.
        '''
        areas = [r.area for r in regionprops(label_image)]

        areas.sort()
        if len(areas) > 2:
            for region in regionprops(label_image):
                if region.area < areas[-2]:
                    for coordinates in region.coords:                
                           label_image[coordinates[0], coordinates[1]] = 0
        binary = label_image > 0
        if plot == True:
            plt.imshow(binary, cmap=plt.cm.bone)
            plt.show() 
        '''
        Step 6: Closure operation with a disk of radius 2. This operation is 
        to keep nodules attached to the lung wall.
        '''
        selem = disk(12)
        binary = binary_closing(binary, selem)
        if plot == True:
            plt.imshow(binary, cmap=plt.cm.bone)
            plt.show()
        '''
        Step 7: Fill in the small holes inside the binary mask of lungs.
        '''
        edges = roberts(binary)
        binary = ndi.binary_fill_holes(edges)
        if plot == True:
            plt.imshow(binary, cmap=plt.cm.bone)
            plt.show()
        '''
        Step 8: Erosion operation with a disk of radius 2. This operation is 
        seperate the lung nodules attached to the blood vessels.
        '''
        selem = disk(2)
        binary = binary_erosion(binary, selem)
        if plot == True:
            plt.imshow(binary, cmap=plt.cm.bone)
            plt.show() 
       
        '''
        Step 9: Superimpose the binary mask on the input image.
        '''
        get_high_vals = binary == 0
        im[get_high_vals] = 0
        if plot == True:
            plt.imshow(im, cmap=plt.cm.bone)
            plt.show() 
            
        return im


class WatershedSegmentation(SegmentationAlgorithm):
    def __init__(self, threshold=-400):
        super(WatershedSegmentation, self).__init__(threshold)

    def get_segmented_lungs(self, image, plot=False):
        # TODO: might add the logic for plotting the filters applied in the process
        #Creation of the markers as shown above:
        marker_internal, marker_external, marker_watershed = self.generate_markers(image)
        
        #Creation of the Sobel-Gradient
        sobel_filtered_dx = ndi.sobel(image, 1)
        sobel_filtered_dy = ndi.sobel(image, 0)
        sobel_gradient = np.hypot(sobel_filtered_dx, sobel_filtered_dy)
        sobel_gradient *= 255.0 / np.max(sobel_gradient)
        
        #Watershed algorithm
        watershed = morphology.watershed(sobel_gradient, marker_watershed)
        
        #Reducing the image created by the Watershed algorithm to its outline
        outline = ndi.morphological_gradient(watershed, size=(3,3))
        outline = outline.astype(bool)
        
        #Performing Black-Tophat Morphology for reinclusion
        #Creation of the disk-kernel and increasing its size a bit
        blackhat_struct = [[0, 0, 1, 1, 1, 0, 0],
                           [0, 1, 1, 1, 1, 1, 0],
                           [1, 1, 1, 1, 1, 1, 1],
                           [1, 1, 1, 1, 1, 1, 1],
                           [1, 1, 1, 1, 1, 1, 1],
                           [0, 1, 1, 1, 1, 1, 0],
                           [0, 0, 1, 1, 1, 0, 0]]
        blackhat_struct = ndi.iterate_structure(blackhat_struct, 8)
        #Perform the Black-Hat
        outline += ndi.black_tophat(outline, structure=blackhat_struct)
        
        #Use the internal marker and the Outline that was just created to generate the lungfilter
        lungfilter = np.bitwise_or(marker_internal, outline)
        #Close holes in the lungfilter
        #fill_holes is not used here, since in some slices the heart would be reincluded by accident
        lungfilter = ndi.morphology.binary_closing(lungfilter, structure=np.ones((5,5)), iterations=3)
        
        #Apply the lungfilter (the filtered areas being assigned 0 HU
        segmented = np.where(lungfilter == 1, image, np.zeros(image.shape))

        return segmented


    def generate_markers(self, image):
        #Creation of the internal Marker
        marker_internal = image < self._threshold
        marker_internal = segmentation.clear_border(marker_internal)
        marker_internal_labels = measure.label(marker_internal)
        areas = [r.area for r in measure.regionprops(marker_internal_labels)]
        areas.sort()
        if len(areas) > 2:
            for region in measure.regionprops(marker_internal_labels):
                if region.area < areas[-2]:
                    for coordinates in region.coords:                
                           marker_internal_labels[coordinates[0], coordinates[1]] = 0
        marker_internal = marker_internal_labels > 0
        #Creation of the external Marker
        external_a = ndi.binary_dilation(marker_internal, iterations=10)
        external_b = ndi.binary_dilation(marker_internal, iterations=55)
        marker_external = external_b ^ external_a
        #Creation of the Watershed Marker matrix
        marker_watershed = np.zeros(image.shape, dtype=np.int)
        marker_watershed += marker_internal * 255
        marker_watershed += marker_external * 128
        
        return marker_internal, marker_external, marker_watershed


def get_segmentation_algorithm():
    if config.SEGMENTATION_ALGO == config.MORPHOLOGICAL_OPERATIONS:
        return MorphologicalSegmentation()
    if config.SEGMENTATION_ALGO == config.WATERSHED:
        return WatershedSegmentation()

    return WatershedSegmentation()