import abjad as abj 
from IPython.display import Audio
from subprocess import call
from audioread import audio_open
import scipy
import numpy as np

%load_ext abjad.ext.ipython
%matplotlib inline
%pylab inline
rc('font',**{'family':'sans-serif','sans-serif':['Helvetica'], 'size': 22})
rc('text', usetex=True)

# Bach in G Major
bach_treble = abj.Staff("d''g'8a'b'c''d''4 g'-.g'-. e''c''8d''e''fs''g''4 g'-.g'-. \\\\
                        c''d''8c''b'a'b'4 c''8b'a'g'fs'4 g'8a'b'g'a'2.")
abj.attach(abj.KeySignature('g'), bach_treble)
abj.attach(abj.TimeSignature((3,4)), bach_treble)
abj.attach(abj.spannertools.PhrasingSlur(), bach_treble[0:6])
abj.attach(abj.spannertools.PhrasingSlur(), bach_treble[8:14])
abj.attach(abj.spannertools.PhrasingSlur(), bach_treble[16:22])
abj.attach(abj.spannertools.PhrasingSlur(), bach_treble[22:27])
abj.attach(abj.spannertools.PhrasingSlur(), bach_treble[27:32])

bach_bass = abj.Staff("<g b>2 a4 b2. c' b a2 fs4 g2 c'4 d' b g d' d8 (c' b a)")
abj.attach(abj.KeySignature('g'), bach_bass)
abj.attach(abj.Clef('bass'), bach_bass)

# we're only going to play with the treble score,
# change to the commented line to also get the bass score
#bach = abj.scoretools.Score([bach_treble, bach_bass])
bach = abj.scoretools.Score([bach_treble])

abj.show(bach)

sffile = {'sine':'Sine Wave.sf2', 'piano':'JV1080_Nice_Piano_m.sf2'}
signal = {}
sr = 0

def load_audio_from_file(filename):
    # load the audio data from the file
    temp = bytearray()
    sr = 0
    with audio_open(filename) as audio:
        sr = audio.samplerate
        nchannels = audio.channels
    
        # chunk size can be specified with 'block_samples' (default 1024):
        for chunk in audio.read_data():
            temp.extend(chunk)

    data = frombuffer(temp, dtype='<i2').reshape(-1, nchannels)[:,0]   
    return sr, data


for instrument in sffile:
    midipath = instrument + '.midi'
    abj.persist(bach).as_midi(midipath)
    wavpath  = instrument + '.wav'
    sfpath = sffile[instrument]

    if call(['fluidsynth', '-C', 'no', '-R', 'no', '-g', '1', '-F', wavpath, sfpath, midipath]):
        raise NameError("call to fluidsynth failed")

    sr, signal[instrument]=load_audio_from_file(wavpath)

n_fft=2**14# window_size # default is 2048, pass it to function if you change it
n_chroma=2**12
overlap = 2**2
hop=n_fft/overlap # default is 512, same as above
n_mels=128 # default is 128
maxf = 860; # max_frequency
minf = 300; # min_frequency
maxfi = int((n_fft*maxf)/sr)
minfi = int((n_fft*minf)/sr)
maxf = sr*(maxfi-0.5)/float(n_fft)
minf = sr*(minfi-0.5)/float(n_fft)

imshow_options = {
    'aspect':'auto', 
    'origin':'lower',
    'cmap':cm.pink, 
    'interpolation':'kaiser'
}

def frequency2index(f, n_fft=n_fft):
    return int(n_fft * f / sr)

# TODO: this one still doesn't work correctly
def index2second(i):
    return n_fft/float(2*sr) + i*hop/float(sr)

def stft(signal, fftsize=1024, overlap=4):  
    # short_time_fourier_transform
    hop = fftsize/overlap
    w = scipy.hamming(fftsize)   
    return np.array([np.fft.fft(w*signal[i:i+fftsize]) 
                     for i in range(0, len(signal)-fftsize, hop)])

pylab.rcParams['figure.figsize'] = (0.75*16.18, 0.75*10.0)

Si={}
S={}
for instrument, data in signal.iteritems():
    Si[instrument] = stft(data, fftsize=n_fft, overlap=overlap).T
    S[instrument] = abs(Si[instrument])**2
    
notes = {
    "$D_3$":146.83, "$G_3$":195.99, "$D_4$":293.66, "$F\sharp_4$":369.99, 
    "$G_4$":392.00, "$G\sharp_4$":415.30, "$A_4$":440.00, "$A\sharp_4$":466.16, "$B_4$":493.88, 
    "$C_5$":523.25, "$C\sharp_5$":554.37, "$D_5$":587.32, "$D\sharp_5$":622.25, "$E_5$":659.25,
    "$F_5$":698.46, "$F\sharp_5$":739.99, "$G_5$":783.99, "$G\sharp_5$":830.61
}

Audio('sine.wav')

imshow(
    S['sine'][minfi:maxfi, :],
    extent=(
        index2second(0),
        # TODO: change the following line to be instrument independent
        index2second(len(S[instrument].T)-1),
        minf,
        maxf
    ), 
    **imshow_options
)

labels = ["$G_4$", "$B_4$", "$D_5$", "$E_5$", "$G_5$"]

_ = yticks( [notes[label] for label in labels] + [minf, maxf])
_ = xticks( range(3, 24, 3))
ylabel("frequency (Hz)")
xlabel("time (s)")
grid(color='w')

axr = twinx()
axr.set_ylim([minf, maxf])
_ = yticks( [notes[label] for label in labels], labels)

Audio('piano.wav')

imshow(
    S['piano'][minfi:maxfi, :],
    extent=(
        index2second(0),
        index2second(len(S[instrument].T)-1),
        minf,
        maxf
    ), 
    **imshow_options
)

labels = ["$G_4$", "$B_4$", "$D_5$", "$E_5$", "$G_5$"]

_ = yticks( [notes[label] for label in labels] + [minf, maxf])
_ = xticks( range(3, 24, 3))
ylabel("frequency (Hz)")
xlabel("time (s)")
grid(color='w')

axr = twinx()
axr.set_ylim([minf, maxf])
_ = yticks( [notes[label] for label in labels], labels)

Audio('honk_clean.wav')

sr, honk_data = load_audio_from_file('honk_clean.wav')

maxf = 6000; # max_frequency
minf = 200; # min_frequency
maxfi = int((n_fft*maxf)/sr)
minfi = int((n_fft*minf)/sr)
maxf = sr*(maxfi-0.5)/float(n_fft)
minf = sr*(minfi-0.5)/float(n_fft)

Si['honk_clean'] = stft(honk_data, fftsize=n_fft, overlap=overlap).T
S['honk_clean'] = abs(Si['honk_clean'])**2

imshow(
    S['honk_clean'][minfi:maxfi, :],
    extent=(
        index2second(0),
        # TODO: change the following line to be instrument independent
        index2second(len(S['honk_clean'].T)-1),
        minf,
        maxf
    ), 
    **imshow_options
)

ylabel("frequency (Hz)")
xlabel("time (s)")

Audio('honk_harmonics.wav')

sr, honk_data = load_audio_from_file('honk_harmonics.wav')

maxf = 6000; # max_frequency
minf = 200; # min_frequency
maxfi = int((n_fft*maxf)/sr)
minfi = int((n_fft*minf)/sr)
maxf = sr*(maxfi-0.5)/float(n_fft)
minf = sr*(minfi-0.5)/float(n_fft)

Si['honk_harmonics'] = stft(honk_data, fftsize=n_fft, overlap=overlap).T
S['honk_harmonics'] = abs(Si['honk_harmonics'])**2

imshow(
    S['honk_harmonics'][minfi:maxfi, :],
    extent=(
        index2second(0),
        # TODO: change the following line to be instrument independent
        index2second(len(S['honk_harmonics'].T)-1),
        minf,
        maxf
    ), 
    **imshow_options
)

ylabel("frequency (Hz)")
xlabel("time (s)")