import settings
import pypianoroll as roll
import matplotlib.pyplot as plt
import numpy as np
import os
from math import floor

MIDI_DIRECTORY = settings.midi_path
SAMPLES_DIRECTORY = settings.samples_path
MIDI_RESOLUTION = settings.midi_resolution
BEAT_PER_BATCH = settings.beats_per_sample

samples = np.empty((0,BEAT_PER_BATCH,96,128))

def erase_note_lenth(pianoroll):
    if pianoroll.ndim != 2:
        raise ValueError('pianoroll should be two dimentional')
    now_block = []
    for x in pianoroll:
        this = None
        prev = None
        new_line =[]
        for y in x:
            this = y
            if prev != None:
                if this > 0 and prev > 0:
                    new_line.append(0)
                else:
                    new_line.append(y)
            else:
                new_line.append(y)
            prev = this
        now_block.append(new_line)
    return np.array(now_block)

print('Start convertion')
for midi_file in os.listdir(MIDI_DIRECTORY):
    try:
        print('Reading file: {}'.format(midi_file))
        song = roll.Multitrack('{}/{}'.format(MIDI_DIRECTORY, midi_file))
        # no_drums_mt = roll.Multitrack(tempo=120.0, downbeat=[0, 96, 192, 288], beat_resolution=24)
        intruments_only = roll.Multitrack(tempo=120.0, beat_resolution=24)

        for track in song.tracks:
            if track.is_drum == False:
                print(track.name, track.program)
                intruments_only.append_track(track=track, pianoroll=track.pianoroll)
                instrument_track = track.pianoroll

                # plt.imshow(instrument_track[24*8:24*24].T)
                # plt.savefig('data/0_{}.png'.format(midi_file))

                instrument_track = erase_note_lenth(instrument_track.T).T
                # plt.imshow(instrument_track[24*8:24*24].T)
                # plt.savefig('data/1_{}.png'.format(midi_file))


        # instruments = no_drums_mt.get_merged_pianoroll(mode='sum')

                beats = floor( (instrument_track.shape[0] / MIDI_RESOLUTION) / BEAT_PER_BATCH) * BEAT_PER_BATCH
                notes_for_beats = beats * MIDI_RESOLUTION

                print('beats: ', beats)
                samples_of_song = np.asarray(np.split(instrument_track[:notes_for_beats], beats))
                samples_of_song = samples_of_song.reshape(int(beats/BEAT_PER_BATCH),BEAT_PER_BATCH,96,128)

                print('Converted samples: {}'.format(samples_of_song.shape))
                samples = np.concatenate([samples_of_song,samples], axis=0)
        np.savez_compressed(SAMPLES_DIRECTORY,samples)

    except Exception as error:
        print('Convertion faild: {}'.format(error))
        pass

    finally:
        print('Done!')