'''<b>Smooth</b> smooths (i.e., blurs) images.
<hr>
This module allows you to smooth (blur) images, which can be helpful to
remove artifacts of a particular size.
Note that smoothing can be a time-consuming process.
'''

import numpy as np
import scipy.ndimage as scind
from centrosome.filter import median_filter, bilateral_filter, circular_average_filter
from centrosome.smooth import circular_gaussian_kernel
from centrosome.smooth import fit_polynomial
from centrosome.smooth import smooth_with_function_and_mask

import cellprofiler.cpimage as cpi
import cellprofiler.cpmodule as cpm
import cellprofiler.settings as cps
from cellprofiler.gui.help import HELP_ON_MEASURING_DISTANCES, HELP_ON_PIXEL_INTENSITIES
from cellprofiler.settings import YES, NO

from matplotlib.widgets import Slider, Button, RadioButtons

from scipy import ndimage as ndi

FIT_POLYNOMIAL = 'Fit Polynomial'
MEDIAN_FILTER = 'Median Filter'
GAUSSIAN_FILTER = 'Gaussian Filter'
SMOOTH_KEEPING_EDGES = 'Smooth Keeping Edges'
CIRCULAR_AVERAGE_FILTER = 'Circular Average Filter'
SM_TO_AVERAGE = "Smooth to Average"
REMOVE_OUTLIER = 'Remove single hot pixels'


class SmoothMultichannel(cpm.CPModule):
    module_name = 'Smooth Multichannel'
    category = "Image Processing"
    variable_revision_number = 2

    def create_settings(self):
        self.image_name = cps.ImageNameSubscriber('Select the input image', cps.NONE)

        self.filtered_image_name = cps.ImageNameProvider('Name the output image', 'FilteredImage')

        self.smoothing_method = cps.Choice(
                'Select smoothing method',
                [FIT_POLYNOMIAL, GAUSSIAN_FILTER, MEDIAN_FILTER, SMOOTH_KEEPING_EDGES, CIRCULAR_AVERAGE_FILTER,
                 SM_TO_AVERAGE, REMOVE_OUTLIER], doc="""
            This module smooths images using one of several filters.
            Fitting a polynomial
            is fastest but does not allow a very tight fit compared to the other methods:
            <ul>
            <li><i>%(FIT_POLYNOMIAL)s:</i> This method treats the intensity of the image pixels
            as a polynomial function of the x and y position of
            each pixel. It fits the intensity to the polynomial,
            <i>A x<sup>2</sup> + B y<sup>2</sup> + C xy + D x + E y + F</i>.
            This will produce a smoothed image with a single peak or trough of intensity
            that tapers off elsewhere in the image. For many microscopy images (where
            the illumination of the lamp is brightest in the center of field of view),
            this method will produce an image with a bright central region and dimmer
            edges. But, in some cases the peak/trough of the polynomial may actually
            occur outside of the image itself.</li>
            <li><i>%(GAUSSIAN_FILTER)s:</i> This method convolves the image with a Gaussian whose
            full width at half maximum is the artifact diameter entered.
            Its effect is to blur and obscure features
            smaller than the artifact diameter and spread bright or
            dim features larger than the artifact diameter.</li>
            <li><i>%(MEDIAN_FILTER)s:</i> This method finds the median pixel value within the
            artifact diameter you specify. It removes bright or dim features that are much smaller
            than the artifact diameter.</li>
            <li><i>%(SMOOTH_KEEPING_EDGES)s:</i> This method uses a bilateral filter which
            limits Gaussian smoothing across an edge while
            applying smoothing perpendicular to an edge. The effect
            is to respect edges in an image while smoothing other
            features. <i>%(SMOOTH_KEEPING_EDGES)s</i> will filter an image with reasonable
            speed for artifact diameters greater than 10 and for
            intensity differences greater than 0.1. The algorithm
            will consume more memory and operate more slowly as
            you lower these numbers.</li>
            <li><i>%(CIRCULAR_AVERAGE_FILTER)s:</i> This method convolves the image with
            a uniform circular averaging filter whose size is the artifact diameter entered. This filter is
            useful for re-creating an out-of-focus blur to an image.</li>
            <li><i>%(SM_TO_AVERAGE)s:</i> Creates a flat, smooth image where every pixel
            of the image equals the average value of the original image.</li>
            <li><i>%(REMOVE_OUTLIER)s:</i> Remove single hot pixels</li>
            </ul>""" % globals())

        self.wants_automatic_object_size = cps.Binary(
                'Calculate artifact diameter automatically?', True, doc="""
            <i>(Used only if "%(GAUSSIAN_FILTER)s", "%(MEDIAN_FILTER)s", "%(SMOOTH_KEEPING_EDGES)s",
            "%(CIRCULAR_AVERAGE_FILTER)s" is selected)</i><br>
            Select <i>%(YES)s</i> to choose an artifact diameter based on
            the size of the image. The minimum size it will choose is 30 pixels,
            otherwise the size is 1/40 of the size of the image.
            <p>Select <i>%(YES)s</i> to manually enter an artifact diameter.</p>""" % globals())

        self.object_size = cps.Float(
                'Typical artifact diameter', 16.0, doc="""
            <i>(Used only if choosing the artifact diameter automatically is set to "%(NO)s")</i><br>
            Enter the approximate diameter (in pixels) of the features to be blurred by
            the smoothing algorithm. This value is used to calculate the size of
            the spatial filter. %(HELP_ON_MEASURING_DISTANCES)s
            For most smoothing methods, selecting a
            diameter over ~50 will take substantial amounts of time to process.""" % globals())

        self.sigma_range = cps.Float(
                'Edge intensity difference', 0.1, doc="""
            <i>(Used only if "%(SMOOTH_KEEPING_EDGES)s" is selected)</i><br>
            Enter the intensity step (which indicates an edge in an image) that you want to preserve.
            Edges are locations where the intensity changes precipitously, so this
            setting is used to adjust the rough magnitude of these changes. A lower
            number will preserve weaker edges. A higher number will preserve only stronger edges.
            Values should be between zero and one. %(HELP_ON_PIXEL_INTENSITIES)s""" % globals())

        self.clip = cps.Binary(
                'Clip intensities to 0 and 1?', True, doc="""
            <i>(Used only if %(FIT_POLYNOMIAL)s is selected)</i><br>
            The <i>%(FIT_POLYNOMIAL)s</i> method is the only smoothing option that can yield
            an output image whose values are outside of the values of the
            input image. This setting controls whether to limit the image
            intensity to the 0 - 1 range used by CellProfiler.
            <p>Select <i>%(YES)s</i> to set all output image pixels less than zero to zero
            and all pixels greater than one to one. </p>
            <p>Select <i>%(NO)s</i> to
            allow values less than zero and greater than one in the output
            image.</p>""" % globals())
        
        self.outlierneighbourhood = cps.Integer(
            'Outlier neighbourhood', 3, 3, doc="""
            <i> Used only if %(REMOVE_OUTLIER)s is selected)</i><br>
            The %(REMOVE_OUTLIER)s module needs a radius that is considered
            the neightbourhood of a pixel. Note that this needs to be an odd
            number. If an even number is provided it is rounded up to the next
            odd number.""" % globals())

        self.treshold = cps.Float(
                'Outlier Treshold', 50, 0, doc="""
            <i> (Used only if %(REMOVE_OUTLIER)s is selected)</i><br>
            The %(REMOVE_OUTLIER)s modules needs a treshold to consider
            outliers. If the second heightest pixel of the neightbourhood is
            larger than this threshold, the value is equal to the second
            highest neightbour.
            """ % globals())

    def settings(self):
        return [self.image_name, self.filtered_image_name,
                self.smoothing_method, self.wants_automatic_object_size,
                self.object_size, self.sigma_range, self.clip,
                self.outlierneighbourhood, self.treshold]

    def visible_settings(self):
        result = [self.image_name, self.filtered_image_name,
                  self.smoothing_method]
        if self.smoothing_method.value not in [FIT_POLYNOMIAL, SM_TO_AVERAGE,
                                              REMOVE_OUTLIER]:
            result.append(self.wants_automatic_object_size)
            if not self.wants_automatic_object_size.value:
                result.append(self.object_size)
            if self.smoothing_method.value == SMOOTH_KEEPING_EDGES:
                result.append(self.sigma_range)
        if self.smoothing_method.value == FIT_POLYNOMIAL:
            result.append(self.clip)
        if self.smoothing_method.value == REMOVE_OUTLIER:
            result.append(self.outlierneighbourhood)
            result.append(self.treshold)
        return result

    def run_per_layer(self, image, channel):
        if channel >= 0:
            pixel_data = image.pixel_data[:,:,channel].squeeze()
        else:
            pixel_data = image.pixel_data
        mask = image.mask
        if self.wants_automatic_object_size.value:
            object_size = min(30, max(1, np.mean(pixel_data.shape) / 40))
        else:
            object_size = float(self.object_size.value)
        sigma = object_size / 2.35

        if self.smoothing_method.value == GAUSSIAN_FILTER:
            def fn(image):
                return scind.gaussian_filter(image, sigma,
                                             mode='constant', cval=0)

            output_pixels = smooth_with_function_and_mask(pixel_data, fn,
                                                          mask)
        elif self.smoothing_method.value == MEDIAN_FILTER:
            output_pixels = median_filter(pixel_data, mask,
                                          object_size / 2 + 1)
        elif self.smoothing_method.value == SMOOTH_KEEPING_EDGES:
            sigma_range = float(self.sigma_range.value)
            output_pixels = bilateral_filter(pixel_data, mask,
                                             sigma, sigma_range)
        elif self.smoothing_method.value == FIT_POLYNOMIAL:
            output_pixels = fit_polynomial(pixel_data, mask,
                                           self.clip.value)
        elif self.smoothing_method.value == CIRCULAR_AVERAGE_FILTER:
            output_pixels = circular_average_filter(pixel_data,
                                                    object_size / 2 + 1, mask)
        elif self.smoothing_method.value == SM_TO_AVERAGE:
            if image.has_mask:
                mean = np.mean(pixel_data[mask])
            else:
                mean = np.mean(pixel_data)
                output_pixels = np.ones(pixel_data.shape, pixel_data.dtype) * mean

        elif self.smoothing_method.value == REMOVE_OUTLIER:
            # TODO: implement how this deals with masks.
            nbhood = self.outlierneighbourhood.value
            output_pixels = self.remove_outlier_pixels(pixel_data,
                                                         threshold=self.treshold.value,
                                                         radius=nbhood,
                                                         mode='max')
        else:
            raise ValueError("Unsupported smoothing method: %s" %
                             self.smoothing_method.value)

        return output_pixels

    def run(self, workspace):
        image = workspace.image_set.get_image(self.image_name.value,
                                              must_be_grayscale=False)

        output_pixels =image.pixel_data.copy()

        if len(image.pixel_data.shape) ==3:
            for i in range(image.pixel_data.shape[2]):
                output_pixels[:,:,i] = self.run_per_layer(image, i)
        else:
            output_pixels = self.run_per_layer(image, -1)

        output_image = cpi.Image(output_pixels, parent_image=image)
        workspace.image_set.add(self.filtered_image_name.value,
                                output_image)
        workspace.display_data.pixel_data = image.pixel_data
        workspace.display_data.output_pixels = output_pixels

    def display(self, workspace, figure):
        figure.set_subplots((2, 1))
        ax1 = figure.subplot_imshow_grayscale(0, 0,
                                        workspace.display_data.pixel_data,
                                        "Original: %s" %
                                        self.image_name.value)
        ax2 = figure.subplot_imshow_grayscale(1, 0,
                                        workspace.display_data.output_pixels,
                                        "Filtered: %s" %
                                        self.filtered_image_name.value,
                                        sharexy=figure.subplot(0, 0))
        # import pdb; pdb.set_trace()
        #schan1 = Slider(ax1, 'channel 1', 0, nchannel1,  valinit=1)
        #schan2 = Slider(ax2, 'channel 2', 0, nchannel2, valinit=1)

        # def update(val):
        #     channel1 = schan1.val
        #     channel2 = schan2.val
        #
        # schan1.on_changed(update)
        # schan2.on_changed(update)

    def upgrade_settings(self, setting_values, variable_revision_number,
                         module_name, from_matlab):
        if variable_revision_number < 2:
            setting_values += [3 , # outlier neighbourhood
                             20]  # treshold
        return setting_values, variable_revision_number, from_matlab


    # function
    @staticmethod
    def remove_outlier_pixels(img, threshold=50, mode='median', radius=3):
        """

        >>> a = np.zeros((10, 10))
        >>> b = a.copy()
        >>> b[5,5] = 100
        >>> np.all(a == remove_outlier_pixels(b, threshold=50, mode='max',\
                                              radius=3))
        True
        """
        if (radius % 2) == 0:
           radius += 1
        mask = np.ones((radius, radius))
        mask[int((radius-1)/2),int((radius-1)/2)] = 0

        if mode == 'max':
            img_agg = ndi.generic_filter(img, np.max, footprint=mask)
        elif mode == 'median':
            img_agg = ndi.generic_filter(img, np.median, footprint=mask)
        else:
            raise('Mode must be: max or median')
        img_out = img.copy()
        imgfil = (img-img_agg) > threshold
        img_out[imgfil] = img_agg[imgfil]
        return img_out
