import scipy
import scipy.fftpack as fftpack
import scipy.io.wavfile as wavfile
import numpy as np
import matplotlib.pyplot as plt

# Get MidiFile package from https://code.google.com/p/midiutil/
import MidiFile3

#############################
# Functions
#############################

def load_wav(input_wav):
    "Read the wave file and convert the input to a mono-wave if needed"
    rate, src = wavfile.read(input_wav)
    data = src
    if len(src.shape)==2:
        data = (src[:,0]+src[:,1])/2
    
    startframe = 0
    while data[startframe]==0:
        startframe+=1
    endframe = -1
    while data[endframe]==0:
        endframe-=1
    
    return rate, data[startframe:endframe]

def process_wav():
    "Analyize the wave with FFT"
    for idx in range(totalsteps):
        print('Analyizing slice #',idx,'/',totalsteps)
        
        for pitch in range(21,109):
            target_freq = 2**((pitch-69.)/12.)*440
            
            minframes = rate/target_freq
            multiplier = 2**int(scipy.ceil(scipy.log(float(baseframes)/float(minframes))/scipy.log(2)))
            multiplier = max(multiplier,2)
            multiplier = min(multiplier,512)
            
            anaframes = int(scipy.ceil(minframes*multiplier))
            anaperiod = float(anaframes)/float(rate)
            
            fft = abs(scipy.fft(data[idx*baseframes:idx*baseframes+anaframes]))
            
            if fft[multiplier]>0.:
                out[pitch-21][idx] = scipy.log10(fft[multiplier])

def cleaning_wav():
    "Clean the results of FFT by substracting the pedestal and others"
    for idx in range(totalsteps):
        sort = np.sort(out[:,idx])
        mean = np.mean(sort[:-16])
        std = np.std(sort[:-16])
        for key in range(88):
            out[key][idx] -= mean # pedestal substraction
            if out[key][idx]<std*2.5: # nxsigma theshold
                out[key][idx] = 0.
        for key in range(87,0,-1): # cleaning nearby notes
            if out[key][idx]<out[key-1][idx]: out[key][idx] = 0.
        for key in range(87):
            if out[key][idx]<out[key+1][idx]: out[key][idx] = 0.

    for idx in range(totalsteps): # remove single isolated note
        for key in range(88):
            if out[key][idx]>0.:
                if idx==0 and out[key][idx+1]==0.:
                    out[key][idx] = 0.
                elif idx==totalsteps-1 and out[key][idx-1]==0.:
                    out[key][idx] = 0.
                elif out[key][idx-1]==0. and out[key][idx+1]==0.:
                    out[key][idx] = 0.

    for idx in range(0,totalsteps-2): # remove single isolated hole
        for key in range(88):
            if out[key][idx]>0. and out[key][idx+1]==0. and out[key][idx+2]>0.:
                out[key][idx+1] = (out[key][idx]+out[key][idx+2])/2.

    for idx in range(totalsteps): # remove harmonics
        for key in range(87,11,-1):
            if out[key][idx]>0. and out[key-12][idx]>0. and out[key][idx]/out[key-12][idx]<0.5:
                out[key][idx] = 0.
        for key in range(88-12):
            if out[key][idx]>0. and out[key+12][idx]>0. and out[key][idx]/out[key+12][idx]<0.5:
                out[key][idx] = 0.

def export_midi(output_midi):
    "Save the result to a MIDI file."
    maxamp = np.amax(out)
    minamp = np.amin(out[np.nonzero(out)])
    MyMIDI = MidiFile3.MIDIFile(2)
    
    for track in range(2):
        MyMIDI.addTrackName(track,0,'Track'+str(track))
        MyMIDI.addTempo(track,0, tempo)

    for pitch in range(21,109):
        track = 0
        if pitch < 60: track = 1
        channel = 0
        time = duration = volume = 0.
        
        for idx in range(totalsteps):
            
            if out[pitch-21][idx]>0.:
                if duration==0.:
                    time = float(idx)/basenotefactor
                duration += 1./basenotefactor
                volume += out[pitch-21][idx]
            if duration>0.:
                if out[pitch-21][idx]==0. or idx==totalsteps-1:
                    volume /= (duration*basenotefactor)
                    volume = (volume-minamp)/(maxamp-minamp)*107.+20. # renormalize the volume based on the amplitudes
                    
                    MyMIDI.addNote(track,channel,pitch,time,duration,int(volume))
                    duration = volume = 0.

    binfile = open(output_midi, 'wb')
    MyMIDI.writeFile(binfile)
    binfile.close()

def mkplot_wave(frames):
    
    t = scipy.linspace(0,float(frames)/float(rate),frames, endpoint=False)
    plt.plot(t, data[:frames])
    plt.axis([t.min(), t.max(), data[:frames].min()*1.1, data[:frames].max()*1.1])
    plt.xlabel('Time [sec]')
    plt.ylabel('Amplitude')
    plt.show()

def mkplot_pitch(steps):
    
    dy, dx = 1.,1.
    y, x = np.mgrid[slice(0., 88+dy, dy),slice(0., steps+dx, dx)]
    plt.pcolor(x, y, out[:,:steps], cmap = 'Blues')
    plt.axis([x.min(), x.max(), y.min(), y.max()])
    plt.xlabel('Time [1/32 beats]')
    plt.ylabel('Key')
    plt.colorbar(orientation='horizontal')
    plt.show()


#############################
# Main program starts
#############################

# defining the global variables
# naively set the tempo to 120.
tempo = 120.0
beats_per_sec = tempo/60.
basenotefactor = 8.

rate, data = load_wav('humoresque.wav')

totalframes = len(data)
baseframes = int(scipy.ceil(rate/beats_per_sec/basenotefactor))
totalsteps = totalframes//baseframes

# plotting the wave
#mkplot_wave(totalframes)

# Create the buffer to save the results of FFT
out = np.zeros((88,totalsteps))

# The core FFT work
process_wav()

# cleaning
cleaning_wav()

# plotting the pitch table
#mkplot_pitch(totalsteps)

# and save as a MIDI file...
export_midi('output.mid')