"""A module for separating solid phase from liquid tracers using image processing techniques."""

__licence__ = """
Copyright (C) 2011  www.openpiv.net

This program is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License
along with this program.  If not, see <http://www.gnu.org/licenses/>.
"""

import numpy as np
import scipy.ndimage
from skimage.morphology import disk, erosion, dilation, opening
from skimage.measure import label


def opening_method(original_image, kernel_size, iterations=1, thresh_factor=1.1):
    """Extract separated images based on particle size.

    This method uses an *erosion* filter followed by a *dilation* (aka *opening*), to remove
    small particle traces, leaving only bigger particles in the image.

    The image of small particles is also generated by using the big particles
    image as a mask on the original image. A threshold process is used to 
    intensify the edges of particles in the mask.

    Parameters
    ----------
    original_image : np.ndarray
        Original two-phase input image

    kernel_size : int
        Erosion/dilation stencil width

    iterations : int
        Erosion iterations, default = 1

    thresh_factor : float
        Used to mask big particles, default = 1.1
        Mask condition is defined as :
            pixel value > thresh_factor * local average intensity


    Returns
    -------
    big_particles_img : np.ndarray
        Extracted image of the phase with bigger particles (dispersed phase)

    small_particles_img : np.ndarray
        Extracted image of the phase with smaller particles (carrier phase)

    """
    if kernel_size % 2 != 1:
        raise Exception("kernel_size must be an odd number")

    # Apply opening filter, resulting big particles image
    kernel = disk((kernel_size - 1) / 2)
    temp_image = original_image
    for _ in range(0, iterations):
        temp_image = erosion(temp_image, kernel)
    for _ in range(0, iterations):
        temp_image = dilation(temp_image, kernel)
    big_particles_img = temp_image

    # Find local average intensity
    averaging_kernel_size = 3 * kernel_size
    local_average = scipy.ndimage.gaussian_filter(original_image, averaging_kernel_size)

    # Find big particles mask
    mask = np.where(big_particles_img >= thresh_factor * local_average, 0, 1)

    # Estimate local background intensity (used to replace big particles)
    ## Normalized convolution is used to approximate local background intensity from the original image.
    ## Plain blur would return a high background intensity near big particles, which is not desired.
    ## See http://homepages.inf.ed.ac.uk/rbf/CVonline/LOCAL_COPIES/PIRODDI1/NormConv/NormConv.html
    ## Uniform filter is used for efficiency, due to the large kernel size.
    estimated_background_intensity = np.divide(
        scipy.ndimage.uniform_filter(mask * original_image, averaging_kernel_size),
        scipy.ndimage.uniform_filter(mask.astype(np.float64), averaging_kernel_size),
        where=scipy.ndimage.uniform_filter(
            mask.astype(np.float64), averaging_kernel_size
        )
        != 0,
        out=np.zeros_like(original_image, dtype=np.float64),
    ).astype(np.uint8)

    # Generate small particles image
    small_particles_img = np.where(
        mask > 0, original_image, estimated_background_intensity
    )

    return big_particles_img, small_particles_img


def median_filter_method(original_image, kernel_size):
    """Extract separated images using a median filter
    
    Proposed by Kiger & Pan. Original paper:

    Kiger, K. T., & Pan, C. (2000). PIV technique for the simultaneous measurement of dilute two-phase flows.
    *Journal of Fluids Engineering, Transactions of the ASME,* 122(4), 811–818.
    https://doi.org/10.1115/1.1314864

    Parameters
    ----------
    original_image : np.ndarray
        Original two-phase input image
    
    kernel_size : int
        Filter stencil width (must be an odd number). Denoted by *Nf* in Kiger & Pan.

    Returns
    -------
    big_particles_img : np.ndarray
        Extracted image of the phase with bigger particles (dispersed phase)

    small_particles_img: np.ndarray
        Extracted image of the phase with smaller particles (carrier phase)
    """

    if kernel_size % 2 != 1:
        raise Exception("kernel_size must be an odd number.")

    big_particles_img = scipy.ndimage.median_filter(original_image, kernel_size)
    small_particles_img = np.where(
        original_image > big_particles_img, original_image - big_particles_img, 0
    )

    return big_particles_img, small_particles_img


def khalitov_longmire(
    original_image,
    big_particles_criteria,
    small_particles_criteria,
    blur_kernel_size=1,
    I_sat=230,
    opening_ksize=3,
):
    """Extract separated images using the method proposed by Khalitov & Longmire, 2002.
    
    For detailed information see:

    Khalitov, D., Longmire, E. Simultaneous two-phase PIV by two-parameter phase discrimination.
    *Experiments in Fluids* 32, 252–268 (2002).
    https://doi.org/10.1007/s003480100356

    Parameters
    ----------
    original_image : np.ndarray
        Original two-phase input image
    
    big_particles_criteria : {
        'min_size' : int,
        ['max_size' : int,]
        ['min_brightness' : int,]
        ['max_brightness' : int]
    }
    A dictionary defining big particles criteria. 'min_size' is mandatory.

    small_particles_criteria : {
        ['min_size' : int,]
        'max_size' : int,
        ['min_brightness' : int,]
        ['max_brightness' : int]
    }
    A dictionary defining small particles criteria. 'max_size' is mandatory.

    blur_kernel_size : int
        Stencil width for pre-processing blur. Must be an odd number.

    I_sat : int
        Saturation intensity for object pixels detection process.

    opening_ksize : int
        Stencil width for opening operation used to remove tiny regions from object pixels.
        Set to -1 to skip opening.

    Returns
    -------
    big_particles_img : np.ndarray
        Extracted image of the phase with bigger particles (dispersed phase)

    small_particles_img : np.ndarray
        Extracted image of the phase with smaller particles (carrier phase)
    """
    # Calculate object pixels using 2nd derivative criteria
    object_pixels = khalitov_longmire_get_object_pixels(
        original_image, blur_kernel_size, I_sat, opening_ksize
    )

    # Find connected components & corresponding data
    (
        N,
        labels_image,
        regions_size_array,
        regions_brightness_array,
    ) = khalitov_longmire_analyse_particle_segments(original_image, object_pixels)

    # Apply criteria to regions

    ## Big particles
    is_region_a_big_particle = regions_size_array > big_particles_criteria["min_size"]

    if "max_size" in big_particles_criteria:
        is_region_a_big_particle &= (
            regions_size_array < big_particles_criteria["max_size"]
        )

    if "min_brightness" in big_particles_criteria:
        is_region_a_big_particle &= (
            regions_brightness_array > big_particles_criteria["min_brightness"]
        )

    if "max_brightness" in big_particles_criteria:
        is_region_a_big_particle &= (
            regions_brightness_array < big_particles_criteria["max_brightness"]
        )

    ## Small particles
    is_region_a_small_particle = (
        regions_size_array < small_particles_criteria["max_size"]
    )

    if "min_size" in small_particles_criteria:
        is_region_a_small_particle &= (
            regions_size_array > small_particles_criteria["min_size"]
        )

    if "min_brightness" in small_particles_criteria:
        is_region_a_small_particle &= (
            regions_brightness_array > small_particles_criteria["min_brightness"]
        )

    if "max_brightness" in small_particles_criteria:
        is_region_a_small_particle &= (
            regions_brightness_array < small_particles_criteria["max_brightness"]
        )

    # Generate big particles mask
    big_particles_mask = np.zeros(original_image.shape, dtype=np.uint8)
    for i in range(0, N):
        if is_region_a_big_particle[i]:
            big_particles_mask[labels_image == i] = 1

    # Generate unidentifed particles mask (for efficiency reasons)
    unindentified_particles_mask = np.zeros(original_image.shape, dtype=np.uint8)
    for i in range(0, N):
        if (not is_region_a_big_particle[i]) and (not is_region_a_small_particle[i]):
            unindentified_particles_mask[labels_image == i] = 1

    # Generate small particles mask
    small_particles_mask = (
        object_pixels - big_particles_mask - unindentified_particles_mask
    )

    # Generate final grayscale images
    big_particles_img = big_particles_mask * original_image
    small_particles_img = small_particles_mask * original_image

    return big_particles_img, small_particles_img


def get_particles_size_array(
    original_image, blur_kernel_size=1, I_sat=230, opening_ksize=3
):
    """Returns the array of particle image areas in pixels.

    Used as a quick means to set size limits in Kalitov-Longmire method.

    Usage Example:

        plt.hist( get_particles_size_array(image) )
        plt.title('Particle size distribution')

    Parameters
    ----------
    original_image : np.ndarray
        Original two-phase input image

    blur_kernel_size : int
        Stencil width for pre-processing blur. Must be an odd number.

    I_sat : int
        Saturation intensity for object pixels detection process.

    opening_ksize : int
        Stencil width for opening operation used to remove tiny regions from object pixels.
        Set to -1 to skip opening.

    Returns
    -------
    size_array : np.array
        Array of length N, containing areas of particle regions number 0 to N in pixels.
    """
    # Calculate object pixels using 2nd derivative criteria
    object_pixels = khalitov_longmire_get_object_pixels(
        original_image, blur_kernel_size, I_sat, opening_ksize
    )

    # Find connected components & corresponding data
    _, _, size_array, _ = khalitov_longmire_analyse_particle_segments(
        original_image, object_pixels
    )

    return size_array


def get_size_brightness_map(
    original_image,
    blur_kernel_size=1,
    I_sat=230,
    opening_ksize=3,
    MAX_PARTICLE_SIZE=400,
):
    """Returns the size-brightness map.

    Used as an advanced means to set size and brightness limits in Kalitov-Longmire method.

    Usage Example:

        plt.imshow(im, origin='lower')
        plt.xlabel('Brightness')
        plt.ylabel('Size (px)')
        plt.title('Signal density')

    Parameters
    ----------
    original_image : np.ndarray
        Original two-phase input image

    blur_kernel_size : int
        Stencil width for pre-processing blur. Must be an odd number.

    I_sat : int
        Saturation intensity for object pixels detection process.

    opening_ksize : int
        Stencil width for opening operation used to remove tiny regions from object pixels.
        Set to -1 to skip opening.

    MAX_PARTICLE_SIZE : int
        Particle area upper limit (Y-axis max in the map) in pixels.

    Returns
    -------
    density_map : np.ndarray
        Density map (D) where D[i,j] is log10(size*brightness*number) at size=i & brightness=j.
        See Kalitov & Longmire, 2002 for more information.
    """
    # n[i,j] = Number of particles with size of i and rounded brightness of j
    n = np.zeros((MAX_PARTICLE_SIZE + 1, 256), dtype=np.uint8)

    # Get Object Pixels
    object_pixels = khalitov_longmire_get_object_pixels(
        original_image, blur_kernel_size, I_sat, opening_ksize
    )

    # Find connected components & corresponding data
    N, _, size_array, brightness_array = khalitov_longmire_analyse_particle_segments(
        original_image, object_pixels
    )

    # Round brightness
    rounded_brightness = np.uint8(np.round(brightness_array))

    # Count size-brightness
    for i in range(0, N):
        if size_array[i] >= MAX_PARTICLE_SIZE:
            continue
        n[size_array[i], rounded_brightness[i]] += 1

    # Calculate total signal density map:
    SIZE, BRIGHTNESS = np.meshgrid(range(0, 256), range(0, MAX_PARTICLE_SIZE + 1))
    with np.errstate(divide="ignore"):
        density_map = np.log10(SIZE * BRIGHTNESS * n)
    density_map[np.isneginf(density_map)] = 0
    return density_map


def khalitov_longmire_analyse_particle_segments(original_image, object_pixels):
    """Private function"""
    # Find connected components
    label_image, N = label(object_pixels, connectivity=1, return_num=True)

    # Integrate 1 over regions
    size_array = scipy.ndimage.sum(1, label_image, index=range(0, N + 1))
    size_array = np.uint16(size_array)

    # Integrate intensity over regions
    total_light_array = scipy.ndimage.sum(
        original_image, label_image, index=range(0, N + 1)
    )
    total_light_array = np.float64(total_light_array)

    brightness_array = total_light_array / size_array

    # Drop background region (index 0)
    return N, label_image - 1, size_array[1:], brightness_array[1:]


def khalitov_longmire_get_object_pixels(
    original_image, blur_kernel_size=1, I_sat=230, opening_ksize=3
):
    """Private function"""
    if blur_kernel_size % 2 != 1:
        raise Exception("Blur kernel size must be an odd number.")

    if (opening_ksize > 1) and (opening_ksize % 2 != 1):
        raise Exception("Opening kernel size must be an odd number.")

    # Pre-processing box blur
    if blur_kernel_size > 0:
        original_image = scipy.ndimage.uniform_filter(original_image, blur_kernel_size)

    # Transform to log space
    I = np.float32(original_image)
    with np.errstate(divide="ignore"):
        Ln_I = np.log(I)
    Ln_I[np.isneginf(Ln_I)] = 0

    # Calculate second directional derivatives
    X_deriv2 = scipy.ndimage.convolve(
        Ln_I, np.asarray([[0, 0, 0], [1, -2, 1], [0, 0, 0]])
    )
    Y_deriv2 = scipy.ndimage.convolve(
        Ln_I, np.asarray([[0, 1, 0], [0, -2, 0], [0, 1, 0]])
    )
    D45_deriv2 = scipy.ndimage.convolve(
        Ln_I, np.asarray([[0, 0, 1], [0, -2, 0], [1, 0, 0]])
    )  # 0.707 coeff is ignored as we're only interested in sign of the result.
    D135_deriv2 = scipy.ndimage.convolve(
        Ln_I, np.asarray([[1, 0, 0], [0, -2, 0], [0, 0, 1]])
    )  # 0.707 coeff is ignored as we're only interested in sign of the result.

    # Apply conditions for object pixels
    object_pixels = (
        (X_deriv2 < 0) & (Y_deriv2 < 0) & (D45_deriv2 < 0) & (D135_deriv2 < 0)
    ) | (I > I_sat)
    object_pixels_uint8 = np.uint8(object_pixels)

    # Apply opening filter to remove noisy dots
    if opening_ksize > 1:
        kernel = disk((opening_ksize - 1) / 2)
        object_pixels_uint8 = opening(object_pixels_uint8, kernel)

    return object_pixels_uint8
