"""Filter library for the application

.. moduleauthor:: xlanro00

"""

import numpy as np
import cv2 as cv
from skimage import filters as skiflt
from skimage import restoration as skirest
from skimage import morphology as skimorph
from PIL import Image, ImageFilter
import bm3d


class img_filter:
    ''' Parent class for all the filters.
    '''

    def __init__(self, img):
        '''
        :param  img: Image to be filtered
        '''
        self.img = img


# --------------------- DENOISING FILTERS ---------------------#


class gaussian(img_filter):
    '''Gaussian blur filter from scikit-image.
    Easier to use than opencv version.
    '''

    def __init__(self, img): 
        super().__init__(img)

    def apply(self, params):
        
        # Standard deviation for Gaussian kernel
        sigma = float(params["sigma"]) if params["sigma"] else 1

        self.img = skiflt.gaussian(
            self.img, sigma=sigma, preserve_range=True)
        self.img = np.uint8(self.img)


class median(img_filter):
    ''' Median blur filter from scikit-image.
    Using this over opencv version as that one is limited to 5x5 kernel.
    '''

    def __init__(self, img):
        super().__init__(img)

    def apply(self, params):
        # Size of the median filter kernel
        # Used kernel is disk of size ksize
        ksize = int(params["ksize"]) if params["ksize"] else 3

        
        self.img = skiflt.median(self.img, footprint=skimorph.disk(ksize))
        self.img = np.uint8(self.img)


class bilateral(img_filter):
    ''' Bilateral filter from opencv.
    '''

    def __init__(self, img):
        super().__init__(img)

    def apply(self, params):

        # Diameter of pixel neighborhood used for filtering
        # This determines how fast the filtering is going to be
        diameter = int(params["diameter"]) if params["diameter"] else 3

        # Standard deviation for grayvalue/color distance
        # A larger value results in averaging of pixels with larger radiometric differences
        sigmaColor = int(params["sigmaColor"]) if params["sigmaColor"] else 75

        # Standard deviation for range distance in pixels
        # A larger value results in averaging of pixels with larger spatial differences
        sigmaSpace = int(params["sigmaSpace"]) if params["sigmaSpace"] else 75

        self.img = np.uint8(self.img)
        self.img = cv.bilateralFilter(
            self.img, diameter, sigmaColor, sigmaSpace)


class bilateral_scikit(img_filter):
    ''' Skimage denoise_bilateral filter.
    Averages pixels based on their distance and color similarity.
    Preserves edges while removing unwanted noise.
    Much slower than opencv implementation.
    '''

    def __init__(self, img):
        super().__init__(img)

    def apply(self, params):

        # Standard deviation for grayvalue/color distance.
        # A larger value results in averaging of pixels with larger radiometric differences.
        # Range of values: 0 to 1.
        sigmaColor = float(params["sigmaColor"]
                           ) if params["sigmaColor"] else 0.1

        # Standard deviation for range distance in pixels.
        # A larger value results in averaging of pixels with larger spatial differences.
        # Range of values: 0 to image size, recommend not more than 20 pixels.
        sigmaSpace = float(params["sigmaSpace"]
                           ) if params["sigmaSpace"] else 9.0

        self.img = skirest.denoise_bilateral(
            self.img, sigma_color=sigmaColor, sigma_spatial=sigmaSpace)
        self.img = np.uint8(self.img * 255.0)  # converting back to uint8


class nlmeans(img_filter):
    ''' Non-local means filter from scikit-image.
    '''

    def __init__(self, img):
        super().__init__(img)

    def apply(self, params):

        # Size of patches used for denoising
        patch_size = int(params["patch_size"]) if params["patch_size"] else 5

        # Distance in pixels where to search for patches
        patch_distance = int(params["patch_distance"]
                             ) if params["patch_distance"] else 3

        # Estimated standard deviation of the noise
        sigma = np.mean(skirest.estimate_sigma(self.img))

        # Cut-off distance, higher means more smoothed image
        h = float(params["h"])*sigma if params["h"] else 0.1*sigma

        self.img = skirest.denoise_nl_means(
            self.img, patch_size=patch_size, fast_mode=True, patch_distance=patch_distance, h=h)
        self.img = np.uint8(self.img * 255.0)  # converting back to uint8


class total_variation(img_filter):
    ''' Scikit image denoise_tv_chambolle filter from scikit-image.

    Performs total variation denoising technique based on original Chambolle paper.
    This filter removes fine detail, but preserves details such as edges.
    '''

    def __init__(self, img):
        super().__init__(img)

    def apply(self, params):

        # Denoising weight. Larger values result in more denoising.
        weight = float(params["weight"]) if params["weight"] else 0.1

        self.img = skirest.denoise_tv_chambolle(
            self.img, weight=weight)
        self.img = np.uint8(self.img * 255.0)  # converting back to uint8


class block_match(img_filter):
    '''Block matching filter from bm3d.

    This filter is very slow and should be used only on small images
    '''

    def __init__(self, img):
        super().__init__(img)

    def apply(self, params):
        sigma = float(params["sigma"]) if params["sigma"] else 20

        self.img = bm3d.bm3d(self.img, sigma_psd=sigma,
                             stage_arg=bm3d.BM3DStages.ALL_STAGES)
        self.img = np.uint8(self.img)


class unsharp_mask_scikit(img_filter):
    ''' Unsharp mask filter from scikit. 

    Apply blurring using gaussian filter, then subtract the blurred image from the original image.
    Radius parameter is the sigma parameter of the gaussian filter.
    Amount parameter regulates the strength of the unsharp mask.
    Better results than using this from opencv. 
    '''

    def __init__(self, img):
        super().__init__(img)

    def apply(self, params):
        # radius of the gaussian filter
        radius = int(params["radius"]) if params["radius"] else 3
        # strength of the unsharp mask
        amount = float(params["amount"]) if params["amount"] else 1.0

        self.img = skiflt.unsharp_mask(self.img, radius=radius,
                                       amount=amount, channel_axis=None)
        self.img = np.uint8(self.img * 255.0)  # converting back to uint


class unsharp_mask_pil(img_filter):
    ''' Unsharp mask filter from PIL.
    '''

    def __init__(self, img):
        super().__init__(img)

    def apply(self, params):

        # Blur radius
        radius = int(params["radius"]) if params["radius"] else 2

        # Unsharp strength in percent
        percent = int(params["percent"]) if params["percent"] else 150

        # Threshold controls the minimum brightness change that will be sharpened
        threshold = int(params["threshold"]) if params["threshold"] else 3

        self.img = np.uint8(self.img)
        tmp = Image.fromarray(self.img)
        tmp = tmp.filter(ImageFilter.UnsharpMask(radius, percent, threshold))
        self.img = np.uint8(np.asarray(tmp))

# ------------------- EDGE DETECTION FILTERS -------------------#


class farid(img_filter):
    ''' Farid filter from filters.
    Finds edges of the image.
    '''

    def __init__(self, img):
        super().__init__(img)

    def apply(self, _):

        self.img = skiflt.farid(self.img)
        self.img = np.uint8(self.img*255)


# ------------------ RIDGE EXTRACTION FILTERS ------------------#


class meijering(img_filter):
    ''' Meijering filter from scikit-image filters.
    Finds continuous ridges.
    '''

    def __init__(self, img):
        super().__init__(img)

    def apply(self, _):

        self.img = skiflt.meijering(self.img)
        self.img = np.uint8(self.img * 255.0)


class sato(img_filter):
    ''' Meijering filter from scikit-image filters.
    Exctracts continuous ridges.
    '''

    def __init__(self, img):
        super().__init__(img)

    def apply(self, _):

        self.img = skiflt.sato(self.img)
        self.img = np.uint8(self.img * 255.0)


class hessian(img_filter):
    ''' Hessian filter from scikit-image filters.
    '''

    def __init__(self, img):
        super().__init__(img)

    def apply(self, params):

        sigmas = float(params["sigma"]) if params["sigma"] else 1.2

        self.img = skiflt.hessian(self.img, sigmas=sigmas)
        self.img = np.uint8(self.img * 255.0)

# ------------------- MISCELLANEOUS FILTERS -------------------#


class invert(img_filter):
    ''' Invert the image using bitwise_not from opencv.
    '''

    def __init__(self, img):
        super().__init__(img)

    def apply(self, _):
        self.img = cv.bitwise_not(self.img)
        self.img = np.uint8(self.img)


class scale_values(img_filter):
    ''' Scale values of the image to use the entire range of data type.
    This should remove the line height issues.
    '''

    def __init__(self, img):
        super().__init__(img)

    def apply(self, _):
        # scale once for inverted image and once for original
        # this is done to get the full value range of the data type
        # which might help getting exact line height on stl model
        tmp = cv.bitwise_not(self.img.astype(np.uint8))
        coef = 255 / np.max(tmp)
        tmp = tmp * coef
        self.img = cv.bitwise_not(tmp.astype(np.uint8))
        coef = 255 / np.max(tmp)
        tmp = tmp * coef


class binarize(img_filter):
    ''' Binarization filter from opencv.
    '''
    def init(self, img):
        super().__init__(img)

    def apply(self, params):
        threshold = int(params["threshold"]) if params["threshold"] else 128

        self.img = cv.threshold(self.img, threshold, 255, cv.THRESH_BINARY)[1]
        self.img = np.uint8(self.img)


class binarize_otsu(img_filter):
    ''' Otsu binarization filter from opencv.
    '''
    def init(self, img):
        super().__init__(img)
    
    def apply(self, _):
        self.img = cv.threshold(self.img, 0, 255, cv.THRESH_BINARY + cv.THRESH_OTSU)[1]
        self.img = np.uint8(self.img)


class add_margin(img_filter):
    ''' Add margin to the image.
    '''
    def init(self, img):
        super().__init__(img)

    def apply(self, params):
        margin = int(params["margin"]) if params["margin"] else 10
        color = int(params["color"]) if params["color"] else 255

        self.fig.set_size_inches(
            ((self.width + 2 * margin) / self.dpi, (self.height + 2 * margin) / self.dpi))
        self.img = cv.copyMakeBorder(
            self.img, margin, margin, margin, margin, cv.BORDER_CONSTANT, value=color)
        self.height, self.width = self.img.shape
        self.img = np.uint8(self.img)


# ---------------------- MORPHOLOGICAL OPS --------------------------#

class erode(img_filter):
    ''' Erosion morphological operation from OpenCV module.
    '''

    def __init__(self, img):
        super().__init__(img)

    def apply(self, params):

        # get an ellipse kernel
        kernel = np.matrix(params["kernel"]) if params["kernel"] else cv.getStructuringElement(
            cv.MORPH_ELLIPSE, (3, 3))

        self.img = cv.morphologyEx(
            np.uint8(self.img), op=cv.MORPH_ERODE, kernel=kernel)
        self.img = np.uint8(self.img)


class dilate(img_filter):
    ''' Dilation morphological operation from OpenCV module.
    '''

    def __init__(self, img):
        super().__init__(img)

    def apply(self, params):

        # get an ellipse kernel
        kernel = np.matrix(params["kernel"]) if params["kernel"] else cv.getStructuringElement(
            cv.MORPH_ELLIPSE, (3, 3))

        self.img = cv.morphologyEx(
            np.uint8(self.img), op=cv.MORPH_DILATE, kernel=kernel)
        self.img = np.uint8(self.img)
