#!/usr/bin/env python3

import settings
import pypianoroll as roll
import numpy as np
import os
from tqdm import tqdm
from math import floor
import sys
from collections import defaultdict
import pickle
from music21 import converter, instrument, note, chord, stream
import music21

midi_folder_path = sys.argv[1]
output_path = sys.argv[2]

def to_sequence(midi_path):
    seq_by_instrument = defaultdict( lambda : [] )
    midi_file = music21.converter.parse(midi_path)
    stream = music21.instrument.partitionByInstrument(midi_file)
    for part in stream:
        for event in part:
            if part.partName != None:
                # TODO: add note lenght as parameter
                if isinstance(event, music21.note.Note):
                    # to_export_event = (str(event.pitch), event.quarterLength)
                    to_export_event = str(event.pitch)
                    seq_by_instrument[part.partName].append(to_export_event)
                elif isinstance(event, music21.chord.Chord):
                    to_export_event = ' '.join(str(note) for note in event.pitches)
                    # to_export_event = (' '.join(str(note) for note in event.pitches), event.quarterLength)
                    seq_by_instrument[part.partName].append(to_export_event)

    X_train_by_instrument = defaultdict( lambda : [] )
    y_train_by_instrument = defaultdict( lambda : [] )

    for instrument, sequence in seq_by_instrument.items():
        for i in range(len(sequence)-8) :
            X_train_by_instrument[instrument].append(np.array(sequence[i: i + 8])) # <seq lenth
            y_train_by_instrument[instrument].append(np.array(sequence[i + 8]))
            # TODO: Notes to integers

    return X_train_by_instrument, y_train_by_instrument

def main():
    print('Exporting...')

    train_X = defaultdict( lambda : [] )
    train_y = defaultdict( lambda : [] )

    for directory, subdirectories, files in os.walk(midi_folder_path):
        for midi_file in tqdm(files):
            midi_file_path = os.path.join(directory, midi_file)
            _X_train, _y_train = to_sequence(midi_file_path)

            for (X_key, X_value), (y_key, y_value) in zip(_X_train.items(), _y_train.items()):
                train_X[X_key].extend(np.array(X_value))
                train_y[y_key].extend(np.array(y_value))

    # this is for intrument separation
    print('Saving...')
    if not os.path.exists(output_path):
        os.makedirs(output_path)
    for (X_key, X_value), (y_key, y_value) in tqdm(zip(train_X.items(), train_y.items())):
        if X_key == y_key:
            np.savez_compressed('{}/{}.npz'.format(output_path, X_key), np.array(X_value), np.array(y_value))

    print('Done!')

if __name__ == '__main__':
    main()