#!/usr/bin/env python3

import settings
import pypianoroll as roll
import matplotlib.pyplot as plt
import numpy as np
import os
from tqdm import tqdm
from math import floor
import sys
import pickle
from tqdm import tqdm
from tqdm import trange
from collections import defaultdict
import bz2
import pickle

def to_samples(midi_file_path, midi_res=settings.midi_resolution):

    # 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

    fill_empty_array = lambda : np.empty((0, 96, 128))
    samples_by_instrument = defaultdict(fill_empty_array)
    all_beats = np.empty((0, 96, 128))

    for track in roll.Multitrack(midi_file_path).tracks:
        if not track.is_drum:
            key = track.program + 1

            # 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 every instrument
            samples_by_instrument[track.program + 1] = np.concatenate([track_beats, samples_by_instrument[ track.program + 1]], axis=0)
        else:
            # TODO: add code for drums samples
            pass
    return samples_by_instrument

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

# TODO: this function is running too slow.
def delete_empty_samples(sample_pack):
    non_empty_arrays = []
    for sample in sample_pack:
        if sample.sum() != 0:
            non_empty_arrays.append(sample)

    return np.array(non_empty_arrays)


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

    from collections import defaultdict
    fill_empty_array = lambda : np.empty((0, 96, 128))
    samples_pack_by_instrument = defaultdict(fill_empty_array)

    sample_pack = np.empty((0,settings.midi_resolution,128))

    for midi_file in tqdm(os.listdir(settings.midi_dir)):
        midi_file_path = '{}/{}'.format(settings.midi_dir, midi_file)
        midi_samples = to_samples(midi_file_path)
        if midi_samples is None:
            continue

        # this is for intrument separation
        for key, value in midi_samples.items():
            value = delete_empty_samples(value)
            samples_pack_by_instrument[key] = np.concatenate((samples_pack_by_instrument[key], value), axis=0)

    # save as compressed pickle (sample-dictionary)
    # sfile = bz2.BZ2File('data/samples.pickle', 'w')
    # pickle.dump(dict(samples_pack_by_instrument), sfile)

    # this is for intrument separation
    print('Saving...')
    for key, value in tqdm(samples_pack_by_instrument.items()):
        np.savez_compressed('data/samples/{}.npz'.format(settings.midi_program[key]), value)

        # Give a preview of what samples looks like
        fig, axes = plt.subplots(nrows=10, ncols=10, figsize=(20, 20))
        for idx, ax in enumerate(axes.ravel()):
            n = np.random.randint(0, value.shape[0])
            sample = value[n]
            ax.imshow(sample, cmap = plt.get_cmap('gray'))
        plt.savefig('data/samples/{}.png'.format(settings.midi_program[key]))

    print('Done!')

if __name__ == '__main__':
    main()