from __future__ import  division
from scipy import signal

# 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
        return 20*log10(abs( fft.fft(W,Nfft)/sqrt(Nfft) ) )

fs = 64 # sampling frequency
t = arange(0,2,1/fs)
f = 10  # one signal
deltaf = 2.3 # second nearby frequency

Nf = 512
fig,ax = subplots(2,1,sharex=True,sharey=True)
fig.set_size_inches((9,4))

x=10*cos(2*pi*f*t) + 10*cos(2*pi*(f+deltaf)*t) # equal amplitudes
X = fft.fft(x,Nf)/sqrt(Nf)
ax[0].plot(linspace(0,fs,len(X)),abs(X),'-o',ms=3.)
ax[0].set_ylabel(r'$|X(k)|$',fontsize=18)
ax[0].set_xlim(xmax = fs/2)
ax[0].grid()
ax[0].text(0.5,0.5,'Same amplitudes',
            transform=ax[0].transAxes,
            backgroundcolor='Lightyellow',
            fontsize=16)

x=cos(2*pi*f*t) + 10*cos(2*pi*(f+deltaf)*t) # one has 10x the amplitude
X = fft.fft(x,Nf)/sqrt(Nf)
ax[1].plot(linspace(0,fs,len(X)),abs(X),'-o',ms=3.)
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].grid()
ax[1].text(0.5,0.5,'Different amplitudes',
                transform=ax[1].transAxes,
                backgroundcolor='lightyellow',
                fontsize=16)
ax[1].annotate('Smaller signal',
            fontsize=12,xy=(f,abs(X)[int(f/fs*Nf)]),
            xytext=(2,15),
            arrowprops={'facecolor':'m','alpha':.3});

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

fig,axs = subplots(2,1)
fig.set_size_inches((10,5))

Nf = 128
nsamp = 64
window = signal.triang(nsamp)
rectwin = ones(nsamp)
ax = axs[0]

ax.plot(arange(nsamp),window,'m-o',label='triangular')
ax.stem(arange(nsamp),rectwin,basefmt='b-',label='rectangular')
ax.set_ylim(ymax = 1.1,ymin=-0.1)
ax.legend(loc=0)

x=array(dftmatrix(64,64)[:,5].real).flatten() # convert to numpy array
n = arange(len(x))
window = signal.triang(len(x))

ax = axs[1]
ax.plot(n,x,'-o',label='signal',ms=3.,alpha=0.3)
ax.plot(n,window*x,'-or',label='windowed\nsignal',ms=5.)
ax.set_ylabel('Amplitude',fontsize=16)
ax.set_xlabel('sample index',fontsize=16)
ax.legend(loc=0)
ax.grid()
ax2 = ax.twinx()
ax2.fill_between(n,window,alpha=0.2,color='m')
ax2.set_ylabel('Window function',fontsize=16);

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


Nf = 512
fig,ax = subplots(3,1,sharex=True,sharey=True)
fig.set_size_inches((10,6))

x=cos(2*pi*f*t) + 10*cos(2*pi*(f+deltaf)*t)
X = fft.fft(x,Nf)/sqrt(Nf)
ax[0].plot(linspace(0,fs,len(X)),db20(X),'-o',ms=3.)
ax[0].set_xlim(xmax = fs/2)
ax[0].grid()
ax[0].text(0.5,0.5,'rectangular window',
                transform=ax[0].transAxes,
                backgroundcolor='lightyellow',
                fontsize=16)
ax[0].annotate('Smaller signal',
            fontsize=12,xy=(f,db20(X)[int(f/fs*Nf)]),
            xytext=(1,30),
            arrowprops={'facecolor':'m'})

w = signal.triang(len(x))
X = fft.fft(x*w,Nf)/sqrt(Nf)
ax[1].plot(linspace(0,fs,len(X)),db20(X),'-o',ms=3.)
ax[1].set_xlim(xmax = fs/2)
ax[1].grid()
ax[1].text(0.5,0.5,'triangular window',
                transform=ax[1].transAxes,
                backgroundcolor='lightyellow',
                fontsize=16)
ax[1].annotate('Smaller signal',
            fontsize=12,xy=(f,db20(X)[int(f/fs*Nf)]),
            xytext=(1,30),
            arrowprops={'facecolor':'m'})

w = signal.hamming(len(x))
X = fft.fft(x*w,Nf)/sqrt(Nf)
ax[2].plot(linspace(0,fs,len(X)),db20(X),'-o',ms=3.)
ax[2].set_xlabel('Frequency (Hz)',fontsize=18)
ax[2].set_xlim(xmax = fs/2)
ax[2].grid()
ax[2].set_ylim(ymin=-20)
ax[2].text(0.5,0.5,'Hamming window',
                transform=ax[2].transAxes,
                backgroundcolor='lightyellow',
                fontsize=16)
ax[2].annotate('Smaller signal',
            fontsize=12,xy=(f,db20(X)[int(f/fs*Nf)]),
            xytext=(1,30),
            arrowprops={'facecolor':'m'});

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