#!/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

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

def to_samples(multitrack, midi_res=settings.midi_resolution, how='by_group'):

    #how = 'by_group', 'by_instrument', 'merged',

    # TODO: add transpositions of every sample to every possible key transposition
    # np.roll(sample, pitch_interval, axis=1) for transposition
    # np.roll(sample, time_steps, axis=0) for time shifting

    samples_by_instrument = defaultdict( lambda : [] )

    for track in multitrack.tracks:

        key = settings.midi_group[track.program + 1] if not track.is_drum else 'Drums'

        # this makes pack of samples of N x 96 x 128 shape
        number_of_beats = floor(track.pianoroll.shape[0] / midi_res)
        track_pianoroll = track.pianoroll[: number_of_beats * midi_res]
        track_beats = track_pianoroll.reshape(number_of_beats, midi_res, 128)

        # save collected pack of data to dictionary with samples packs for groups of instruments
        for sample in track_beats:
            if sample.sum() != 0:
                samples_by_instrument[key].append(sample)

        # TODO: add posibility of choosing between saving samples to groups of instrument, or to every instrument separatly or with no differance
        # TODO: add option, for looking only for one instrument/group
        # TODO: add option for colecting, more than one beat per sample (min 4)

    return samples_by_instrument

def to_midi(samples, output_path=settings.generated_midi_path, program=0, tempo=120, is_drum=False, beat_resolution=settings.beat_resolution):
    tracks = [roll.Track(samples, program=program, is_drum=is_drum)]
    return_midi = roll.Multitrack(tracks=tracks, tempo=tempo, downbeat=[0, 96, 192, 288], beat_resolution=beat_resolution)
    roll.write(return_midi, output_path)
    return return_midi

# TODO: Make optial function to erase information of note lenth - ??
def ignore_note_lenght():
    pass

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

    samples_pack_by_instrument = defaultdict( lambda : list() )

    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)
            #load midi ro pypianoroll - Multirack
            try:
                multitrack = roll.parse(midi_file_path)
            except:
                # IDEA: Log errors, and save to file?
                continue

            for key, value in to_samples(multitrack).items():
                samples_pack_by_instrument[key].extend(value)

    # this is for intrument separation
    print('Saving...')
    if not os.path.exists(output_path):
        os.makedirs(output_path)
    for key, value in tqdm(samples_pack_by_instrument.items()):
        np.savez_compressed('{}/{}.npz'.format(output_path, key), np.array(value))

    print('Done!')

if __name__ == '__main__':
    main()