# coding=utf-8

"""

Active contour model

"""

import cellprofiler.module
import cellprofiler.object
import cellprofiler.setting
import numpy
import scipy.ndimage
import skimage.draw
import skimage.morphology
import skimage.filters
import skimage.measure
import skimage.segmentation


DIFFERENTIAL_METHOD = "Partial differential equation Chan-Vese"
MORPH_GEODESIC_METHOD = "Morphological geodesic"
MORPH_CHAN_VESE_METHOD = "Morphological Chan-Vese"
LEVEL_SET_CIRCLE = "circle"
LEVEL_SET_CHECKERBOARD = "checkerboard"


class ActiveContourModel(cellprofiler.module.ImageSegmentation):
    module_name = "Active contour model"

    variable_revision_number = 1

    def create_settings(self):
        super(ActiveContourModel, self).create_settings()

        self.method = cellprofiler.setting.Choice(
            text="Active contour method",
            choices=[DIFFERENTIAL_METHOD,
                     MORPH_GEODESIC_METHOD,
                     MORPH_CHAN_VESE_METHOD],
            value=DIFFERENTIAL_METHOD,
        )

        # Shared by all three methods
        self.iterations = cellprofiler.setting.Integer(
            text="Iterations",
            value=20
        )

        # Only shared by PDE and GEODESIC
        self.threshold = cellprofiler.setting.Float(
            text="Threshold",
            value=0
        )

        # Pre-processing settings
        self.pde_alpha = cellprofiler.setting.Float(
            text="Alpha",
            value=0.2
        )

        self.advanced_settings = cellprofiler.setting.Binary(
            text="Advanced settings for PDE method",
            value=False
        )

        self.phi_bound = cellprofiler.setting.Float(
            text="Phi bound",
            value=1.2
        )

        self.pre_threshold = cellprofiler.setting.Float(
            text="Pre-thresholding factor",
            value=0.9
        )

        self.connectivity = cellprofiler.setting.Integer(
            text="Label connectivity",
            value=1
        )

        self.cfl_factor = cellprofiler.setting.Float(
            text="CFL condition maintenance factor",
            value=0.45
        )

        self.sdf_smoothing = cellprofiler.setting.Float(
            text="SDF smoothing factor",
            value=0.5
        )

        # Shared morph settings
        self.alpha = cellprofiler.setting.Float(
            text="Alpha",
            value=100.0
        )

        self.sigma = cellprofiler.setting.Float(
            text="Sigma",
            value=5.0
        )

        self.level_set = cellprofiler.setting.Choice(
            text="Initial level set",
            choices=[LEVEL_SET_CIRCLE, LEVEL_SET_CHECKERBOARD],
            value=LEVEL_SET_CIRCLE
        )

        self.adv_level_set = cellprofiler.setting.Binary(
            text="Advanced level set options",
            value=False
        )

        self.circle_center = cellprofiler.setting.Coordinates(
            text="Circle level set center",
            value=(-1,-1)
        )

        self.circle_radius = cellprofiler.setting.Float(
            text="Circle level set radius",
            value=-1.
        )

        self.checkerboard_size = cellprofiler.setting.Integer(
            text="Checkerboard level set square size",
            value=5,
            minval=0
        )

        self.smoothing = cellprofiler.setting.Integer(
            text="Smoothing",
            value=1,
            minval=0
        )

        # Geodesic settings
        self.balloon = cellprofiler.setting.Float(
            text="Ballon force",
            value=0.
        )

        # Morph Chan-Vese settings
        self.lambda1 = cellprofiler.setting.Float(
            text="Outer region weight",
            value=1.
        )

        self.lambda2 = cellprofiler.setting.Float(
            text="Inner region weight",
            value=1.
        )

    def settings(self):
        __settings__ = super(ActiveContourModel, self).settings()

        return __settings__ + [
            self.method,
            self.iterations,
            self.threshold,
            self.pde_alpha,
            self.advanced_settings,
            self.phi_bound,
            self.pre_threshold,
            self.connectivity,
            self.cfl_factor,
            self.sdf_smoothing,
            self.alpha,
            self.sigma,
            self.level_set,
            self.adv_level_set,
            self.circle_center,
            self.circle_radius,
            self.checkerboard_size,
            self.smoothing,
            self.balloon,
            self.lambda1,
            self.lambda2
        ]

    def visible_settings(self):
        __settings__ = super(ActiveContourModel, self).settings()

        # Shared by all three methods
        __settings__ += [
            self.method,
            self.iterations,
        ]

        # Add PDE settings
        if self.method.value == DIFFERENTIAL_METHOD:
            __settings__ += [
                self.pde_alpha,
                self.threshold,
                self.advanced_settings,
            ]
            if self.advanced_settings.value:
                __settings__ += [
                    self.phi_bound,
                    self.pre_threshold,
                    self.connectivity,
                    self.cfl_factor,
                    self.sdf_smoothing
                ]

        # Add common settings
        elif self.method.value in [MORPH_GEODESIC_METHOD, MORPH_CHAN_VESE_METHOD]:
            __settings__ += [
                self.alpha,
                self.sigma,
                self.level_set,
                self.adv_level_set,
            ]

            # Advanced settings for level setting
            if self.adv_level_set.value:
                if self.level_set.value == LEVEL_SET_CIRCLE:
                    __settings__ += [
                        self.circle_center,
                        self.circle_radius
                    ]
                elif self.level_set.value == LEVEL_SET_CHECKERBOARD:
                    __settings__ += [
                        self.checkerboard_size
                    ]

            # Add common smoothing function after level setting options
            __settings__ += [
                self.smoothing
            ]
            # Add individual settings
            if self.method.value == MORPH_GEODESIC_METHOD:
                __settings__ += [
                    self.threshold,
                    self.balloon
                ]
            elif self.method.value == MORPH_CHAN_VESE_METHOD:
                __settings__ += [
                    self.lambda1,
                    self.lambda2
                ]
        return __settings__

    def run(self, workspace):
        x_name = self.x_name.value

        y_name = self.y_name.value

        images = workspace.image_set

        x = images.get_image(x_name)

        x_data = x.pixel_data

        y_data = None

        if self.method.value == DIFFERENTIAL_METHOD:
            thresholding = skimage.filters.threshold_otsu(x_data)

            thresholding = thresholding * self.pre_threshold.value

            binary = x_data > thresholding

            y_data, phi = chan_vese(x_data,
                                    binary,
                                    alpha=self.pde_alpha.value,
                                    iterations=self.iterations.value,
                                    threshold=self.threshold.value,
                                    phi_bound=self.phi_bound.value,
                                    cfl_factor=self.cfl_factor.value,
                                    sdf_smoothing=self.sdf_smoothing.value)

        else:
            level_set = str(self.level_set.value)
            # Change default settings only if asked to
            if self.adv_level_set.value:

                if self.level_set.value == LEVEL_SET_CIRCLE:
                    # TODO: fix this for 3d images
                    center_x = self.circle_center.x
                    center_y = self.circle_center.y
                    radius = self.circle_radius.value
                    assert (center_x == center_y == -1) or (center_x != -1 and center_y != -1), \
                        "Coordinates must either both be -1 (default) or both positive"

                    if center_x == center_y == -1:
                        center = None
                    else:
                        if len(x_data.shape) == 3:
                            raise NotImplementedError("3D center selection not yet implemented")
                        center = (center_x, center_y)
                    level_set = skimage.segmentation.circle_level_set(x_data.shape,
                                                                      center=center,
                                                                      radius=radius)
                elif self.level_set.value == LEVEL_SET_CHECKERBOARD:
                    square_size = self.checkerboard_size.value
                    level_set = skimage.segmentation.checkerboard_level_set(x_data.shape,
                                                                            square_size=square_size)

            if self.method.value == MORPH_GEODESIC_METHOD:
                # Perform preprocessing for geodesic method
                pre_process = skimage.segmentation.inverse_gaussian_gradient(x_data,
                                                                             alpha=self.alpha.value,
                                                                             sigma=self.sigma.value)

                threshold_policy = 'auto' if self.threshold.value == 0 else self.threshold.value

                y_data = skimage.segmentation \
                    .morphological_geodesic_active_contour(pre_process,
                                                           iterations=self.iterations.value,
                                                           smoothing=self.smoothing.value,
                                                           init_level_set=level_set,
                                                           threshold=threshold_policy,
                                                           balloon=self.balloon.value
                                                           )

            elif self.method.value == MORPH_CHAN_VESE_METHOD:
                y_data = skimage.segmentation \
                    .morphological_chan_vese(x_data,
                                             iterations=self.iterations.value,
                                             smoothing=self.smoothing.value,
                                             init_level_set=level_set,
                                             lambda1=self.lambda1.value,
                                             lambda2=self.lambda2.value
                                             )

        y_data = skimage.measure.label(y_data, connectivity=self.connectivity.value)

        objects = cellprofiler.object.Objects()

        objects.segmented = y_data

        objects.parent_image = x

        workspace.object_set.add_objects(objects, y_name)

        self.add_measurements(workspace)

        if self.show_window:
            workspace.display_data.x_data = x_data

            workspace.display_data.y_data = y_data

            workspace.display_data.dimensions = x.dimensions


epsilon = numpy.finfo(numpy.float).eps


def chan_vese(image, mask, iterations, alpha, threshold, phi_bound, cfl_factor, sdf_smoothing):
    image = skimage.img_as_float(image)

    # -- Create a signed distance map (SDF) from mask
    phi = bwdist(mask) - bwdist(1 - mask) + mask - 0.5

    # --main loop
    iteration = 0

    stop = False

    previous_mask = mask

    c = 0

    while iteration < iterations and not stop:
        # get the curve's narrow band
        index = numpy.flatnonzero(numpy.logical_and(phi <= phi_bound, phi >= -phi_bound))

        if len(index) > 0:
            interior_points = numpy.flatnonzero(phi <= 0)

            exterior_points = numpy.flatnonzero(phi > 0)

            interior_mean = numpy.sum(image.flat[interior_points]) / (len(interior_points) + epsilon)

            exterior_mean = numpy.sum(image.flat[exterior_points]) / (len(exterior_points) + epsilon)

            force = (image.flat[index] - interior_mean) ** 2 - (image.flat[index] - exterior_mean) ** 2

            curvature = get_curvature(phi, index)

            gradient_descent = force / numpy.max(numpy.abs(force)) + alpha * curvature

            # -- maintain the CFL condition
            dt = cfl_factor / (numpy.max(numpy.abs(gradient_descent)) + epsilon)

            # -- evolve the curve
            phi.flat[index] += dt * gradient_descent

            # -- Keep SDF smooth
            phi = sussman(phi, sdf_smoothing)

            new_mask = phi <= 0

            c = convergence(previous_mask, new_mask, threshold, c)

            if c <= 5:
                iteration += 1

                previous_mask = new_mask
            else:
                stop = True

        else:
            break

    # -- make mask from SDF
    segmentation = phi <= 0  # -- Get mask from levelset

    return segmentation, phi


def bwdist(a):
    """
    this is an intermediary function, 'a' has only True, False vals,
    so we convert them into 0, 1 values -- in reverse. True is 0,
    False is 1, distance_transform_edt wants it that way.
    """
    return scipy.ndimage.distance_transform_edt(a == 0)


# -- compute curvature along SDF
def get_curvature(phi, index):
    dimz, dimy, dimx = phi.shape
    zyx = numpy.array([numpy.unravel_index(i, phi.shape) for i in index])  # get subscripts
    z = zyx[:, 0]
    y = zyx[:, 1]
    x = zyx[:, 2]

    # -- get subscripts of neighbors
    zm1 = z - 1
    ym1 = y - 1
    xm1 = x - 1
    zp1 = z + 1
    yp1 = y + 1
    xp1 = x + 1

    # -- bounds checking
    zm1[zm1 < 0] = 0
    ym1[ym1 < 0] = 0
    xm1[xm1 < 0] = 0
    zp1[zp1 >= dimz] = dimz - 1
    yp1[yp1 >= dimy] = dimy - 1
    xp1[xp1 >= dimx] = dimx - 1

    # -- get central derivatives of SDF at x,y
    dx = (phi[z, y, xm1] - phi[z, y, xp1]) / 2  # (l-r)/2

    dxx = phi[z, y, xm1] - 2 * phi[z, y, x] + phi[z, y, xp1]  # l-2c+r

    dx2 = dx * dx

    dy = (phi[z, ym1, x] - phi[z, yp1, x]) / 2  # (u-d)/2

    dyy = phi[z, ym1, x] - 2 * phi[z, y, x] + phi[z, yp1, x]  # u-2c+d

    dy2 = dy * dy

    dz = (phi[zm1, y, x] - phi[zp1, y, x]) / 2  # (b-f)/2

    dzz = phi[zm1, y, x] - 2 * phi[z, y, x] + phi[zp1, y, x]  # b-2c+f

    dz2 = dz * dz

    # (ul+dr-ur-dl)/4
    dxy = (phi[z, ym1, xm1] + phi[z, yp1, xp1] - phi[z, ym1, xp1] - phi[z, yp1, xm1]) / 4

    # (lf+rb-rf-lb)/4
    dxz = (phi[zp1, y, xm1] + phi[zm1, y, xp1] - phi[zp1, y, xp1] - phi[zm1, y, xm1]) / 4

    # (uf+db-df-ub)/4
    dyz = (phi[zp1, ym1, x] + phi[zm1, yp1, x] - phi[zp1, yp1, x] - phi[zm1, ym1, x]) / 4

    # -- compute curvature (Kappa)
    curvature = ((dxx * (dy2 + dz2) + dyy * (dx2 + dz2) + dzz * (dx2 + dy2) - 2 * dx * dy * dxy - 2 * dx * dz * dxz - 2 * dy * dz * dyz) / (dx2 + dy2 + dz2 + epsilon))

    return curvature


def mymax(a, b):
    return (a + b + numpy.abs(a - b)) / 2


# -- level set re-initialization by the sussman method
def sussman(D, dt):
    # forward/backward differences
    a = D - shiftr(D)  # backward

    b = shiftl(D) - D  # forward

    c = D - shiftd(D)  # backward

    d = shiftu(D) - D  # forward

    e = D - shiftf(D)  # backward

    f = shiftb(D) - D  # forward

    a_p = a
    a_n = a.copy()  # a+ and a-
    b_p = b
    b_n = b.copy()
    c_p = c
    c_n = c.copy()
    d_p = d
    d_n = d.copy()
    e_p = e
    e_n = e.copy()
    f_p = f
    f_n = f.copy()

    a_p[a < 0] = 0
    a_n[a > 0] = 0
    b_p[b < 0] = 0
    b_n[b > 0] = 0
    c_p[c < 0] = 0
    c_n[c > 0] = 0
    d_p[d < 0] = 0
    d_n[d > 0] = 0

    dD = numpy.zeros(D.shape)
    D_neg_ind = numpy.flatnonzero(D < 0)
    D_pos_ind = numpy.flatnonzero(D > 0)

    dD.flat[D_pos_ind] = numpy.sqrt(mymax(a_p.flat[D_pos_ind] ** 2, b_n.flat[D_pos_ind] ** 2)
                                    + mymax(c_p.flat[D_pos_ind] ** 2, d_n.flat[D_pos_ind] ** 2)
                                    + mymax(e_p.flat[D_pos_ind] ** 2, f_n.flat[D_pos_ind] ** 2)
                                    ) - 1

    dD.flat[D_neg_ind] = numpy.sqrt(mymax(a_n.flat[D_neg_ind] ** 2, b_p.flat[D_neg_ind] ** 2)
                                    + mymax(c_n.flat[D_neg_ind] ** 2, d_p.flat[D_neg_ind] ** 2)
                                    + mymax(e_n.flat[D_neg_ind] ** 2, f_p.flat[D_neg_ind] ** 2)
                                    ) - 1

    D = D - dt * numpy.sign(D) * dD

    return D


# -- whole matrix derivatives
def shiftd(m):
    return m[:, range(1, m.shape[1]) + [m.shape[1] - 1], :]


def shiftl(m):
    return m[:, :, range(1, m.shape[2]) + [m.shape[2] - 1]]


def shiftr(m):
    return m[:, :, [0] + range(0, m.shape[2] - 1)]


def shiftu(m):
    return m[:, [0] + range(0, m.shape[1] - 1), :]


def shiftf(m):
    return m[[0] + range(0, m.shape[0] - 1), :, :]


def shiftb(m):
    return m[range(1, m.shape[0]) + [m.shape[0] - 1], :, :]


# Convergence Test
def convergence(p_mask, n_mask, thresh, c):
    if numpy.sum(numpy.abs(numpy.logical_xor(p_mask, n_mask))) < thresh:
        c += 1
    else:
        c = 0

    return c

