import pylab
import numpy
import scipy
import scipy.signal

size = 4096
speed = 512

fft = {}
window = {}
sinus = {}

def printer(function):
    window[function] = eval('scipy.signal.%s(size)' % function)
    pylab.plot(numpy.abs(window[function]))
    pylab.xlim(0, size-1)
    pylab.savefig('%s-%03d.png' % (function, speed))
    pylab.close()

    fft[function] = scipy.fft(window[function])
    pylab.loglog(numpy.abs(fft[function])[0:size/2+1])
    pylab.savefig('%s-fft-%03d.png' % (function, speed))
    pylab.close()

    sinus[function] = scipy.fft(y * window[function])
    pylab.loglog(numpy.abs(sinus[function])[0:size/2+1])
    pylab.savefig('%s-sinus-%03d.png' % (function, speed))
    pylab.close()

x = numpy.array(range(size))*(numpy.pi/speed)
y = numpy.sin(x)
pylab.plot(x, y)
pylab.xlim(0, size * numpy.pi/speed)
pylab.savefig('sinus-%03d.png' % speed)
pylab.close()

printer('boxcar')
printer('hanning')
printer('hamming')
printer('blackman')
printer('bartlett')
printer('triang')
printer('parzen')
printer('bohman')
printer('nuttall')
printer('blackmanharris')
printer('flattop')
printer('barthann')

colors = ('b', 'r', 'g', 'c', 'y', 'm', 'k')
plots = ('boxcar', 'hanning', 'hamming', 'blackman', 'bohman', 'blackmanharris','nuttall')

for plot, color in zip(plots, colors):
    pylab.loglog(numpy.abs(fft[plot])[0:size/2+1], c = color, label=plot)
pylab.legend(loc=3)
pylab.savefig('fft-%03d.png' % speed)
pylab.close()

for plot, color in zip(plots, colors):
    pylab.plot(numpy.abs(window[plot]), c = color, label=plot)
pylab.legend(loc=0)
pylab.savefig('window-%03d.png' % speed)
pylab.close()

for plot, color in zip(plots, colors):
    pylab.loglog(numpy.abs(sinus[plot][0:size/2+1]), c = color, label=plot)
pylab.legend(loc=3)
pylab.savefig('sinus-fft-%03d.png' % speed)
pylab.close()
