###  functions useful for this part

import os, cv2
import numpy as np
import pandas as pd
from PIL import Image, ImageOps
from scipy.signal import convolve, fftconvolve, convolve2d

# Interactivity
import ipywidgets as widgets
from ipywidgets import interact, interact_manual

import plotly.express as px                           # Advanced visualization library (similar to Seaborn)
import plotly.graph_objects as go

from matplotlib.pylab import *
import matplotlib.pyplot as plt
from matplotlib import cm
from mpl_toolkits import mplot3d
%matplotlib inline
%config InlineBackend.figure_format = 'retina'        # Improve the sharpness of figures
plt.rcParams["figure.figsize"] = 12.8, 9.6

import warnings                                        # "Do not disturb" mode
warnings.filterwarnings('ignore')


# Adding Gaussian noise
def gaussian_noise(img, mean = 0.0, std = 10.0):
    noisy_gauss = img + np.random.normal(mean, std, img.shape)
    return np.array(np.clip(noisy_gauss, 0, 255), dtype='uint8')


def display_1x2(src, dst):
    fig = plt.figure(figsize=(12, 6))
    plt.subplot(121), plt.imshow(src, cmap="gray", origin="upper"), plt.title("Original image")
    plt.subplot(122), plt.imshow(dst, cmap="gray", origin="upper"), plt.title("Filtered image")
    plt.show()


def display_2x2(img1, img2, img3, img4):
    fig = plt.figure(figsize=(12, 12))

    plt.subplot(221), plt.imshow(img1, cmap="gray", origin="upper"), plt.title("Original")
    plt.subplot(222), plt.imshow(img2, cmap="gray", origin="upper"), plt.title("Filtered (FFT)")
    plt.subplot(223), plt.imshow(img3, cmap="hot", extent=(-img3.shape[0]//2, img3.shape[0]//2, -img3.shape[1]//2, img3.shape[1]//2)), plt.title("[Fourier] Original")
    plt.subplot(224), plt.imshow(img4, cmap="hot", extent=(-img4.shape[0]//2, img4.shape[0]//2, -img4.shape[1]//2, img4.shape[1]//2)), plt.title("[Fourier] Filtered")
    plt.show()


def kernel_heatmap(kernel, cm="gray"):
    fig = px.imshow(kernel, color_continuous_scale=cm)
    fig.show()


def kernel_3D(kernel):
    z_data = pd.DataFrame(kernel)
    fig = go.Figure(data=[go.Surface(z=z_data.values)])

    fig.update_layout(title='Kernel', autosize=True,
                      width=500, height=500,
                      margin=dict(l=65, r=50, b=65, t=90))

    fig.show()


def kernel_3D_continuous(data, color_map="viridis"):

    x,y,z = data
    X1 = np.reshape(x, -1)
    Y1 = np.reshape(y, -1)
    Z1 = np.reshape(z, -1)

    # Normalize the colors based on Z value
    norm = plt.Normalize(z.min(), z.max())
    colors = cm.jet(norm(z))
    ax = plt.axes(projection='3d')
    ax.get_proj = lambda: np.dot(mplot3d.Axes3D.get_proj(ax), np.diag([1, 1, 1, 1]))
    ax.plot_trisurf(X1, Y1, Z1, cmap=color_map)
kernel_size_slider = widgets.IntSlider(min=3, max=7, step=2, value=3)

img_path = "img/"

@interact_manual
def mean_filter(image=[f for f in os.listdir(img_path) if os.path.isfile(os.path.join(img_path, f))], kernel_size=kernel_size_slider):
    img = np.array(Image.open(img_path + image).convert("L")).astype("uint8")

    #Applying noise and then filtering
    noisy = gaussian_noise(img)

    # Convolution kernel creation
    kernel = 1/9 * np.ones((kernel_size,kernel_size), dtype="uint8")

    # Denoising
    denoised_average = convolve2d(noisy, kernel, mode="same")

    # we display
    display_1x2(noisy, denoised_average)

kernel_size_slider = widgets.IntSlider(min=3, max=7, step=2, value=3)

@interact_manual
def median_filter(image=[f for f in os.listdir(img_path) if os.path.isfile(os.path.join(img_path, f))], kernel_size=kernel_size_slider):
    img = np.array(Image.open(img_path + image).convert("L")).astype("uint8")

    noisy = gaussian_noise(img)
    denoised_median = cv2.medianBlur(noisy, ksize=kernel_size)

    # we display
    display_1x2(noisy, denoised_median)

from scipy.stats import multivariate_normal

def Gaussian(mu=[0.0, 0.0], sigma=[.5, .5]):
    n = sigma[0] + sigma[1] *3
    x, y = np.mgrid[-n:n:30j, -n:n:30j]
    xy = np.column_stack([x.flat, y.flat])
    covariance = np.diag(np.array(sigma)**2)
    z = multivariate_normal.pdf(xy, mean=np.array(mu), cov=covariance)
    z = z.reshape(x.shape)
    return x,y,z

@interact
def gaussiant_filter(sigmaX=(0.5,5,0.5), sigmaY=(0.5,5,0.5)):
    kernel_3D_continuous(Gaussian([0.0, 0.0], [sigmaX, sigmaY]))

kernel_size_slider = widgets.IntSlider(min=3, max=7, step=2, value=3)

@interact
def gaussiant_filter(image=[f for f in os.listdir(img_path) if os.path.isfile(os.path.join(img_path, f))], kernel_size=kernel_size_slider, noise=(0,30,1), sigma=(0.5,5,0.5)):
    img = np.array(Image.open(img_path + image).convert("L")).astype("uint8")

    # Application of Gaussian noise followed by Gaussian filtering.
    noisy = gaussian_noise(img, 0, noise)
    denoised_gaussian = cv2.GaussianBlur(noisy, ksize=(kernel_size,kernel_size), sigmaX=sigma, sigmaY=sigma, borderType=cv2.BORDER_DEFAULT)

    # we display
    display_1x2(noisy, denoised_gaussian)

@interact_manual
def mean_filter1_fft(image=[f for f in os.listdir(img_path) if os.path.isfile(os.path.join(img_path, f))], mask_size=(0,100,10)):
    img = np.array(Image.open(img_path + image).convert("L"))

## Spatial --> Frequential
    # We apply the discrete Fourier transform
    dft = np.fft.fft2(img)
# Shift the result to center it
    dft_shift = np.fft.fftshift(dft)
# Calculation of magnitude spectrum (Fourier representation)
    magnitude_img = 20 * np.log(np.abs(dft_shift))

# Filtering: creation of a square black binary mask, except at its center
    rows, cols = img.shape
    crow, ccol = rows//2 , cols//2
    mask = np.zeros((rows,cols), np.uint8)
    mask[crow-mask_size:crow+mask_size, ccol-mask_size:ccol+mask_size] = 1
## Apply the mask
    img_filtered = dft_shift * mask
    # Et calcule sa représentation de Fourier
    magnitude_img_filtered = 20 * np.log(np.abs(img_filtered))

### Frequency --> Spatial
    # Shift inverse
    img_filtered = np.fft.ifftshift(img_filtered)
# TF reversal
    img_filtered = np.fft.ifft2(img_filtered)
    img_filtered = np.abs(img_filtered)

    display_2x2(img, img_filtered, magnitude_img, magnitude_img_filtered)

@interact_manual
def mean_filter1_fft(image=[f for f in os.listdir(img_path) if os.path.isfile(os.path.join(img_path, f))], mask_size=(0,100,10)):
    img = np.array(Image.open(img_path + image).convert("L"))

    # Spatial --> Frequential
    dft = cv2.dft(np.float32(img), flags=cv2.DFT_COMPLEX_OUTPUT)
    dft_shift = np.fft.fftshift(dft)

    # Calculation of magnitude spectrum (Fourier representation)
    magnitude_img = 20 * np.log(cv2.magnitude(dft_shift[:,:,0], dft_shift[:,:,1]))

    # Filtering: creation of a square black binary mask, except at its center
    rows, cols = img.shape
    crow, ccol = rows//2 , cols//2
    mask = np.zeros((rows,cols,2), np.uint8)
    mask[crow-mask_size:crow+mask_size, ccol-mask_size:ccol+mask_size] = 1

    # Apply the mask
    img_filtered = dft_shift * mask
    # And calculates its Fourier representation
    magnitude_img_filtered = 20 * np.log(cv2.magnitude(img_filtered[:,:,0], img_filtered[:,:,1]))

    # Frequential --> Spatial
    img_filtered = np.fft.ifftshift(img_filtered)
    img_filtered = cv2.idft(img_filtered)
    img_filtered = cv2.magnitude(img_filtered[:,:,0], img_filtered[:,:,1])

    display_2x2(img, img_filtered, magnitude_img, magnitude_img_filtered)

kernel_size_slider = widgets.IntSlider(min=3, max=7, step=2, value=3)

@interact_manual
def mean_filter3_fft(image=[f for f in os.listdir(img_path) if os.path.isfile(os.path.join(img_path, f))], kernel_size=kernel_size_slider):
    img = np.array(Image.open(img_path + image).convert("L")).astype("uint8")

    kernel = 1/kernel_size * np.ones((kernel_size,kernel_size), dtype="uint8")

    noisy = gaussian_noise(img, 0, 15)

    fft_filtered = fftconvolve(noisy, kernel, mode='same')

    display_1x2(noisy, fft_filtered)

### <span style="color:crimson">**[<u>Exercise</u>]** Your turn:</span>
***
<div style="color:DarkSlateBlue">  
    
1. **Create an interactive method to generate a square image (in grayscale) whose values will be assigned according to the equation $f(x,y) = \sin(2\pi (\lambda_x*x + \lambda_y*y))$ with sliders to modify the values of $\lambda_x$ and $\lambda_y$ (between 0 and 1, with a step of 0.1).**.
    
> <u>Trick</u>:
```python
n = 300
xs = np.linspace(-2*np.pi, 2*np.pi, n)
ys = np.linspace(-2*np.pi, 2*np.pi, n)
x, y = np.meshgrid(xs, ys)
z = #2D sin wave calculation
```

    
2. **Create a method for displaying the resulting image in both the spatial and frequency domains (absolute shift logarithmic representation).
    
> <u>Trick</u>: you can draw inspiration from the `kernel_heatmap` method for display in the spatial domain.
    
> Vary $\lambda_x$ and $\lambda_y$ and observe the effect on the TF magnitude spectrum.

> Add an option to binarize the generated image.
    

3. **Create an interactive method allowing:**
    * Select an image
    * apply Gaussian or Salt&Pepper noise
    * Denoise/smooth by choosing a filter (average, median, Gaussian) and filter size.
    * Display the result as a 2x3 grid:

|            |                |                 |
|:----------:|:--------------:|:---------------:|
| image      | noisy_image    | denoised_image  |
| fft_image  | fft_noisy_image| fft_denoised_image|

> Observe the differences in denoising quality of different filters depending on the type of noise.
    
    
4. **Add the possibility of generating circular binary masks (whose radius can be adjusted interactively), with a white center and black exterior**.

> Apply the mask to the frequency representation of your image and observe the result on the recomposed image.
    
    
5. **Modify your code so that the mean and Gaussian filters are applied "manually": create the kernel and apply it via the `fftconvolve` method in `Scipy`.**.
    
    
6. **Modify your code so that you can apply the filter to a color image, by implementing 3 methods:**.
    1. Apply it to each of the image channels and combine the result
    2. Apply it to the luminance channel of an appropriate color space
    
> Display the 3 results side-by-side and compare.
    
    
7. **Propose a solution for extracting edges/high frequencies from an image using only a Gaussian filter.
    

8. **Extend the principle of question 7. by applying it to the difference of two Gaussian kernels, i.e. the same image smoothed according to two Gaussian kernels (with the first image less blurred than the second).
    
> You've just applied a Gaussian Difference (**DoG**), which is an edge extraction technique.
    
8.  Propose a basic method for quantifying the difference in quality between the image and its denoised version, and use this method to compare the effectiveness of different filters according to the type of noise applied.**.
    
</div>