#!/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)
            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):
    print('Deleting empty samples...')
    temp_sample_pack = sample_pack
    index_manipulator = 1
    for index, sample in enumerate(sample_pack):
        if sample.sum() == 0:
            temp_sample_pack = np.delete(temp_sample_pack, index-index_manipulator, axis=0)
            index_manipulator = index_manipulator + 1
    print('Deleted {} empty samples'.format(index_manipulator-1))
    return temp_sample_pack

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)):
        print(midi_file)
        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():
            samples_pack_by_instrument[key] = np.concatenate((samples_pack_by_instrument[key], value), axis=0)

    # TODO: Delete empty samples - optimize
    # sample_pack = delete_empty_samples(sample_pack)

    # 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()