import numpy as np
import matplotlib.pyplot as plt
from scipy.io import wavfile
from scipy.linalg import dft
from matplotlib import cm


L = 1024 # Window size
N = 2048 # Lenght of X
H = 512 # Hop size

dft = dft(L)
haan = np.hanning(L)
dft = np.multiply(dft, haan)

k = (int(N/H) - 1)*L
A = np.zeros((k,N))
print(A.shape)
# print(A[0:L][0:L].shape)
for i in range((int(N/H) - 1)):
    A[i*L:(i+1)*L, 0:L] = dft
    A[i*L:(i+1)*L, :] = np.roll(A[i*L:(i+1)*L, :], i*H, axis=1)

plt.imshow(np.absolute(A), interpolation='None', cmap=cm.coolwarm)
plt.colorbar()
plt.show()