# Name: *** Enter your name here ***
#
# Student-ID: *** Enter your student ID here ***
#
# Rename the file to <yourname>02.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: Convolution theorem
#
def exercise1():

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

    # Task 1B: Create an array storing the theoretical values for the complex
    # Fourier coefficients of `cos^2(3x)`

    # *** Insert your Python code here ***
    
    # Task 1C: Sample `cos^2(3x)` using 128 sampling points in [0, 2pi) and
    # compute its Fourier transform using NumPy's FFT

    # *** Insert your Python code here ***

    # Task 1D: Use the convolution theorem to compute the Fourier transform of
    # `cos^2(3x)`. Use 128 sampling points in [0, 2pi) to sample `cos(3x)`.
    # Evaluate the Fourier transform of `cos^2(3x)` by convolving the Fourier
    # transform of `cos(3x)` with itself.

    # *** Insert your Python code here ***

    # Task 1E: Plot the theoretical Fourier coefficients from 1B and the values
    # computed in 1C and 1D
    
    # *** Insert your Python code here ***

    
# Exercise 2: Sampling `cos(x) + 2 sin(4x)`
#
def exercise2():

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

    # Task 2A: Sample `cos(x) + 2 sin(4x)` using various numbers of sampling
    # points covering the interval [0, 2pi)
    
    # *** Insert your Python code here ***

    # Task 2B: Resample the downsampled signals and plot them. Use the Fourier
    # transform to do the interpolation. Implement the resampling procedure as
    # a local Python function `resample(f, N)` with two arguments: the
    # downsampled signal `f` (a NumPy array) and the desired size of the
    # resampled signal `N` (an integer larger than `len(f)`). Plot the original
    # signal as a transparent black line and the resampled signal as a red line
    
    # *** Insert your Python code here ***


# Exercise 3: NumPy oneliners
#
def exercise3():

    # loading oneliner inputs and results

    oneliners = np.load('./oneliners.npz', allow_pickle=True)

    # Please read the following explanation:
    #
    # In most tasks, you should generate an output that is stored in a
    # variable `result`. For example, consider the following task:
    #
    # Let x, y be two arrays, build a new array whose entries are the sum of
    # the entries in x and y
    #
    x, y = oneliners['input_test1']

    # This task can be solved by the following oneliner
    result = x + y

    # Check
    print('Test 1:', end=' ')
    print(np.allclose(result, oneliners['output_test1']))
    
    # However in the last two tasks, you need to change the *content* of an
    # array. Consider the following example task:
    #
    # Let x by an array, subtract 1 from all entries.
    #
    x = oneliners['input_test2']

    # This task can be solved by the following oneliner:
    x -= 1

    # Check
    print('Test 2:', end=' ')
    print(np.allclose(x, oneliners['output_test2']), end='\n\n')
    
    # Please do not change the code above this line

    ##########################################################################
    
    # Task 3A: Compute the inner product of `x` and `y`
    #
    x, y = oneliners['inputA']

    # *** Please replace the following line. Store the result in a variable
    # named `result` ***
    result = np.inf 
    
    print('Task 3A:', end=' ')
    print(np.allclose(result, oneliners['outputA']))

    # Task 3B: Form a matrix with elements x[i] + y[j]
    #
    x, y = oneliners['inputB']

    # *** Please replace the following line ***
    result = np.inf 
    
    print('Task 3B:', end=' ')
    print(np.allclose(result, oneliners['outputB']))

    # Task 3C: Compute the inner product of matrices `x` and `y`
    #
    x, y = oneliners['inputC']

    # *** Please replace the following line ***
    result = np.inf 
    
    print('Task 3C:', end=' ')
    print(np.allclose(result, oneliners['outputC']))

    # Task 3D: Compute the FFT of vector `x`
    #
    x = oneliners['inputD']

    # *** Please replace the following line ***
    result = np.inf 
    
    print('Task 3D:', end=' ')
    print(np.allclose(result, oneliners['outputD']))

    # Task 3E: Compute the FFT of image `x`
    #
    x = oneliners['inputE']
    
    # *** Please replace the following line ***
    result = np.inf 
    
    print('Task 3E:', end=' ')
    print(np.allclose(result, oneliners['outputE']))

    # Task 3F: Compute the full convolution of two vectors
    #
    x, y = oneliners['inputF']

    # *** Please replace the following line ***
    result = np.inf 
    
    print('Task 3F:', end=' ')
    print(np.allclose(result, oneliners['outputF']))

    # Task 3G: Create the DFT matrix
    #
    x, y = oneliners['inputG']

    # *** Please replace the following line ***
    result = np.inf 
    
    print('Task 3G:', end=' ')
    print(np.allclose(result, oneliners['outputG']))

    # Task 3H: Count the number of entries that are larger than 1
    #
    x = oneliners['inputH']

    # *** Please replace the following line ***
    result = np.inf 
    
    print('Task 3H:', end=' ')
    print(np.allclose(result, oneliners['outputH']))
    
    # Task 3I: Set all entries < 0 to 0
    #
    x = oneliners['inputI']

    # *** Please insert your code, change the content of x directly ***
    
    print('Task 3I:', end=' ')
    print(np.allclose(x, oneliners['outputI']))
    
    # Task 3J: Set all entries that are > -1 and < 1., to 0.
    #
    x = oneliners['inputJ']

    # *** Please insert your code, change the content of x directly ***
    
    print('Task 3J:', end=' ')
    print(np.allclose(x, oneliners['outputJ']))
    
    
# 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)