#%qtconsole

from __future__ import  division

Nf = 64 # N- DFT size
fs = 64 # sampling frequency
f = 10  # one signal
t = arange(0,1,1/fs) # time-domain samples
deltaf = 1/2. # second nearby frequency

fig,ax = subplots(2,1,sharex=True,sharey=True)
fig.set_size_inches((8,3))

x=cos(2*pi*f*t) + cos(2*pi*(f+2)*t) # 2 Hz frequency difference
X = fft.fft(x,Nf)/sqrt(Nf)
ax[0].plot(linspace(0,fs,Nf),abs(X),'-o')
ax[0].set_title(r'$\delta f = 2$',fontsize=18)
ax[0].set_ylabel(r'$|X(k)|$',fontsize=18)
ax[0].grid()

x=cos(2*pi*f*t) + cos(2*pi*(f+deltaf)*t) # delta_f frequency difference
X = fft.fft(x,Nf)/sqrt(Nf)
ax[1].plot(linspace(0,fs,Nf),abs(X),'-o')
ax[1].set_title(r'$\delta f = 1/2$',fontsize=14)
ax[1].set_ylabel(r'$|X(k)|$',fontsize=18)
ax[1].set_xlabel('Frequency (Hz)',fontsize=18)
ax[1].set_xlim(xmax = fs/2)
ax[1].set_ylim(ymax=6)
ax[1].grid()

# fig.savefig('figure_00@.png', bbox_inches='tight', dpi=300)

Nf = 64*2
fig,ax = subplots(2,1,sharex=True,sharey=True)
fig.set_size_inches((8,4))

X = fft.fft(x,Nf)/sqrt(Nf)
ax[0].plot(linspace(0,fs,len(X)),abs(X),'-o',ms=3.)
ax[0].set_title(r'$N=%d$'%Nf,fontsize=18)
ax[0].set_ylabel(r'$|X(k)|$',fontsize=18)
ax[0].grid()

Nf = 64*4
X = fft.fft(x,Nf)/sqrt(Nf)
ax[1].plot(linspace(0,fs,len(X)),abs(X),'-o',ms=3.)
ax[1].set_title(r'$N=%d$'%Nf,fontsize=18)
ax[1].set_ylabel(r'$|X(k)|$',fontsize=18)
ax[1].set_xlabel('Frequency (Hz)',fontsize=18)
ax[1].set_xlim(xmax = fs/2)
ax[1].set_ylim(ymax=6)
ax[1].grid()

# fig.savefig('figure_00@.png', bbox_inches='tight', dpi=300)

t = arange(0,2,1/fs)
x=cos(2*pi*f*t) + cos(2*pi*(f+deltaf)*t)

Nf = 64*2
fig,ax = subplots(2,1,sharex=True,sharey=True)
fig.set_size_inches((8,4))

X = fft.fft(x,Nf)/sqrt(Nf)
ax[0].plot(linspace(0,fs,len(X)),abs(X),'-o',ms=3.)
ax[0].set_title(r'$N=%d$'%Nf,fontsize=18)
ax[0].set_ylabel(r'$|X(k)|$',fontsize=18)
ax[0].grid()

Nf = 64*8
X = fft.fft(x,Nf)/sqrt(Nf)
ax[1].plot(linspace(0,fs,len(X)),abs(X),'-o',ms=3.)
ax[1].set_title(r'$N=%d$'%Nf,fontsize=18)
ax[1].set_ylabel(r'$|X(k)|$',fontsize=18)
ax[1].set_xlabel('Frequency (Hz)',fontsize=18)
ax[1].set_xlim(xmax = fs/2)
ax[1].set_ylim(ymax=6)
ax[1].grid()

# fig.savefig('figure_00@.png', bbox_inches='tight', dpi=300)

def abs_sinc(k=None,N=64,Ns=32):
    if k is None: k = arange(0,N-1)
    y = where(k == 0, 1.0e-20, k)
    return abs(sin( Ns*2*pi/N*y)/sin(2*pi*y/N))/sqrt(N)

fig,ax=subplots()
fig.set_size_inches((8,3))

ax.plot(abs_sinc(N=512,Ns=10),label='duration=10')
ax.plot(abs_sinc(N=512,Ns=20),label='duration=20')
ax.set_xlabel('DFT Index',fontsize=18)
ax.set_ylabel(r'$|X(\Omega_k)|$',fontsize=18)
ax.set_title('Rectangular Windows DFTs',fontsize=18)
ax.grid()
ax.legend(loc=0);

# fig.savefig('figure_00@.png', bbox_inches='tight', dpi=300)

def dftmatrix(Nfft=32,N=None):
    'construct DFT matrix'
    k= np.arange(Nfft)
    if N is None: N = Nfft
    n = arange(N)
    U = matrix(exp(1j* 2*pi/Nfft *k*n[:,None])) # use numpy broadcasting to create matrix
    return U/sqrt(Nfft)

Nf = 32 # DFT size
U = dftmatrix(Nf,Nf)
x = U[:,12].real # input signal
X = U.H*x # DFT of input
rect = ones((Nf/2,1)) # short rectangular window
z = x[:Nf/2] # product of rectangular window and x (i.e. chopped version of x) 
R = dftmatrix(Nf,Nf/2).H*rect # DFT of rectangular window
Z = dftmatrix(Nf,Nf/2).H*z    # DFT of product of x_n and r_n

idx=arange(Nf)-arange(Nf)[:,None] # use numpy broadcasting to setup summand's indices
idx[idx<0]+=Nf # add periodic Nf to negative indices for wraparound
a = arange(Nf) # k^th frequency index

fig,ax = subplots(4,8,sharex=True,sharey=True)
fig.set_size_inches((12,5))
for i,j in enumerate(ax.flat):
    #markerline, stemlines, baseline = j.stem(arange(Nf),abs(R[idx[:,i],0])/sqrt(Nf))
    #setp(markerline, 'markersize', 3.)
    j.fill_between(arange(Nf),1/sqrt(Nf)*abs(R[idx[:,i],0]).flat,0,alpha=0.3)
    markerline, stemlines, baseline =j.stem(arange(Nf),abs(X))
    setp(markerline, 'markersize', 4.)
    setp(markerline,'markerfacecolor','r')
    setp(stemlines,'color','r')
    j.axis('off')
    j.set_title('k=%d'%i,fontsize=8)
    
# fig.savefig('figure_00@.png', bbox_inches='tight', dpi=300)


fig,ax=subplots()
fig.set_size_inches((7,3))
ax.plot(a,abs(R[idx,0]*X)/sqrt(Nf), label=r'$|Z_k|$ = $X_k\otimes_N R_k$')
ax.plot(a, abs(Z),'o',label=r'$|Z_k|$ by DFT')
ax.set_xlabel('DFT index,k',fontsize=18)
ax.set_ylabel(r'$|Z_k|$',fontsize=18)
ax.set_xticks(arange(ax.get_xticks().max()))
ax.tick_params(labelsize=8)
ax.legend(loc=0)
ax.grid()
# fig.savefig('figure_00@.png', bbox_inches='tight', dpi=300)