"""Utilities for manipulating multitrack pianorolls."""
from __future__ import absolute_import, division, print_function

from copy import deepcopy

import numpy as np

from pypianoroll.multitrack import Multitrack
from pypianoroll.track import Track


def _check_supported(obj):
    """
    Raise TypeError if the object is not a :class:`pypianoroll.Multitrack` or
    :class:`pypianoroll.Track` object.

    """
    if not isinstance(obj, (Multitrack, Track)):
        raise TypeError(
            "Support only `pypianoroll.Multitrack` and `pypianoroll.Track` objects"
        )


def check_pianoroll(arr):
    """
    Check if the array is a standard pianoroll matrix. Return True if it is.
    Otherwise, return False. Raise TypeError if the input object is not a numpy
    array.

    """
    if not isinstance(arr, np.ndarray):
        raise TypeError("`arr` must be of np.ndarray type")
    if not (np.issubdtype(arr.dtype, np.bool_) or np.issubdtype(arr.dtype, np.number)):
        return False
    if arr.ndim != 2:
        return False
    if arr.shape[1] != 128:
        return False
    return True


def assign_constant(obj, value):
    """
    Assign a constant value to the nonzeros in the pianoroll(s). If a pianoroll
    is not binarized, its data type will be preserved. If a pianoroll is
    binarized, it will be casted to the type of `value`.

    Arguments
    ---------
    value : int or float
        The constant value to be assigned to the nonzeros of the
        pianoroll(s).

    """
    _check_supported(obj)
    obj.assign_constant(value)


def binarize(obj, threshold=0):
    """
    Return a copy of the object with binarized pianoroll(s).

    Parameters
    ----------
    threshold : int or float
        Threshold to binarize the pianoroll(s). Default to zero.

    """
    _check_supported(obj)
    copied = deepcopy(obj)
    copied.binarize(threshold)
    return copied


def clip(obj, lower=0, upper=127):
    """
    Return a copy of the object with pianoroll(s) clipped by a lower bound and
    an upper bound specified by `lower` and `upper`, respectively.

    Parameters
    ----------
    lower : int or float
        The lower bound to clip the pianoroll. Default to 0.
    upper : int or float
        The upper bound to clip the pianoroll. Default to 127.

    """
    _check_supported(obj)
    copied = deepcopy(obj)
    copied.clip(lower, upper)
    return copied


def copy(obj):
    """Return a copy of the object."""
    _check_supported(obj)
    copied = deepcopy(obj)
    return copied


def downsample(obj, factor):
    """
    Return a copy of the object with pianoroll(s) downsampled by the given
    factor.

    Parameters
    ----------
    factor : int
        The ratio between the original beat resolution and the desired beat
        resolution.

    """
    _check_supported(obj)
    copied = deepcopy(obj)
    copied.downsample(factor)
    return copied


def load(filename):
    """
    Return a :class:`pypianoroll.Multitrack` object loaded from a .npz file.

    Parameters
    ----------
    filename : str
        The filename of the .npz file.

    """
    return Multitrack(filename)


def pad(obj, pad_length):
    """
    Return a copy of the object with pianoroll padded with zeros at the end
    along the time axis.

    Parameters
    ----------
    pad_length : int
        The length to pad along the time axis with zeros.

    """
    _check_supported(obj)
    copied = deepcopy(obj)
    copied.pad(pad_length)
    return copied


def pad_to_multiple(obj, factor):
    """
    Return a copy of the object with its pianoroll padded with zeros at the end
    along the time axis with the minimal length that make the length of the
    resulting pianoroll a multiple of `factor`.

    Parameters
    ----------
    factor : int
        The value which the length of the resulting pianoroll will be a multiple
        of.

    """
    _check_supported(obj)
    copied = deepcopy(obj)
    copied.pad_to_multiple(factor)
    return copied


def pad_to_same(obj):
    """
    Return a copy of the object with shorter pianorolls padded with zeros at the
    end along the time axis to the length of the pianoroll with the maximal
    length.

    """
    if not isinstance(obj, Multitrack):
        raise TypeError("Support only `pypianoroll.Multitrack` class objects")
    copied = deepcopy(obj)
    copied.pad_to_same()
    return copied


def parse(filename, beat_resolution=24, name="unknown"):
    """
    Return a :class:`pypianoroll.Multitrack` object loaded from a MIDI (.mid,
    .midi, .MID, .MIDI) file.

    Parameters
    ----------
    filename : str
        The filename of the MIDI file.

    """
    return Multitrack(filename, beat_resolution=beat_resolution, name=name)


def plot(obj, **kwargs):
    """
    Plot the object. See :func:`pypianoroll.plot_multitrack` and
    :func:`pypianoroll.plot_track` for full documentation.

    """
    _check_supported(obj)
    return obj.plot(**kwargs)


def save(filepath, obj, compressed=True):
    """
    Save the object to a .npz file.

    Parameters
    ----------
    filepath : str
        The path to save the file.
    obj: `pypianoroll.Multitrack` objects
        The object to be saved.

    """
    if not isinstance(obj, Multitrack):
        raise TypeError("Support only `pypianoroll.Multitrack` class objects")
    obj.save(filepath, compressed)


def transpose(obj, semitone):
    """
    Return a copy of the object with pianoroll(s) transposed by `semitones`
    semitones.

    Parameters
    ----------
    semitone : int
        Number of semitones to transpose the pianoroll(s).

    """
    _check_supported(obj)
    copied = deepcopy(obj)
    copied.transpose(semitone)
    return copied


def trim_trailing_silence(obj):
    """
    Return a copy of the object with trimmed trailing silence of the
    pianoroll(s).

    """
    _check_supported(obj)
    copied = deepcopy(obj)
    length = copied.get_active_length()
    copied.pianoroll = copied.pianoroll[:length]
    return copied


def write(obj, filepath):
    """
    Write the object to a MIDI file.

    Parameters
    ----------
    filepath : str
        The path to write the MIDI file.

    """
    if not isinstance(obj, Multitrack):
        raise TypeError("Support only `pypianoroll.Multitrack` class objects")
    obj.write(filepath)
