# -*- coding: utf-8 -*-
"""
Created on Mon Jul 17 23:27:58 2017

@author: danjr
"""

'''
Script to turn a midi file into a piano roll-like numpy array.

It'll likely handle multiple simultaneous notes within a single midi track
incorrectly.
'''

import mido
import numpy as np
import matplotlib.pyplot as plt


# midi file to read
filepath = r'bach\chorales\01ausmei.mid'
midi_file = mido.MidiFile(filepath)

'''
Initialize the piano roll matrix
'''
num_notes = 100

# find all the tempos in the piece, and as the worst-case scenario assume the overall tempo is the fastest one
tempos = list(map(lambda x: x.tempo, filter(lambda message: message.type=='set_tempo', [item for sublist in midi_file.tracks for item in sublist])))
tempo = min(tempos)
#pretty sure this function just multiplies all these numbers together but oh well
num_ticks = int(mido.second2tick(midi_file.length,midi_file.ticks_per_beat, tempo))
piano_roll = np.zeros([num_notes,num_ticks*2])

for ti,track in enumerate(midi_file.tracks):    
    
    num_messages = len(track)
    
    '''
    Look for note_on/note_off pairs, and turn "on" the corresponding values in
    piano_roll.
    '''    
    for mi in range(num_messages):
        message = track[mi]
        
        if (message.type=='note_on') and (message.velocity!=0):
            '''
            This message designates the start of a note. Once the message
            designating the end of the note is found, the piano roll will be
            updated to include the note.
            '''
            on_message = message
            
            '''
            Find the tick marker for the start of the note. This is the sum of
            all the previous messages' timings, plus this timing.
            '''
            times_before_on = list(map(lambda message: message.time,filter(lambda message: message.type=='note_on' or message.type=='note_off',track[0:mi+1])))
            start_tick = np.sum(times_before_on)            
            
            # look for the corresponding "note_off" message in the subsequent messages
            found_note_off = False
            mii = mi + 1 # start by looking at the next message; this will likely be the "off" message
            while found_note_off == False:
                
                # check to see if this message is the "off" message
                if ((track[mii].type == 'note_off') | (hasattr(track[mii], 'velocity') and track[mii].velocity==0)) and (track[mii].note == message.note):
                    '''
                    The corresponding "off" message has been found.
                    Find the duration of the note (the starting and ending
                    ticks) and update piano_roll.
                    '''
                    
                    found_note_off = True
                    off_message = track[mii]
                    
                    '''
                    The note's off tick is the sum of all previous messages' 
                    ticks plus this timing.
                    '''
                    times_before_off = list(map(lambda message: message.time,filter(lambda message: message.type=='note_on' or message.type=='note_off',track[0:mii+1])))                    
                    end_tick = np.sum(times_before_off)
                    
                    piano_roll[on_message.note,start_tick:end_tick] = ti # for now, set the value to the track number so different tracks are distinguishable in the matrix plot

                else:
                    '''
                    This is not the corresponding "off" message. Update the
                    index to look at the next message.
                    '''
                    mii = mii+1
                    
                if mii == num_messages:
                    '''
                    In case it finds an "orphan" note_on message -- for now,
                    just ignore it.
                    '''
                    found_note_off = True # just so it stops looking
                    print('Warning, did not find an "off" message for the following message:')
                    print(message)
                    
'''
Remove the initial empty columns -- why do these exist?
'''                  
print('Removing initial empty columns.')      
sum_cols = np.sum(piano_roll,axis=0)
first_non_empty = np.nonzero(sum_cols)[0][0]
piano_roll = piano_roll[:,first_non_empty::]
print('The piano roll now has shape:')
print(piano_roll.shape)

'''
Remove the trailing empty columns. Shouldn't be an issue once dynamic sizing
is implemented.
'''
print('Removing trailing empty columns.')      
sum_cols = np.sum(piano_roll,axis=0)
last_non_empty = np.nonzero(sum_cols)[0][-1]
piano_roll = piano_roll[:,0:last_non_empty]
print('The piano roll now has shape:')
print(piano_roll.shape)

# plot the raw piano roll
fig = plt.figure()
ax_orig = fig.add_subplot(111)
ax_orig.imshow(piano_roll,aspect='auto',origin='bottom')
ax_orig.set_title('Original piano roll')

'''
See which columns are duplicates of the previous one.
'''
is_dup = np.zeros(piano_roll.shape[1])
for tick in range(piano_roll.shape[1]):
    if tick>0 and np.array_equal(piano_roll[:,tick],piano_roll[:,tick-1]):
        is_dup[tick] = 1
        
'''
Find the minimum spacing between columns that aren't unique. This will define
the maximum amount of down-sampling possible without losing information.
'''        
new_cols = np.where(is_dup==0) # indices of columns that aren't just duplicates of the previous
new_cols_diff = np.diff(new_cols)
resample_to = 30

# plot the distribution of spacings      
fig = plt.figure()
ax = fig.add_subplot(111)
ax.hist(new_cols_diff.transpose(),bins=np.arange(np.max(new_cols_diff)))
ax.set_ylabel('Number of occurances')
ax.set_xlabel('Spacing between non-duplicate columns in piano_roll')


'''
Downsample
'''

def any_notes_on(chunk):
    '''
    given a vertical slice of a piano roll, return which notes are on for any
    portion of the slice
    '''
    
    note_occurances = np.sum(chunk,axis=1)
    on_notes = np.array(list(map(lambda x: 1 if x>0 else 0, note_occurances)))
    return on_notes

# with resample_to defined, get the divisions of the "chunks" to downsample
spacings = np.arange(0,np.shape(piano_roll)[1],resample_to,dtype=int)
end_idxs = spacings[1::]
start_idxs = spacings[:-1]
# illustrate the chunks on the original matrix plot of the piano roll.
# this can be used to check visually that they downsampling is in the right place
[ax_orig.axvline(s) for s in spacings]

# for each chunk generate the resampled array, and add to the resampled roll
piano_roll_resampled = np.zeros([num_notes,len(start_idxs)])
for i,start_idx in enumerate(start_idxs):
    end_idx = end_idxs[i]
    chunk = piano_roll[:,start_idx:end_idx-1]
    on_notes = any_notes_on(chunk)
    piano_roll_resampled[:,i] = on_notes

print('The downsampled piano roll has shape:')
print(piano_roll_resampled.shape)
    
# Plot the resampled piano roll.
fig = plt.figure()
ax = fig.add_subplot(111)
ax.imshow(piano_roll_resampled,aspect='auto',origin='bottom')
ax.set_title('Resampled piano roll')