<reponame>robgon-art/pianonet
import numpy as np
from pypianoroll import Track, Multitrack

from pianonet.core.midi_tools import play_midi_from_file
from pianonet.core.custom_multitrack import CustomMultitrack


class Pianoroll(object):
    """
    An wrapped array representing piano key states in time. The first axis represents the time step, each of which has
    an associated key state of 128 keys. This implementation represents keys as binary on or off states at each time
    step, so the array is a 2D set of booleans.

    Example pianoroll array state:

            t = 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, ...

            C   0  1  1  1  1  0  0  0  0  0
            B   0  0  0  0  0  1  1  1  1  1
            A#  1  0  0  0  0  0  0  0  0  0
            A   0  1  1  1  1  0  0  0  0  0
            G#  0  0  0  0  0  0  0  0  0  0

    General MIDI notes: Note index 60 is middle C, program = 0 is piano instrument, the max velocity allowed is 128.
                        A midi file can store up to 16 different tracks.
    """

    def __init__(self, initializer, use_custom_multitrack=False):
        """
        initializer: A string that is a path to a midi file or an array of shape (time_steps, 128).
        """

        if isinstance(initializer, str):
            midi_file_path = initializer
            self.load_from_midi_file(midi_file_path, use_custom_multitrack)
        else:
            np_array = initializer
            self.array = np.copy(np_array)

    def load_from_midi_file(self, midi_file_path, use_custom_multitrack):
        """
        midi_file_path: String that is path to a midi file to load. This midi file is assumed to have a beat resolution
                        of 24.

        A merged and binarized numpy array (time_steps, 128) in shape is loaded into self.array.
        """

        if use_custom_multitrack:
            multitrack = CustomMultitrack(filename=midi_file_path)
        else:
            multitrack = Multitrack(filename=midi_file_path)

        multitrack.check_validity()

        if multitrack.beat_resolution != 24:
            raise Exception("Beat resolution of all midi files should be 24. Encountered beat resolution is " + str(
                multitrack.beat_resolution))

        multitrack.merge_tracks(track_indices=[i for i in range(len(multitrack.tracks))], mode='max', program=0,
                                remove_merged=True)

        pianoroll_array = (multitrack.tracks[0].pianoroll > 0.0).astype('bool')

        if (pianoroll_array.shape[0] == 0) or (pianoroll_array.shape[1] != 128):
            raise Exception(
                "Shape of pianoroll array should be (timesteps, 128), where timesteps > 0. Encountered shape is " + str(
                    pianoroll_array.shape))

        self.array = pianoroll_array

    def get_multitrack(self):
        """
        Returns the pypianoroll Multitrack representation of the midi created from the pianoroll array. This track has
        one piano track with an assumed tempo of 120 and a beat resolution of 24.
        """

        track = Track(pianoroll=self.array, program=0, is_drum=False)

        return Multitrack(tracks=[track], tempo=120, downbeat=None, beat_resolution=24)

    def save_to_midi_file(self, file_path):
        """
        Saves the Pianoroll instance to midi file on disc.

        file_path: Where to save the file.
        """

        self.get_multitrack().write(filename=file_path)

    def play(self):
        """
        Play the notes represented in self.array as a midi audio output.
        """

        play_midi_from_file(multitrack=self.get_multitrack())

    def stretch(self, stretch_fraction):
        """
        stretch_fraction: float >= 0.0 indicating how much to stretch the pianoroll

        stratch_fraction = 1.0 returns the original pianoroll
        stretch_fraction < 1.0 shortens (speeds up) the pianoroll
        stretch_fraction > 1.0 lengthens (slows down) the pianoroll

        This method modifies self.array in place. This method works as well as possible when stretch_fraction >= 0.5 as
        measured by similarity in the overlapping area of played notes. It could be improved for stretch_fraction < 0.5
        by taking the weighted most common state of the skipped notes, which could be > 2.0 in this case.
        """

        if stretch_fraction == 1.0:
            return

        time_steps_in_original_array = self.array.shape[0]
        time_steps_in_stretched_array = round(time_steps_in_original_array * stretch_fraction)

        if time_steps_in_stretched_array == 0:
            raise Exception("Cannot have zero timesteps in stretched pianoroll.")

        stretch_time_fractions_array = np.arange(0.0, 1.0, (1.0 / time_steps_in_stretched_array))

        original_array_indices = np.round(stretch_time_fractions_array * time_steps_in_original_array).astype('int')

        original_array_indices = np.clip(original_array_indices, a_min=None, a_max=(time_steps_in_original_array - 1))

        self.array = self.array[original_array_indices, :]

    def get_stretched(self, stretch_fraction):
        """
        stretch_fraction: float >= 0.0 indicating how much to stretch the pianoroll

        Same as stretch method, but returns a copy of the result.
        """

        new_pianoroll = self.get_copy()

        new_pianoroll.stretch(stretch_fraction=stretch_fraction)

        return new_pianoroll

    def add_zero_padding(self, left_padding_timesteps=0, right_padding_timesteps=0):
        """
        left_padding_timesteps: How many empty time steps to add to the left side of the pianoroll array
        right_padding_timesteps: How many empty time steps to add to the right side of the pianoroll array
        """

        total_timesteps = left_padding_timesteps + right_padding_timesteps + self.array.shape[0]

        padded_array = np.zeros((total_timesteps, self.array.shape[1]), dtype='bool')

        padded_array[left_padding_timesteps:self.array.shape[0] + left_padding_timesteps, :] = self.array

        self.array = padded_array

    def trim_silence_off_ends(self):
        """
        Crop start and end of self.array where there are only zero values.
        """

        note_counts_at_each_timestep_array = np.sum(self.array, axis=1)

        non_zero_indices = np.argwhere(note_counts_at_each_timestep_array)

        first_non_zero_index = non_zero_indices[0][0]
        last_non_zero_index = non_zero_indices[-1][0]

        if last_non_zero_index <= first_non_zero_index:
            raise Exception("The entire pianoroll is silence. Trimming silence would result in an empty array.")

        self.array = self.array[first_non_zero_index:last_non_zero_index + 1]

    def get_num_timesteps(self):
        """
        Returns how many timesteps are in the pianoroll as an integer.
        """

        return self.array.shape[0]

    def get_copy(self):
        """
        Returns copy of this pianoroll instance.
        """

        return Pianoroll(self.array.copy())

    def __getitem__(self, val):
        """
        val: A slice denoting what timesteps of the pianoroll to keep.

        A new pianoroll instance is returned with the requested slice.

        Example: p[10:20] returns a new pianoroll instance with the key state sets between 10 inclusive and 20 exclusive.
        """

        if isinstance(val, slice):
            return Pianoroll(self.array[val])

import numpy as np
from pypianoroll import Track, Multitrack

from pianonet.core.midi_tools import play_midi_from_file
from pianonet.core.custom_multitrack import CustomMultitrack


class Pianoroll(object):
    """
    An wrapped array representing piano key states in time. The first axis represents the time step, each of which has
    an associated key state of 128 keys. This implementation represents keys as binary on or off states at each time
    step, so the array is a 2D set of booleans.

    Example pianoroll array state:

            t = 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, ...

            C   0  1  1  1  1  0  0  0  0  0
            B   0  0  0  0  0  1  1  1  1  1
            A#  1  0  0  0  0  0  0  0  0  0
            A   0  1  1  1  1  0  0  0  0  0
            G#  0  0  0  0  0  0  0  0  0  0

    General MIDI notes: Note index 60 is middle C, program = 0 is piano instrument, the max velocity allowed is 128.
                        A midi file can store up to 16 different tracks.
    """

    def __init__(self, initializer, use_custom_multitrack=False):
        """
        initializer: A string that is a path to a midi file or an array of shape (time_steps, 128).
        """

        if isinstance(initializer, str):
            midi_file_path = initializer
            self.load_from_midi_file(midi_file_path, use_custom_multitrack)
        else:
            np_array = initializer
            self.array = np.copy(np_array)

    def load_from_midi_file(self, midi_file_path, use_custom_multitrack):
        """
        midi_file_path: String that is path to a midi file to load. This midi file is assumed to have a beat resolution
                        of 24.

        A merged and binarized numpy array (time_steps, 128) in shape is loaded into self.array.
        """

        if use_custom_multitrack:
            multitrack = CustomMultitrack(filename=midi_file_path)
        else:
            multitrack = Multitrack(filename=midi_file_path)

        multitrack.check_validity()

        if multitrack.beat_resolution != 24:
            raise Exception("Beat resolution of all midi files should be 24. Encountered beat resolution is " + str(
                multitrack.beat_resolution))

        multitrack.merge_tracks(track_indices=[i for i in range(len(multitrack.tracks))], mode='max', program=0,
                                remove_merged=True)

        pianoroll_array = (multitrack.tracks[0].pianoroll > 0.0).astype('bool')

        if (pianoroll_array.shape[0] == 0) or (pianoroll_array.shape[1] != 128):
            raise Exception(
                "Shape of pianoroll array should be (timesteps, 128), where timesteps > 0. Encountered shape is " + str(
                    pianoroll_array.shape))

        self.array = pianoroll_array

    def get_multitrack(self):
        """
        Returns the pypianoroll Multitrack representation of the midi created from the pianoroll array. This track has
        one piano track with an assumed tempo of 120 and a beat resolution of 24.
        """

        track = Track(pianoroll=self.array, program=0, is_drum=False)

        return Multitrack(tracks=[track], tempo=120, downbeat=None, beat_resolution=24)

    def save_to_midi_file(self, file_path):
        """
        Saves the Pianoroll instance to midi file on disc.

        file_path: Where to save the file.
        """

        self.get_multitrack().write(filename=file_path)

    def play(self):
        """
        Play the notes represented in self.array as a midi audio output.
        """

        play_midi_from_file(multitrack=self.get_multitrack())

    def stretch(self, stretch_fraction):
        """
        stretch_fraction: float >= 0.0 indicating how much to stretch the pianoroll

        stratch_fraction = 1.0 returns the original pianoroll
        stretch_fraction < 1.0 shortens (speeds up) the pianoroll
        stretch_fraction > 1.0 lengthens (slows down) the pianoroll

        This method modifies self.array in place. This method works as well as possible when stretch_fraction >= 0.5 as
        measured by similarity in the overlapping area of played notes. It could be improved for stretch_fraction < 0.5
        by taking the weighted most common state of the skipped notes, which could be > 2.0 in this case.
        """

        if stretch_fraction == 1.0:
            return

        time_steps_in_original_array = self.array.shape[0]
        time_steps_in_stretched_array = round(time_steps_in_original_array * stretch_fraction)

        if time_steps_in_stretched_array == 0:
            raise Exception("Cannot have zero timesteps in stretched pianoroll.")

        stretch_time_fractions_array = np.arange(0.0, 1.0, (1.0 / time_steps_in_stretched_array))

        original_array_indices = np.round(stretch_time_fractions_array * time_steps_in_original_array).astype('int')

        original_array_indices = np.clip(original_array_indices, a_min=None, a_max=(time_steps_in_original_array - 1))

        self.array = self.array[original_array_indices, :]

    def get_stretched(self, stretch_fraction):
        """
        stretch_fraction: float >= 0.0 indicating how much to stretch the pianoroll

        Same as stretch method, but returns a copy of the result.
        """

        new_pianoroll = self.get_copy()

        new_pianoroll.stretch(stretch_fraction=stretch_fraction)

        return new_pianoroll

    def add_zero_padding(self, left_padding_timesteps=0, right_padding_timesteps=0):
        """
        left_padding_timesteps: How many empty time steps to add to the left side of the pianoroll array
        right_padding_timesteps: How many empty time steps to add to the right side of the pianoroll array
        """

        total_timesteps = left_padding_timesteps + right_padding_timesteps + self.array.shape[0]

        padded_array = np.zeros((total_timesteps, self.array.shape[1]), dtype='bool')

        padded_array[left_padding_timesteps:self.array.shape[0] + left_padding_timesteps, :] = self.array

        self.array = padded_array

    def trim_silence_off_ends(self):
        """
        Crop start and end of self.array where there are only zero values.
        """

        note_counts_at_each_timestep_array = np.sum(self.array, axis=1)

        non_zero_indices = np.argwhere(note_counts_at_each_timestep_array)

        first_non_zero_index = non_zero_indices[0][0]
        last_non_zero_index = non_zero_indices[-1][0]

        if last_non_zero_index <= first_non_zero_index:
            raise Exception("The entire pianoroll is silence. Trimming silence would result in an empty array.")

        self.array = self.array[first_non_zero_index:last_non_zero_index + 1]

    def get_num_timesteps(self):
        """
        Returns how many timesteps are in the pianoroll as an integer.
        """

        return self.array.shape[0]

    def get_copy(self):
        """
        Returns copy of this pianoroll instance.
        """

        return Pianoroll(self.array.copy())

    def __getitem__(self, val):
        """
        val: A slice denoting what timesteps of the pianoroll to keep.

        A new pianoroll instance is returned with the requested slice.

        Example: p[10:20] returns a new pianoroll instance with the key state sets between 10 inclusive and 20 exclusive.
        """

        if isinstance(val, slice):
            return Pianoroll(self.array[val])
An wrapped array representing piano key states in time. The first axis represents the time step, each of which has     an associated key state of 128 keys. This implementation represents keys as binary on or off states at each time     step, so the array is a 2D set of booleans.      Example pianoroll array state:              t = 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, ...              C   0  1  1  1  1  0  0  0  0  0             B   0  0  0  0  0  1  1  1  1  1             A#  1  0  0  0  0  0  0  0  0  0             A   0  1  1  1  1  0  0  0  0  0             G#  0  0  0  0  0  0  0  0  0  0      General MIDI notes: Note index 60 is middle C, program = 0 is piano instrument, the max velocity allowed is 128.                         A midi file can store up to 16 different tracks. initializer: A string that is a path to a midi file or an array of shape (time_steps, 128). midi_file_path: String that is path to a midi file to load. This midi file is assumed to have a beat resolution                         of 24.          A merged and binarized numpy array (time_steps, 128) in shape is loaded into self.array. Returns the pypianoroll Multitrack representation of the midi created from the pianoroll array. This track has         one piano track with an assumed tempo of 120 and a beat resolution of 24. Saves the Pianoroll instance to midi file on disc.          file_path: Where to save the file. Play the notes represented in self.array as a midi audio output. stretch_fraction: float >= 0.0 indicating how much to stretch the pianoroll          stratch_fraction = 1.0 returns the original pianoroll         stretch_fraction < 1.0 shortens (speeds up) the pianoroll         stretch_fraction > 1.0 lengthens (slows down) the pianoroll          This method modifies self.array in place. This method works as well as possible when stretch_fraction >= 0.5 as         measured by similarity in the overlapping area of played notes. It could be improved for stretch_fraction < 0.5         by taking the weighted most common state of the skipped notes, which could be > 2.0 in this case. stretch_fraction: float >= 0.0 indicating how much to stretch the pianoroll          Same as stretch method, but returns a copy of the result. left_padding_timesteps: How many empty time steps to add to the left side of the pianoroll array         right_padding_timesteps: How many empty time steps to add to the right side of the pianoroll array Crop start and end of self.array where there are only zero values. Returns how many timesteps are in the pianoroll as an integer. Returns copy of this pianoroll instance. val: A slice denoting what timesteps of the pianoroll to keep.          A new pianoroll instance is returned with the requested slice.          Example: p[10:20] returns a new pianoroll instance with the key state sets between 10 inclusive and 20 exclusive.