import numpy as np
import matplotlib.pyplot as plt

#%%###########################################################################
def dft_k(x, k):
    N = len(x)
    t = np.linspace(0, 1, N, False)
    exp = np.exp(1j * 2 * np.pi * k * t)
    return (x * exp).sum()

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

def plot(x, y=None, size=0, show=1, ax_equal=False, **kw):
    if y is None:
        plt.plot(x, **kw)
    else:
        plt.plot(x, y, **kw)
    _scale_plot(plt.gcf(), plt.gca(), size=size, show=show, ax_equal=ax_equal)

def scat(x, y=None, size=0, show=1, ax_equal=False, s=18, **kw):
    if y is None:
        plt.scatter(np.arange(len(x)), x, s=s, **kw)
    else:
        plt.scatter(x, y, s=s, **kw)
    _scale_plot(plt.gcf(), plt.gca(), size=size, show=show, ax_equal=ax_equal)

def _scale_plot(fig, ax, size=True, show=False, ax_equal=False):
    xmin, xmax = ax.get_xlim()
    rng = xmax - xmin
    ax.set_xlim(xmin + .02 * rng, xmax - .02 * rng)
    if size:
        fig.set_size_inches(15, 7)
    if ax_equal:
        yabsmax = max(np.abs([*ax.get_ylim()]))
        mx = max(yabsmax, max(np.abs([xmin, xmax])))
        ax.set_xlim(-mx, mx)
        ax.set_ylim(-mx, mx)
        fig.set_size_inches(8, 8)
    if show:
        plt.show()
#%%###########################################################################
N = 128
f_signal, f_winding = 1, 0.5

t = np.linspace(0, 1, N, False)
x = np.cos(2 * np.pi * f_signal * t)
exp = np.exp(1j * 2 * np.pi * f_winding * t)
coef, theta, r = allfft(x)

dftk = (x * exp)[:128]
scat(dftk.real, dftk.imag, size=0, ax_equal=True, show=0)

re, im = dftk.sum().real, dftk.sum().imag
plt.quiver([0], [0], [re], [im], color='r', scale=8500 / max(np.abs([re, im])))
plt.show()
#%%##################################################################
x_p = np.pad(x, [0, len(x)])
coef_p, theta_p, r_p = allfft(x_p)
scat(r_p, show=1)
#%%##################################################################
coef_pk = [dft_k(x, k) for k in np.arange(0, len(x), 0.5)]
theta_pk, r_pk = np.angle(coef_pk), np.abs(coef_pk)
print(np.mean(abs(r_pk - r_p)))
#%%###########################################################################
t_p = np.linspace(0, 1, 2 * N, False)
f_winding_p = 1
exp_p = np.exp(1j * 2 * np.pi * f_winding_p * t_p)

dftk_p = (x_p * exp_p)
scat(dftk_p.real, dftk_p.imag, size=0, ax_equal=True, show=0)

rp, ip = dftk_p.sum().real, dftk_p.sum().imag
plt.quiver([0], [0], [rp], [ip], color='r', scale=8500 / max(np.abs([rp, ip])))
plt.show()
#%%##################################################################
scat(x_p, show=0)
plot(exp_p.real, color='orange')