#!/usr/bin/env python3
# vim:ts=4:sts=4:sw=4:expandtab

import itertools
import matplotlib.pyplot as plt 
import numpy
import wave
import sys
import time
import struct

import audio

ANALYZE_LENGTH = 0.12
ANALYZE_TOTAL_LENGTH = 30

numpy.set_printoptions(threshold=sys.maxsize)

source = None
if len(sys.argv) > 1:
    source = audio.wave_open(sys.argv[1])
else:
    source = audio.mic_open()

fig = plt.figure()

#for i in range(int(ANALYZE_TOTAL_LENGTH/ANALYZE_LENGTH)):
for i in itertools.count():

    res = audio.audio_read(source, int(ANALYZE_LENGTH * audio.FRAMERATE))
    if len(res) != int(ANALYZE_LENGTH * audio.FRAMERATE):
        break
    fre = numpy.fft.rfft(res)
    fre = numpy.abs(fre)
    if i*ANALYZE_LENGTH >= 1 and int(i*ANALYZE_LENGTH) != int((i-1)*ANALYZE_LENGTH):
        fig.clear()
        a = fig.add_subplot(311)
        b = fig.add_subplot(312)
        c = fig.add_subplot(313)
        a.plot(range(len(res)), res)
        b.plot(range(len(fre)), fre)
        c.plot(range(len(fre[0:200])), fre[0:200])
        fig.canvas.draw()
        plt.pause(0.001)

    prob = fre/numpy.sum(fre)
    exp = (numpy.arange(len(prob))*prob).sum()
    var = (numpy.abs((numpy.arange(len(prob))-exp))*prob).sum()
    print(exp/ANALYZE_LENGTH, var)
    print(numpy.argmax(fre)/ANALYZE_LENGTH)

plt.show()

audio.audio_close(source)