import cv2
import numpy as np
import matplotlib.pyplot as plt


class frequency_filter:
    def __init__(self, img_path):
        self.path = img_path

    def Ideal_low_pass_filtering_smooth(self):
        img = cv2.imread(self.path)
        img = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
        img = np.float32(img)
        dft = cv2.dft(img, flags=cv2.DFT_COMPLEX_OUTPUT)
        dft = np.fft.fftshift(dft)
        rows, cols = img.shape
        mask = np.zeros((rows, cols, 2), dtype=np.uint8)
        mask[int(rows / 2) - 20:int(rows / 2) + 20, int(cols / 2) - 20:int(cols / 2) + 20] = 1
        f = dft * mask
        f = np.fft.ifftshift(f)
        img = cv2.idft(f)
        img = cv2.magnitude(img[:, :, 0], img[:, :, 1])
        filename = "saved Img/ideal_low_pass.bmp"
        plt.imsave(filename, img, cmap='gray')

    def butterworth(self, D0, N):
        D0 = int(D0)
        N = int(N)
        img = cv2.imread(self.path)
        img = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
        img = np.float32(img)
        dft = cv2.dft(img, flags=cv2.DFT_COMPLEX_OUTPUT)
        dft = np.fft.fftshift(dft)
        rows, cols = img.shape
        mask = np.zeros((rows, cols, 2))
        for i in range(rows):
            for j in range(cols):
                D = np.sqrt((i - rows / 2) ** 2 + (j - cols / 2) ** 2)
                mask[i, j, 0] = mask[i, j, 1] = 1.0 / (1.0 + ((D / D0) ** (2 * N)))
        f = dft * mask
        f = np.fft.ifftshift(f)
        img = cv2.idft(f)
        img = cv2.magnitude(img[:, :, 0], img[:, :, 1])
        filename = "saved Img/butterworth.bmp"
        plt.imsave(filename, img, cmap='gray')

    def gauss(self, D0):
        D0 = int(D0)
        img = cv2.imread(self.path)
        img = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
        img = np.float32(img)
        dft = cv2.dft(img, flags=cv2.DFT_COMPLEX_OUTPUT)
        dft = np.fft.fftshift(dft)
        rows, cols = img.shape
        crow, ccol = int(rows / 2), int(cols / 2)
        mask = np.zeros((rows, cols, 2))
        for i in range(rows):
            for j in range(cols):
                D = np.sqrt((i - crow) ** 2 + (j - ccol) ** 2)
                mask[i, j, 0] = mask[i, j, 1] = np.exp(-(D * D) / (2 * D0 * D0))
        f = dft * mask
        f = np.fft.ifftshift(f)
        img = cv2.idft(f)
        img = cv2.magnitude(img[:, :, 0], img[:, :, 1])
        filename = "saved Img/gauss.bmp"
        plt.imsave(filename, img, cmap='gray')