#! /usr/bin/env python

import sys, os.path
import numpy as np
import aubio

if __name__ == '__main__':

    if len(sys.argv) < 2:
        print ("usage: %s <inputfile>" % sys.argv[0])
        sys.exit(1)
    input_filename = sys.argv[1]

    win_size = 8192 #2048 #1024
    hop_size = win_size // 2
    samplerate = 44100
    bins_per_octave = 72 # 24 # 6, 12, 24, 36, 48, 72...

    a_source = aubio.source(input_filename, samplerate, hop_size)
    a_pvoc = aubio.pvoc(win_size, hop_size)
    a_constantq = aubio.constantq(win_size, samplerate, bins_per_octave)

    numbins = a_constantq.get_numbins()

    print ('created constant-q with %d win_size, %d numbins, %d Hz' %
            (win_size, numbins, samplerate))

    read = hop_size

    cqts = np.ndarray((0,numbins)).astype(aubio.float_type)

    while read:
        samples, read = a_source()
        fftgrain = a_pvoc(samples)
        cqt = a_constantq(fftgrain)
        cqts = np.vstack([cqts, cqt])

    print ('finished computing constant-q over %d frames of %d bins' % cqts.shape)
    try:
        import matplotlib.pyplot as plt
        import matplotlib.ticker as ticker
        plt.rc('font', size=8)
        plt.rc('axes', titlesize=8)

        fig, ax = plt.subplots(1)
        ax.imshow(cqts.T, aspect='auto', origin='bottom', cmap=plt.cm.gray_r)
        ax.axis([0, len(cqts), 0, len(cqts[0])])

        time_step = hop_size / float(samplerate)
        freq_start = 19.0 # TODO get actual low_freq
        freq_step = 12./bins_per_octave
        freq_end = freq_start + len(cqts[0]) * freq_step
        #print ('plotting', end=" ")
        #print ('x from 0 to %.2f (sec),' % (time_step * len(cqts)), end=" ")
        #print ('y from %.2f to %.2f (midi note)' % (freq_start, freq_end))

        ax.fmt_xdata = lambda x: "%.1f" % (x * time_step)
        x_formatter = ticker.FuncFormatter(lambda x, pos: ax.fmt_xdata(x) )
        ax.xaxis.set_major_formatter(x_formatter)
        if 0:
            total_duration = time_step * len(cqts)
            if total_duration >= 60: rate = 10.
            elif total_duration >= 10: rate = 5.
            elif total_duration >= 3: rate = 1.
            elif total_duration >= 1: rate = .2
            elif total_duration < 1: rate = .1
            #ax.xaxis.set_major_locator(ticker.FixedLocator(np.arange(0, len(cqts), rate/time_step)))
            ax.xaxis.set_major_locator(ticker.MultipleLocator(base=rate/time_step))
        else:
            ax.xaxis.set_major_locator(ticker.AutoLocator())
            #ax.xaxis.set_major_locator(ticker.MaxNLocator(nbins=10, symmetric=True))
            #ax.xaxis.set_major_locator(ticker.MaxNLocator(integer=True, steps=[1]))
            #ax.xaxis.set_major_locator(ticker.MultipleLocator(base=1./time_step))

        ax.set_xlabel ('Time (s)')

        ax.fmt_ydata = lambda y: "%.1f" % (y * freq_step + freq_start)
        y_formatter = ticker.FuncFormatter(lambda y, pos: ax.fmt_ydata(y))
        ax.yaxis.set_major_formatter(y_formatter)
        ax.yaxis.set_major_locator(ticker.MultipleLocator(base=bins_per_octave))
        ax.set_ylabel ('Frequency (midi)')

        ax.set_title('Constant-Q tranform of %s' % os.path.basename(input_filename))

        plt.show()
    except ImportError:
        print ('install matplotlib to plot results')