from __future__ import division
from scipy import signal

fo = 1 # signal frequency
fs = 32 # sample frequency
Ns = 32 # number of samples
x = sin( 2*pi*fo/fs*arange(Ns)) # sampled signal
fig,axs= subplots(3,2,sharex='col',sharey='col')
fig.set_size_inches((12,6))
subplots_adjust(hspace=.3)

ax=axs[0,0]
ax.plot(arange(Ns),x,label='signal')
ax.plot(arange(Ns)+Ns,x,label='extension')
ax.set_ylabel('Amplitude',fontsize=14)
ax.set_title('Continuous at endpoints',fontsize=16)

ax=axs[0,1]
N=Ns  #chosen so DFT bin is exactly on fo
Xm = abs(fft.fft(x,N))
idx = int(fo/(fs/N))
ax.stem(arange(N)/N*fs,Xm,basefmt='b-')
ax.plot( fo, Xm[idx],'o')
ax.set_ylabel(r'$|X_k|$',fontsize=18)
ax.set_xlim(xmax = 5)
ax.set_ylim(ymin=-1)
ax.text(0.3,0.5,'No spectral leakage',fontsize=16,
        transform=ax.transAxes,
        bbox={'fc':'y','alpha':.3})
ax.grid()

# loss of continuity case

fo = 1.1 # signal frequency
x = sin( 2*pi*fo/fs*arange(Ns)) # sampled signal

ax=axs[1,0]
ax.plot(arange(Ns),x,label='signal')
ax.plot(arange(Ns)+Ns,x,label='extension')
ax.set_ylabel('Amplitude',fontsize=14)
ax.set_title('Discontinuous at endpoints',fontsize=16)

ax=axs[1,1]
Xm = abs(fft.fft(x,N))
idx = int(fo/(fs/N))
ax.stem(arange(N)/N*fs,Xm,basefmt='b-')
ax.plot( fo, Xm[idx],'o')
ax.set_xlabel('Frequency(Hz)',fontsize=16)
ax.set_ylabel(r'$|X_k|$',fontsize=18)
ax.text(0.3,0.5,'Spectral Leakage',fontsize=16,
        transform=ax.transAxes,
        bbox={'fc':'y','alpha':.3})
ax.set_xlim(xmax = 5)
ax.set_ylim(ymin=-1)
ax.grid()

x = x*signal.triang(Ns,2)
ax=axs[2,0]
ax.plot(arange(Ns),x,label='signal')
ax.plot(arange(Ns)+Ns,x,label='extension')
ax.set_xlabel('sample index',fontsize=14)
ax.set_ylabel('Amplitude',fontsize=14)
ax.set_title('Window restores continuity at endpoints',fontsize=14)

ax=axs[2,1]
Xm = abs(fft.fft(x,N))
idx = int(fo/(fs/N))
ax.stem(arange(N)/N*fs,Xm,basefmt='b-')
ax.plot( fo, Xm[idx],'o')
ax.set_xlabel('Frequency(Hz)',fontsize=16)
ax.set_ylabel(r'$|X_k|$',fontsize=18)
ax.text(0.3,0.5,'Leakage with window',fontsize=16,
        transform=ax.transAxes,
        bbox={'fc':'y','alpha':.3})
ax.set_xlim(xmax = 5)
ax.set_ylim(ymin=-1)
ax.grid()

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


# some useful functions 
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)

def db20(W,Nfft=None):
    'Given DFT, return power level in dB'
    if Nfft is None: # assume W is DFT 
        return 20*log10(abs(W))
    else: # assume time-domain passed, so need DFT
        DFT= fft.fft(array(W).flatten(),Nfft)/sqrt(Nfft)
        return 20*log10(abs(DFT.flatten()))

U=dftmatrix(64) 
u=U[:,6].real*sqrt(2) # create test sinusoid
fo = 2*pi/64*6 # in radians/sec
nz=randn(64,1) # noise samples
w=signal.triang(64) # window function

fig,ax= subplots(2,1)
fig.set_size_inches((10,5))
subplots_adjust(hspace=.3)
n = arange(len(u))
ax[0].plot(n,u.real,label='before window',lw=2)
ax[0].set_ylabel('Amplitude',fontsize=16)
ax[0].plot(n,diag(w)*u.real,label='after window',lw=3.)
ax[0].fill_between(n,array(u).flat, array(diag(w)*u).flat,alpha=0.3)
ax[0].legend(loc=0,fontsize=12)
ax[0].set_xlabel('n')
ax[0].grid()
ax[0].annotate('Lost signal due to window',fontsize=14, bbox={'fc':'b','alpha':.3},
      xy=(11,0.1),
      xytext=(30,40), textcoords='offset points',
      arrowprops={'facecolor':'b','alpha':.3})

N=256 # DFT size for plot
idx = int(fo/(2*pi/N))
ax[1].plot(db20(u,N),label='DFT before window')
ax[1].plot(db20(diag(w)*u,N),label='DFT after window')
ax[1].set_ylim(ymin=-40,ymax=-5)
ax[1].set_xlim(xmax=40)
ax[1].set_ylabel(r'$20\log_{10}|X_k|$',fontsize=12)
ax[1].set_xlabel(r'$k$',fontsize=14)
ax[1].annotate('Loss of signal power',fontsize=14,xy=(22,-13),
            xytext=(2,-15),
            arrowprops={'facecolor':'m','alpha':.3})
pkU = db20(u,N)[idx]
pkW = db20(diag(w)*u,N)[idx]
ax[1].annotate('',xy=(idx,pkW),xytext=(idx,pkU),fontsize=12,
               arrowprops={'arrowstyle':'<->','color':'m'})
ax[1].legend(loc=0,fontsize=12)
ax[1].grid()

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


from matplotlib.patches import Rectangle

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

N = 256 # DFT size
idx = int(fo/(2*pi/N))
Xm = abs(fft.fft(array(diag(w)*u).flatten(),N)/sqrt(N))**2
ax.plot(Xm,'-o')
ax.add_patch(Rectangle((idx-10/2,0),width=10,height=Xm[idx],alpha=0.3))
ax.set_xlim(xmax = N/4)
ax.set_ylabel(r'$|W_k|^2$',fontsize=18)
ax.set_xlabel(r'$k$',fontsize=18)
ax.set_title('Equivalent Noise Bandwidth',fontsize=18)
ax.annotate('Area of rectangle\n is windowed\nnoise power\n'\
              +r'$\sigma_\nu^2 \mathbf{w}^T \mathbf{w}$',
            fontsize=14,
            xy=(idx,Xm.max()/2.),
            xytext=(40,Xm.max()/2.),
            arrowprops={'facecolor':'m','alpha':.3});
ax.annotate('',ha='center',fontsize=24,
             xy=(idx+10/2,Xm.max()*1.05),
             xytext=(idx-10/2,Xm.max()*1.05),
             arrowprops=dict(arrowstyle='<->'))
ax.annotate('',ha='center',fontsize=24,
             xy=(15,0),
             xytext=(15,Xm.max()),
             arrowprops=dict(arrowstyle='<->'))
ax.text( 1, Xm.max()/2,r'$ |W_0|^2\sigma_\nu^2 $',fontsize=20,bbox={'fc':'gray','alpha':.3})
ax.text( idx-5, Xm.max()*1.1,r'$B_{neq}$',fontsize=18,bbox={'fc':'gray','alpha':.3})
ax.set_ylim(ymax = Xm.max()*1.2)
ax.grid()

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