# -*- coding: utf-8 -*-
import numpy as np
import matplotlib.pyplot as plt

#%%##########################################################################
def _consts(sigma):
    cs = (1 + np.exp(-sigma ** 2) - 2 * np.exp(-.75 * sigma ** 2)) ** (-.5)
    ks = np.exp(-.5 * sigma ** 2)
    return cs, ks

def morlet(N, sigma=5):
    cs, ks = _consts(sigma)
    n = np.linspace(-8, 8, N)
    return np.exp(-.5 * n ** 2) * (np.exp(1j * sigma * n) - ks
                                   ) * cs * np.pi ** (-.25)

def morlet_fourier(N, sigma):
    cs, ks = _consts(sigma)
    k = np.linspace(0, N - 1, N)
    return (np.exp(-.5 * (sigma - k) ** 2) - ks * np.exp(-.5 * k ** 2)
            ) * cs * np.pi ** (-.25)

def plot(x, complex=False):
    plt.plot(x)
    if complex:
        plt.plot(x.imag)
        plt.legend(['real', 'imag'], fontsize=14)
    plt.show()
#%%##########################################################################
N, sigma = 256, 5
morl     = morlet(N, sigma)
morl_f   = morlet_fourier(N, sigma)
morl_fft = np.fft.fft(morl)
#%%##########################################################################
plot(morl, complex=True)
plot(morl_f)
plot(morl_fft, complex=True)
plot(np.abs(morl_fft))