from music21 import stream

def generate():
    # loads the notes used to train the model from the notes file
    with open('./notes', 'rb') as filepath:
        notes = pickle.load(filepath)

    n_vocab = len(set(notes))
    pitchnames = sorted(set(item for item in notes))
    
    network_input, normalized_input = prepare_sequences(notes, n_vocab)
    # print(normalized_input.shape)

    model = create_network(normalized_input, n_vocab)
    # uses generate_notes to get sequences based on prediction output
    prediction_output = generate_notes(model, network_input, pitchnames, n_vocab)
    # creates the midi file from these sequences
    create_midi(prediction_output)

def prepare_sequences(notes, n_vocab):
    pitchnames = sorted(set(item for item in notes))
    note_to_int = dict((note, number) for number, note in enumerate(pitchnames))

    # the length of each sequence can be altered to improve results
    sequence_length = 35

    network_input = []
    output = []
    for i in range(0, len(notes) - sequence_length, 1):
        sequence_in = notes[i:i + sequence_length]
        sequence_out = notes[i + sequence_length]
        network_input.append([note_to_int[char] for char in sequence_in])
        output.append(note_to_int[sequence_out])

    n_patterns = len(network_input)

    # reshaping and normalizing the input
    normalized_input = numpy.reshape(network_input, (n_patterns, sequence_length, 1)) / float(n_vocab)

    return (network_input, normalized_input)

def create_network(network_input, n_vocab):
    
    model = tf.keras.Sequential([                                 
                                  tf.keras.layers.LSTM(256, input_shape = (network_input.shape[1], network_input.shape[2]), return_sequences=True, recurrent_dropout=0.3),
                                  tf.keras.layers.LSTM(256, return_sequences=True, recurrent_dropout=0.3),
                                  tf.keras.layers.LSTM(256),
                                  tf.keras.layers.BatchNormalization(),
                                  tf.keras.layers.Dropout(0.3),
                                  tf.keras.layers.Dense(256),
                                  tf.keras.layers.ReLU(),
                                  tf.keras.layers.BatchNormalization(),
                                  tf.keras.layers.Dropout(0.3),
                                  tf.keras.layers.Dense(n_vocab, activation = 'softmax')                           
  ]
  )
  
    model.load_weights('./weights.hdf5')
    return model

# generates notes based on a random sequence of notes
def generate_notes(model, network_input, pitchnames, n_vocab):
    # randomly picks a starting point for the prediction
    start = numpy.random.randint(0, len(network_input)-1)

    # maps integers to pitchnames of notes
    int_to_note = dict((number, note) for number, note in enumerate(pitchnames))

    # sets pattern to be the arbitrarily chosen sequence
    pattern = network_input[start]
    # stores generated sequences
    prediction_output = []

    # generates 100 notes based on the sequence in an auto-regressive mode
    for note_index in range(100):

        # the prediction input is the sequence reformatted for input and normalized
        prediction_input = numpy.reshape(pattern, (1, len(pattern), 1)) / float(n_vocab)

        # predictions stores a vector with the probabilities that the next note will be possible note
        prediction = model.predict(prediction_input, verbose=0)
        # stores the index of the most likely next note
        index = numpy.argmax(prediction)
        # stores note
        result = int_to_note[index]
        # adds the note to prediction_output
        prediction_output.append(result)

        pattern.append(index)
        pattern = pattern[1:len(pattern)]

    return prediction_output

# uses the prediction output to create a midi file
def create_midi(prediction_output):
    # the offset between notes
    offset = 0
    # stores the notes that will be in the midi
    output_notes = []

    # create note and chord objects based on the values generated by the model
    for pattern in prediction_output:
        # if pattern is a chord, separate it into notes, create note objects for each
        # of the notes with their information, offset the chord by the value stored
        # in offset, and add the chord to output_notes as a chord object
        if ('.' in pattern) or pattern.isdigit():
            notes_in_chord = pattern.split('.')
            notes = []
            for current_note in notes_in_chord:
                new_note = note.Note(int(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)
        # if the pattern is a note, create a note object with its information,
        # offset it by the value stored in offset, and add it to output_notes
        else:
            new_note = note.Note(pattern)
            new_note.offset = offset
            new_note.storedInstrument = instrument.Piano()
            output_notes.append(new_note)

        # increase the offset to separate notes (otherwise they stack)
        offset += 0.5

    # converts the notes to a music21 stream object
    midi_stream = stream.Stream(output_notes)

    print (len((output_notes)))

    # writes the stream object to a file
    import random
    file_number = random.randint(0, 1000)
    PATH = './output_files/output' + str(file_number) + '.midi'
    midi_stream.write('midi', fp=PATH) 

if __name__ == '__main__':
    generate()