#!python3
#!/usr/bin/env python3
''' This module generates a sample, and create a midi file.

Usage:
    >>> ./generate.py [trained_model_path] [output_path]

'''
import settings
import sys
import random
import pickle

import numpy as np
import tensorflow as tf
import pypianoroll as roll
import matplotlib.pyplot as plt
from tqdm import trange, tqdm
from music21 import converter, instrument, note, chord, stream
from keras.layers import Input, Dense, Conv2D
from keras.models import Model
from keras.layers import Input, Dense, Conv2D, Flatten, LSTM, Dropout, TimeDistributed, RepeatVector
from keras.models import Model, Sequential


def choose_by_prob(list_of_probs):
    ''' This functions a list of values and assumed
    that if the value is bigger it should by returned often
    
    It was crated to give more options to choose than argmax function,
    thus is more than one way that you can develop a melody.
    
    Returns a index of choosen value from given list.
    '''
    sum_prob = np.array(list_of_probs).sum()
    prob_normalized = [x/sum_prob for x in list_of_probs]
    cumsum = np.array(prob_normalized).cumsum()
    prob_cum = cumsum.tolist()
    random_x = random.random()
    for i, x in enumerate(prob_cum):
        if random_x < x:
            return i

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

# load model and dictionary that can translate back index_numbers to notes
# this dictionary is generated with model
print('Loading... {}'.format(trained_model_path))
model = pickle.load(open(trained_model_path, 'rb'))
int_to_note, n_vocab, seq_len = pickle.load(open('{}_dict'.format(trained_model_path), 'rb'))

seed = [random.randint(0,n_vocab) for x in range(seq_len)] 
music = []

print('Generating...')
for i in trange(124):
    predicted_vector = model.predict(np.array(seed).reshape(1,seq_len,1))
    # using best fitted note
    # predicted_index = np.argmax(predicted_vector)
    # using propability distribution for choosing note
    # to prevent looping
    predicted_index = choose_by_prob(predicted_vector)
    music.append(int_to_note[predicted_index])
    seed.append(predicted_index)
    seed = seed[1:1+seq_len]


print('Saving...')
offset = 0
output_notes = []
for _event in tqdm(music):
    event, note_len = _event.split(';')
    if (' ' in event) or event.isdigit():
        notes_in_chord = event.split(' ')
        notes = []
        for current_note in notes_in_chord:
            new_note = note.Note(current_note)
            new_note.storedInstrument = instrument.Piano()
            notes.append(new_note)
        new_chord = chord.Chord(notes)
        new_chord.offset = offset
        output_notes.append(new_chord)
    else:
        new_note = note.Note(event)
        new_note.offset = offset
        new_note.storedInstrument = instrument.Piano()
        output_notes.append(new_note)

    offset += float(note_len)

midi_stream = stream.Stream(output_notes)

midi_stream.write('midi', fp='{}.mid'.format(output_path))

print('Done!')