"""! @file filters.py
    @brief Filter library for the application
    @author xlanro00
"""

import numpy as np
#import matplotlib.pyplot as plt
import cv2 as cv
from skimage import filters as skiflt
from skimage import restoration as skirest


# Parent class for all the filters
class filter:
    ''' Parent class for all the filters.
    '''

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


class convolve(filter):
    ''' Convolve with custom kernel using opencv.
    If no kernel is given, use default 3x3 kernel for averaging.
    Possibly useful for custom filters.
    '''

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

    def apply(self, params):
        ksize = int(params["ksize"]) if params["ksize"] else 3
        kernel = np.array(params["kernel"]) if params["kernel"] else np.ones(
            (ksize, ksize), np.float32) / np.sqrt(ksize)

        print("with params: ksize: " +
              str(ksize) + " kernel: \n" + str(kernel))
        self.img = cv.filter2D(self.img, -1, kernel)


class blur(filter):
    ''' Blur filter from OpenCV.
    Performs averaging of the image.
    '''
    def __init__(self, img):
        super().__init__(img)

    def apply(self, params):
        # TODO remove try-except
        if(params["anchor"]):
            try:
                anchor = tuple(map(int, params["anchor"].split(',')))
            except AttributeError:
                anchor = tuple(params["anchor"])
        else:
            anchor = (-1, -1)
        ksize = int(params["ksize"]) if params["ksize"] else 3

        print("with params: ksize: " +
              str(ksize) + " anchor: " + str(anchor))
        self.img = cv.blur(self.img, ksize=(ksize, ksize), anchor=anchor)


class gaussian(filter):
    ''' Gaussian blur filter from OpenCV.
    '''
    def __init__(self, img):
        super().__init__(img)

    def apply(self, params):
        ksize = int(params["ksize"]) if params["ksize"] else 3
        sigmaX = float(params["sigmaX"]) if params["sigmaX"] else 0
        sigmaY = float(params["sigmaY"]) if params["sigmaY"] else 0

        print("with params: ksize: " + str(ksize) +
              " sigmaX: " + str(sigmaX) + " sigmaY: " + str(sigmaY))
        self.img = cv.GaussianBlur(self.img, (ksize, ksize), sigmaX, sigmaY)


class median(filter):
    ''' Median blur filter from OpenCV.
    '''
    def __init__(self, img):
        super().__init__(img)

    def apply(self, params):
        ksize = int(params["ksize"]) if params["ksize"] else 3

        print("with params: ksize: " + str(ksize))
        self.img = cv.medianBlur(np.float32(self.img), ksize)


class bilateral(filter):
    ''' Bilateral filter from OpenCV.
    '''
    def __init__(self, img):
        super().__init__(img)

    def apply(self, params):
        # Set default values
        d = int(params["d"]) if params["d"] else 1
        sigmaColor = int(params["sigmaColor"]) if params["sigmaColor"] else 75
        sigmaSpace = int(params["sigmaSpace"]) if params["sigmaSpace"] else 75

        print("with params: d: " + str(d) + " sigmaColor: " +
              str(sigmaColor) + " sigmaSpace: " + str(sigmaSpace))
        self.img = cv.bilateralFilter(self.img, d, sigmaColor, sigmaSpace)


class denoise(filter):
    # TODO possibly not necessary
    def __init__(self, img):
        super().__init__(img)

    def apply(self, params):
        h = int(params["h"]) if params["h"] else 10
        tWS = int(params["templateWindowSize"]
                  ) if params["templateWindowSize"] else 7
        sWS = int(params["searchWindowSize"]
                  ) if params["searchWindowSize"] else 21

        print("with params: h: " + str(h) +
              " tWS: " + str(tWS) + " sWS: " + str(sWS))
        self.img = np.uint8(self.img)
        self.img = cv.fastNlMeansDenoising(
            self.img, h, tWS, sWS)


class denoise_bilateral(filter):
    ''' Scikit image denoise_bilateral filter.

    Performs bilateral denoising technique on the image.
    Averages pixels based on their distance and color similarity.
    Preserves edges while removing unwanted noise.
    '''

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

    def apply(self, params):
        sigmaColor = float(params["sigmaColor"]
                           ) if params["sigmaColor"] else 0.1
        sigmaSpace = float(params["sigmaSpace"]
                           ) if params["sigmaSpace"] else 15.0
        channelAxis = int(params["channelAxis"]
                          ) if params["channelAxis"] else None
        iterations = int(params["iterations"]) if params["iterations"] else 1

        print("with params: sigma_color: " + str(sigmaColor) +
              " sigma_spatial: " + str(sigmaSpace) + " channel_axis: " +
              str(channelAxis) + " iterations: " + str(iterations))

        for i in range(iterations):
            self.img = skirest.denoise_bilateral(
                self.img, sigma_color=sigmaColor,
                sigma_spatial=sigmaSpace, channel_axis=channelAxis)


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

    Performs total variation denoising technique on the image.
    This filter removes fine detail, but preserves edges.
    '''

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

    def apply(self, params):
        weight = float(params["weight"]) if params["weight"] else 0.1
        channelAxis = int(params["channelAxis"]
                          ) if params["channelAxis"] else None
        iterations = int(params["iterations"]) if params["iterations"] else 1

        print("with params: weight: " + str(weight) +
              " channel_axis: " + str(channelAxis) + " iterations: " + str(iterations))
        for i in range(iterations):
            self.img = skirest.denoise_tv_chambolle(
                self.img, weight=weight, channel_axis=channelAxis)


class sharpen(filter):
    ''' Convolution with a sharpening kernel using opencv.
    '''
    # TODO possibly unnecessary, because unsharp masking is working better
    def __init__(self, img):
        super().__init__(img)

    def apply(self, params):
        kernel = np.matrix(params["kernel"]) if params["kernel"] else np.array(
            [[0, -1, 0], [-1, 5, -1], [0, -1, 0]])

        print("with params: kernel: \n" + str(kernel))
        self.img = cv.filter2D(self.img, ddepth=-1, kernel=kernel)


class unsharp_mask(filter):
    ''' Unsharp mask filter from opencv.

    First blur the image a little bit, then calculate Laplacian of the image to get the edges.
    Scale the Laplacian and subtract it from the original image.
    '''

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

    def apply(self, params):
        strength = float(params["strength"]) if params["strength"] else 1.0
        ksize = int(params["ksize"]) if params["ksize"] else 3

        blurred = cv.medianBlur(self.img, ksize)
        lap = cv.Laplacian(blurred, cv.CV_32F)

        print("with params: strength: " +
              str(strength) + " ksize: " + str(ksize))
        self.img = blurred - strength*lap


class unsharp_mask_scikit(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 opencv module. 
    '''

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

    def apply(self, params):
        radius = int(params["radius"]) if params["radius"] else 3
        amount = float(params["amount"]) if params["amount"] else 1
        channelAxis = int(params["channelAxis"]
                          ) if params["channelAxis"] else None

        #self.img = cv.cvtColor(self.img, cv.COLOR_GRAY2RGB)
        print("with params: radius: " +
              str(radius) + " amount: " + str(amount))
        self.img = skiflt.unsharp_mask(
            self.img, radius=radius, amount=amount, channel_axis=channelAxis)
        #self.img = cv.cvtColor(self.img, cv.COLOR_RGB2GRAY)


class morph(filter):
    ''' General morphological operations from OpenCV.

    Can be used with MORPH_OPEN, MORPH_CLOSE, MORPH_DILATE, MORPH_ERODE and more as 'op'.
    '''

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

    def apply(self, params):
        kernel = np.matrix(params["kernel"]) if params["kernel"] else np.ones(
            (3, 3), np.uint8)
        iterations = int(params["iterations"]) if params["iterations"] else 1
        op = getattr(cv, params["op"]) if params["op"] else cv.MORPH_OPEN
        if(params["anchor"]):
            # TODO remove try-except
            try:
                anchor = tuple(map(int, params["anchor"].split(',')))
            except AttributeError:
                anchor = tuple(params["anchor"])
        else:
            anchor = (-1, -1)

        print("with params: kernel: \n" + str(kernel) + " anchor: " +
              str(anchor) + " iterations: " + str(iterations) + " op: " + str(op))
        self.img = cv.morphologyEx(
            np.uint8(self.img), op=op, kernel=kernel,
            anchor=anchor, iterations=iterations)
