import numpy as np
import matplotlib.pyplot as plt

#%%###########################################################################
def allfft(x):
    coef = np.fft.fft(x)
    return coef, np.angle(coef), np.abs(coef)

def plot(x, y=None, show=False, **kw):
    if y is None:
        plt.plot(x, **kw)
    else:
        plt.plot(x, y, **kw)
    _scale(plt.gcf(), plt.gca(), show=show)

def scat(x, y=None, show=False, **kw):
    if y is None:
        plt.scatter(np.arange(len(x)), x, **kw)
    else:
        plt.scatter(x, y, **kw)
    _scale(plt.gcf(), plt.gca(), show=show)

def _scale(fig, ax, show=False):
    xmin, xmax = ax.get_xlim()
    rng = xmax - xmin
    ax.set_xlim(xmin + .02 * rng, xmax - .02 * rng)
    if show:
        plt.show()

#%%###########################################################################
M = 1 + 128
fc = 2 / M
t  = np.linspace(0, 1, M)
ti = np.arange(M)
wh = .54 - .46 * np.cos(2 * np.pi * t)  # Hamming
sinc = np.sin(2 * np.pi * fc * (ti - M / 2)) / (ti - M / 2)
if M % 2 == 0:
    sinc[M // 2] = 2 * np.pi * fc

whs = wh * sinc / np.pi
#%%###########################################################################
coef_hs, theta_hs, r_hs = allfft(whs)
#%%
scat(np.log10(r_hs)     * 20, show=1, color='g')
scat(np.log10(r_hs[:9]) * 20, show=1, color='g')
#%%###########################################################################
tx = np.linspace(0, 1, 128, endpoint=False)
s1 = 10 * np.cos(2 * np.pi * 1 * tx)
s2 = 1  * np.cos(2 * np.pi * 5 * tx)
s = s1 + s2
coef_s, theta_s, r_s = allfft(s)

conv_len = len(s) + M - 1
sp   = np.pad(s,   [0, conv_len - len(s)])
whsp = np.pad(whs, [0, conv_len - M])

coef_sp,  theta_sp,  r_sp  = allfft(sp)
coef_hsp, theta_hsp, r_hsp = allfft(whsp)
#%%###########################################################################
plot(s1)
plot(s2, show=1)
plot(s,  show=1, color='k')
#%%
out = np.convolve(sp, whsp)[:conv_len]
coef_o, theta_o, r_o = allfft(out)
#%%###########################################################################
plot(s1)
plot(out, color='purple')
plt.gca().set_xlim(-2, 257)
plt.show()

scat(r_o, color='purple', show=1)