# Name: *** Enter your name here ***
#
# Student-ID: *** Enter your student ID here ***
#
# Rename the file to <yourname>01.py
#
# Submit your solutions via moodle
#
import numpy as np
import scipy.signal as sig
import matplotlib.pylab as plt

# Please do not change the code above this line

# Exercise 1: Fourier representation of sinusoidal signals
#
def exercise1():

    # 1D arrays
    f = np.array([0, 1, 2, 3, 2, 1, 0])
    g = np.array([1, 2, 3])
    h = np.array([1, -2, 1])

    # 2D arrays
    F = np.array([[0, 0, 0, 0, 0, 0, 0], 
                  [0, 0, 0, 1, 0, 0, 0], 
                  [0, 0, 1, 2, 1, 0, 0], 
                  [0, 1, 2, 2, 2, 1, 0], 
                  [0, 0, 1, 2, 1, 0, 0], 
                  [0, 0, 0, 1, 0, 0, 0], 
                  [0, 0, 0, 0, 0, 0, 0]])

    G = np.array([[1, 2, 3],
                  [2, 3, 0],
                  [3, 0, 0]])

    H = np.array([[0, 1, 0],
                  [1, -4, 1],
                  [0, 1, 0]])

    # Please do not change the code above this line

    # Task 1B: compute the (full) convolution with `np.convolve` (for 1D arrays)
    # and `scipy.signal.convolve2d` (for 2D arrays)

    print('1D convolutions:\n' + '=' * 50)
    print('f o g =', np.convolve(f, g))
    print('g o f =', np.convolve(g, f))
    print('f o h =', np.convolve(f, h))
    print('f o f =', np.convolve(f, f))

    print('\n2D convolutions:\n' + '=' * 50)
    print('F o G =')
    print(sig.convolve2d(F, G))
    print('F o H =')
    print(sig.convolve2d(F, H))

    # Task 1C: Compute `f o g` (where `o` denotes the convolution) by applying
    # the convolution theorem

    # zero-pad the filter `g` so as to match the size of the signal `f`
    #
    g_pad = np.pad(g, (len(f)-len(g)) // 2)

    # move the center of the padded filter to the start of the array
    #
    g_shift = np.fft.ifftshift(g_pad)

    # map the signal and modified filter to the frequency domain
    #
    f_ft, g_ft = np.fft.fft(f), np.fft.fft(g_shift)

    # multiply the (transformed) arrays elementwise
    #
    conv_ft = f_ft * g_ft

    # map the result back to the spatial domain
    #
    conv = np.fft.ifft(conv_ft)

    # compare the result with the result obtained in 1B
    #
    print('1D convolutions agree? {}'.format(
        np.allclose(conv, np.convolve(f, g, 'same')))
    )
    
    # Task 1D: Compute `F o G` by applying the convolution theorem. Proceed as
    # in 1C by using the 2D versions of fft and ifft
    #
    G_pad = np.pad(G, (len(F)-len(G)) // 2)
    G_shift = np.fft.ifftshift(G_pad)
    F_ft, G_ft = np.fft.fft2(F), np.fft.fft2(G_shift)
    conv_ft = F_ft * G_ft
    conv = np.fft.ifft2(conv_ft)
    print('2D convolutions agree? {}'.format(
        np.allclose(conv, sig.convolve2d(F, G, 'same')))
    )
    
            
# Exercise 2: Denoising with a sliding average
#
def exercise2():

    # close all figures
    plt.close('all')

    # generation of a noisy signal
    
    # fix PRNG seed
    np.random.seed(42)
    
    # unperturbed signal
    N = 1000
    x = np.linspace(0, 5*np.pi, N, endpoint=False)
    signal = np.sin(x) + np.sin(2*x) + np.cos(5*x)

    # generate noise with desired SNR
    snr = 5.
    sigma = np.sqrt(np.var(signal) / snr)
    noise = sigma * np.random.randn(N)

    # perturbed signal
    y = signal + noise

    # please do not change the code above this line

    # Task 2A: Compute the moving average of `y` where a data point is replaced
    # by the average of the previous 10 data points, the data point itself and
    # the following 10 data points.
    #
    n = 10
    w = 2 * n + 1

    # solution 1
    y_A = np.zeros(len(y) + w - 1)
    counts = np.zeros(len(y) + w -1)
    for i in range(w):
        y_A[i:i+len(y)] += y
        counts[i:i+len(y)] += 1
    y_A /= counts
    y_A = y_A[n:-n]

    # solution 2
    y_A = np.array([np.mean(y[max(i-n, 0):i+n+1]) for i in range(len(y))])
    
    # Task 2B: Compute the moving average of `y` by using the fact that a
    # sliding average can be expressed as the convolution of the signal with a
    # rectangular pulse. 

    # sum up signal inside window
    y_B = np.convolve(y, np.ones(w), 'same')
    
    # compute number of contributing data points

    # 1) lazy solution
    norm = np.convolve(np.ones_like(y), np.ones(w), 'same')

    # 2) alternative solution
    norm = np.ones_like(y) * w
    norm[:n] = np.arange(n+1, w)
    norm[-(n+1):] = np.arange(w, n, -1)

    # divide by number of data points contributing to average
    y_B /= norm

    # Task C: Plot the results obtained in 2A and 2B and compare them.

    fig, axes = plt.subplots(3, 1, figsize=(10, 6), sharex='all', sharey='all')
    for ax in axes[:2]:
        ax.plot(x, y, color='k', lw=3, alpha=0.3)

    axes[0].set_title('2A: sliding mean')
    axes[0].plot(x, y_A, color='r')

    axes[1].set_title('2B: convolution')
    axes[1].plot(x, y_B, color='b')

    axes[2].set_title('2A vs 2B')
    axes[2].plot(x, y_A, color='r', lw=6, label='2A', alpha=0.5)
    axes[2].plot(x, y_B, color='b', label='2B')
    axes[2].legend()
    axes[2].set_xlabel(r'location $x$')

    fig.tight_layout()


# Exercise 3: Filtering an image
#
def exercise3():

    # close all figures
    plt.close('all')

    # Task 3A: Load 'noisycells.tif' with plt.imread and convert its datatype
    # to double. Filter the image by averaging all pixels inside a patch of
    # size 5x5. 

    image = plt.imread('noisycells.tif').astype(float)

    n = 2
    w = 2 * n + 1
    
    image_A = np.zeros_like(image)
    for i in range(image.shape[0]):
        for j in range(image.shape[1]):
            image_A[i, j] = np.mean(image[max(0, i-n):i+n+1,
                                          max(0, j-n):j+n+1])

    # Task 3B. Filter the image by using a convolution (sig.convolve2d); see
    # also Task 2B for a 1D version of the same task.
            
    image_B = sig.convolve2d(image, np.ones((w, w)), 'same')
    image_B /= sig.convolve2d(np.ones_like(image), np.ones((w, w)), 'same')

    # Task 3C. Generate a binary image by thresholding the image from 3A (or
    # 3B) at a minimum intensity of 50. 

    # assuming that intensities lie in [0, 255] which is the case for tif image
    # but not for png or jpg file
    image_C = image_A > 50

    # better (i.e. more general) solution
    image_C = image_A > image_A.max() / 5.
        
    # Task 3D. Apply filter `H` from Exercise 1 to the image obtained in 3C. 
    # Binarize the resulting image by setting all intensities different from
    # zero to `True` and all pixels with intensity zero to `False`
            
    H = np.array([[0, 1, 0],
                  [1, -4, 1],
                  [0, 1, 0]])
    
    image_D = sig.convolve(image_C, H, 'same')
    image_D = image_D == 0
    
    # Task 3E. Plot all images

    images = [image, image_A, image_B, image_C, image_D]
    titles = ['Original', 'Task 3A', 'Task 3B', 'Task 3C', 'Task 3D']

    plt.rc('font', size=12)
    fig, axes = plt.subplots(2, 3, figsize=(10, 6))
    axes = list(axes.flat)
    for ax in axes:
        ax.axis('off')
    for ax, im, title in zip(axes, images, titles):
        ax.set_title(title)
        ax.imshow(im, cmap=plt.cm.gray)
    fig.tight_layout()

    
# Test code - you shouldn't change this part of the template

if __name__ == '__main__':

    exercises = {'1': (exercise1, ),
                 '2': (exercise2, ),
                 '3': (exercise3, ), 
                 }

    choice = input('Please enter the number of the exercise: ').upper()
    if choice not in exercises:
        print('Please choose among {0}'.format(
            list(exercises.keys())))
    else:
        func, *args = exercises[choice]
        print('Task {0}, running "{1}"'.format(
            choice, func.__name__))
        result = func(*args)
        if result is not None:
            print('Result:', result)