import numpy as nupy
import matplotlib.pyplot as plt
import scipy.io.wavfile as scw
import scipy as scy

rate, data = scw.read("bokfink_c.wav")
sig = data[:,0]
dt = 1/rate
sz = data.shape
np = sz[0]
np2 = int(np/2)
ts = nupy.arange(np)*dt
npfft = 256
npfft2 = int(npfft/2)
ovlp = .3
novl = int((1.-ovlp)*npfft)
nfft = (np - npfft)/novl
amp = nupy.ones((npfft2+1, nfft))
nts = nupy.ones(nfft+1)
for i in range(int(nfft)):
    si = i*novl
    fi = si + npfft
    _sig = nupy.copy(sig[si:fi])*scy.hanning(npfft)
    nts[i] = ts[si]
    amp[:, i] = (nupy.abs(2*nupy.fft.fft(_sig)/np))[0:npfft2+1]
nts[nfft] = ts[-1]
freq = nupy.arange(npfft2 + 1)/(npfft*dt)

plt.subplot(2,1,1)
plt.plot(ts, sig)
plt.xlabel("Time [s]")
plt.ylabel("Amplitude")
plt.subplots_adjust(bottom=0.1, right=0.8, top=0.9)

plt.subplot(2,1,2)
plt.pcolor(nts, freq, nupy.log10(amp), cmap=plt.cm.jet, vmin=-2.5, vmax=1.)
plt.xlabel("Time [s]")
plt.ylabel("Frequency [Hz]")
plt.subplots_adjust(bottom=0.1, right=0.8, top=0.9)

cax = plt.axes([0.85, 0.1, 0.025, 0.37])
plt.colorbar(cax=cax)