#!/usr/bin/env python

# 2022-04: Python2 to Python3 conversion by Allison Hsiang

import cv2
import numpy as np

from PIL import Image, ImageEnhance
from scipy import ndimage
from skimage import morphology


def morphOpen(image, disk_size):
    '''
    Performs morphological opening (erosion followed by dilation using same
    structuring element); equivalent of 'imopen' in MATLAB.
    '''
    strel = morphology.disk(disk_size)
    erosion = cv2.erode(image,strel)
    dilation = cv2.dilate(erosion,strel)
    return dilation


def rgb2gray(image):
    '''
    Converts image in RGB colorspace to grayscale.
    '''
    gray = cv2.cvtColor(image,cv2.COLOR_RGB2GRAY)
    return gray


def rgb2gray_luminosity(image, maxBand):
    '''
    Converts image from RGB to grayscale using luminosity method.
    See: https://mmuratarat.github.io/2020-05-13/rgb_to_grayscale_formulas
    '''
    channels = (image[:,:,0], image[:,:,1], image[:,:,2])
    #gray = 0.2126 * R + 0.7152 * G + 0.0722 * B
    #gray = gray.astype(np.uint8)
    gray = channels[maxBand]

    return gray


def rgbFilter(image):
    '''
    Takes an image in RGB colorspace, finds the strongest RGB color band (i.e.,
    the band with the highest value), and discards the other color bands.
    '''
    rgbBands = [image[:,:,i] for i in range(3)]
    # Determine most intense color band
    maxBand = np.argmax([x.sum() for x in rgbBands])
    # Discard less intense bands
    for i in range(3):
        if not i == maxBand:
            image[:,:,i] = 0
    return image, maxBand


def contrastAdjust(image, contrast_adjustment):
    '''
    Enchances contrast in image.
    '''
    if not contrast_adjustment == 0:
        newImg = Image.fromarray(image)
        contrast = ImageEnhance.Contrast(newImg)
        adjusted = contrast.enhance(contrast_adjustment)
        adjArray = np.array(adjusted)
        adjArray[adjArray < 50] = 0 # Remove low values
        return adjArray
    else:
        return image


def convertBW(image, image_name, threshold_adjustment):
    '''
    Converts a grayscale image to black and white.
    '''
    threshold,image_bw = cv2.threshold(image, 0, 255, cv2.THRESH_BINARY+cv2.THRESH_OTSU)
    if threshold + threshold_adjustment >= 0:
        image_bw = cv2.threshold(image,threshold + threshold_adjustment,255,cv2.THRESH_BINARY)[1]
    else:
        print ('INFO: Threshold adjustment value too large! Using automatic value for', image_file)
    return image_bw


def adaptiveThreshold(image, adaptive_kernel):
    '''
    Applies adaptive Gaussian thresholding.
    '''
    midpoint = int((image.shape[1]-image.shape[0])/2)
    pixel_neighborhood_size = image.shape[0]+midpoint+1 if (image.shape[0]+midpoint) % 2 == 0 else image.shape[0]+midpoint

    #se = cv2.getStructuringElement(cv2.MORPH_RECT, (adaptive_kernel,adaptive_kernel))
    se = cv2.getStructuringElement(cv2.MORPH_RECT, (int(image.shape[0]/5),int(image.shape[0]/5)))
    bg = cv2.morphologyEx(image, cv2.MORPH_DILATE, se)
    out_gray = cv2.divide(image, bg, scale=255)
    image_gauss = cv2.adaptiveThreshold(out_gray.astype('uint8'),
                                     255,
                                     cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
                                     cv2.THRESH_BINARY,
                                     pixel_neighborhood_size, #21
                                     4)

    image_gauss = cv2.morphologyEx(image_gauss, cv2.MORPH_OPEN, np.ones((5,5), np.uint8)) # Morphological opening to smooth edges
    return image_gauss


def fillHoles(image):
    '''
    Files holes contained within a binary object.
    '''
    image_filled = ndimage.morphology.binary_fill_holes(image/255.0).astype("uint8") * 255
    return image_filled


def clearBorder(image, size):
    '''
    Deletes binary objects that touch the border of the image.
    '''
    h,w = size[:2]
    image_noise = np.array(image.copy())
    mask = np.zeros((h+2,w+2),np.uint8)

    # Loop through border pixels as seeds
    for x in range(w):
        temp = cv2.floodFill(image_noise,mask,(x,0),0)
        temp = cv2.floodFill(image_noise,mask,(x,h-1),0)
    for y in range(h):
        temp = cv2.floodFill(image_noise,mask,(0,y),0)
        temp = cv2.floodFill(image_noise,mask,(w-1,y),0)

    image_border = np.asarray(image_noise)
    return image_border


def removeNoise(image, type):
    '''
    Removes all noise from a binary image by deleting all contained areas except
    the largest (simple) or by deleting all contained areas that are smaller
    than 1% the size of the largest object (relative).
    '''
    labels,area_index = ndimage.label(image)
    sizes = ndimage.sum(image,labels,range(1,area_index + 1))
    max_patch = np.where(sizes == sizes.max())[0] + 1
    max_index = np.zeros(area_index + 1, np.uint8)
    max_index[max_patch] = 1
    max_object = max_index[labels]

    if type == 'simple':
        image_clean = max_object * 255
    if type == 'relative':
        thresh_size = sizes.max() * 0.01
        keep_patches = np.where(sizes > thresh_size)[0] + 1
        keep_indices = np.zeros(area_index + 1, np.uint8)
        keep_indices[keep_patches] = 1
        keep_objects = keep_indices[labels]
        image_clean = keep_objects * 255

    return image_clean


def smoothImage(image, disk_size):
    '''
    Smooths the outline of a binary object by using morphological opening with
    a user-defined kernel size.
    '''
    smoothed = ndimage.binary_opening(image,structure=np.ones((disk_size,disk_size))).astype('uint8') * 255
    return smoothed


""" # Unused functions from old pipeline
def genGammaMap(gamma):
    invGamma = gamma ** -1
    gamma_map = np.array([((i / 255.0) ** invGamma) * 255 for i in np.arange(0,256)]).astype("uint8")
    return gamma_map


def gammaFilter(image,gamma):
    gamma_map = genGammaMap(gamma)
    return cv2.LUT(image,gamma_map)


# Threshold value version; newer version keeps largest labeled area
def removeNoise(image,size,noise_limit):
    h,w = size[:2]

    labels,areaIndex = ndimage.label(image)
    areas = np.array(ndimage.sum(image,labels,np.arange(labels.max() + 1)))
    areas_scaled = [x / 255.0 for x in areas]
    mask = areas > h * w * noise_limit

    image_clean = mask[labels.ravel()].reshape(labels.shape).astype('uint8') * 255

    return image_clean
"""
