#!/usr/bin/env python
# -*- coding: utf-8 -*-
'''Matching functions'''

import sys
import six
import numpy as np
import numba

from .exceptions import ParameterError
from .utils import valid_intervals

__all__ = ['match_intervals', 'match_events']

@numba.jit(nopython=True, cache=True)
def __jaccard(int_a, int_b):  # pragma: no cover
    '''Jaccard similarity between two intervals

    Parameters
    ----------
    int_a, int_b : np.ndarrays, shape=(2,)

    Returns
    -------
    Jaccard similarity between intervals
    '''
    ends = [int_a[1], int_b[1]]
    if ends[1] < ends[0]:
        ends.reverse()

    starts = [int_a[0], int_b[0]]
    if starts[1] < starts[0]:
        starts.reverse()

    intersection = ends[0] - starts[1]
    if intersection < 0:
        intersection = 0.

    union = ends[1] - starts[0]

    if union > 0:
        return intersection / union

    return 0.0

@numba.jit(nopython=True, cache=True)
def __match_interval_overlaps(query, intervals_to, candidates):  # pragma: no cover
    '''Find the best Jaccard match from query to candidates'''

    best_score = -1
    best_idx = -1
    for idx in candidates:
        score = __jaccard(query, intervals_to[idx])

        if score > best_score:
            best_score, best_idx = score, idx
    return best_idx

@numba.jit(nopython=True, cache=True)
def __match_intervals(intervals_from, intervals_to, strict=True):  # pragma: no cover
    '''Numba-accelerated interval matching algorithm.

    '''
    # sort index of the interval starts
    start_index = np.argsort(intervals_to[:, 0])

    # sort index of the interval ends
    end_index = np.argsort(intervals_to[:, 1])

    # and sorted values of starts
    start_sorted = intervals_to[start_index, 0]
    # and ends
    end_sorted = intervals_to[end_index, 1]

    search_ends = np.searchsorted(start_sorted, intervals_from[:, 1], side='right')
    search_starts = np.searchsorted(end_sorted, intervals_from[:, 0], side='left')

    output = np.empty(len(intervals_from), dtype=numba.uint32)
    for i in range(len(intervals_from)):
        query = intervals_from[i]

        # Find the intervals that start after our query ends
        after_query = search_ends[i]
        # And the intervals that end after our query begins
        before_query = search_starts[i]

        # Candidates for overlapping have to (end after we start) and (begin before we end)
        candidates = set(start_index[:after_query]) & set(end_index[before_query:])

        # Proceed as before
        if len(candidates) > 0:
            output[i] = __match_interval_overlaps(query, intervals_to, candidates)
        elif strict:
            # Numba only lets us use compile-time constants in exception messages
            raise ParameterError
        else:
            # Find the closest interval
            # (start_index[after_query] - query[1]) is the distance to the next interval
            # (query[0] - end_index[before_query])
            dist_before = np.inf
            dist_after = np.inf
            if search_starts[i] > 0:
                dist_before = query[0] - end_sorted[search_starts[i]-1]
            if search_ends[i] + 1 < len(intervals_to):
                dist_after = start_sorted[search_ends[i]+1] - query[1]
            if dist_before < dist_after:
                output[i] = end_index[search_starts[i]-1]
            else:
                output[i] = start_index[search_ends[i]+1]
    return output

def match_intervals(intervals_from, intervals_to, strict=True):
    '''Match one set of time intervals to another.

    This can be useful for tasks such as mapping beat timings
    to segments.

    Each element `[a, b]` of `intervals_from` is matched to the
    element `[c, d]` of `intervals_to` which maximizes the
    Jaccard similarity between the intervals:

        `max(0, |min(b, d) - max(a, c)|) / |max(d, b) - min(a, c)|`

    In `strict=True` mode, if there is no interval with positive
    intersection with `[a,b]`, an exception is thrown.

    In `strict=False` mode, any interval `[a, b]` that has no
    intersection with any element of `intervals_to` is instead
    matched to the interval `[c, d]` which minimizes

        `min(|b - c|, |a - d|)`

    that is, the disjoint interval `[c, d]` with a boundary closest
    to `[a, b]`.

    .. note:: An element of `intervals_to` may be matched to multiple
       entries of `intervals_from`.

    Parameters
    ----------
    intervals_from : np.ndarray [shape=(n, 2)]
        The time range for source intervals.
        The `i` th interval spans time `intervals_from[i, 0]`
        to `intervals_from[i, 1]`.
        `intervals_from[0, 0]` should be 0, `intervals_from[-1, 1]`
        should be the track duration.

    intervals_to : np.ndarray [shape=(m, 2)]
        Analogous to `intervals_from`.

    strict : bool
        If `True`, intervals can only match if they intersect.
        If `False`, disjoint intervals can match.

    Returns
    -------
    interval_mapping : np.ndarray [shape=(n,)]
        For each interval in `intervals_from`, the
        corresponding interval in `intervals_to`.

    See Also
    --------
    match_events

    Raises
    ------
    ParameterError
        If either array of input intervals is not the correct shape

        If `strict=True` and some element of `intervals_from` is disjoint from
        every element of `intervals_to`.

    Examples
    --------
    >>> ints_from = np.array([[3, 5], [1, 4], [4, 5]])
    >>> ints_to = np.array([[0, 2], [1, 3], [4, 5], [6, 7]])
    >>> librosa.util.match_intervals(ints_from, ints_to)
    array([2, 1, 2], dtype=uint32)
    >>> # [3, 5] => [4, 5]  (ints_to[2])
    >>> # [1, 4] => [1, 3]  (ints_to[1])
    >>> # [4, 5] => [4, 5]  (ints_to[2])

    The reverse matching of the above is not possible in `strict` mode
    because `[6, 7]` is disjoint from all intervals in `ints_from`.
    With `strict=False`, we get the following:
    >>> librosa.util.match_intervals(ints_to, ints_from, strict=False)
    array([1, 1, 2, 2], dtype=uint32)
    >>> # [0, 2] => [1, 4]  (ints_from[1])
    >>> # [1, 3] => [1, 4]  (ints_from[1])
    >>> # [4, 5] => [4, 5]  (ints_from[2])
    >>> # [6, 7] => [4, 5]  (ints_from[2])
    '''

    if len(intervals_from) == 0 or len(intervals_to) == 0:
        raise ParameterError('Attempting to match empty interval list')

    # Verify that the input intervals has correct shape and size
    valid_intervals(intervals_from)
    valid_intervals(intervals_to)

    try:
        return __match_intervals(intervals_from, intervals_to, strict=strict)
    except ParameterError:
        six.reraise(ParameterError,
                    ParameterError('Unable to match intervals with strict={}'.format(strict)),
                    sys.exc_info()[2])

def match_events(events_from, events_to, left=True, right=True):
    '''Match one set of events to another.

    This is useful for tasks such as matching beats to the nearest
    detected onsets, or frame-aligned events to the nearest zero-crossing.

    .. note:: A target event may be matched to multiple source events.

    Examples
    --------
    >>> # Sources are multiples of 7
    >>> s_from = np.arange(0, 100, 7)
    >>> s_from
    array([ 0,  7, 14, 21, 28, 35, 42, 49, 56, 63, 70, 77, 84, 91,
           98])
    >>> # Targets are multiples of 10
    >>> s_to = np.arange(0, 100, 10)
    >>> s_to
    array([ 0, 10, 20, 30, 40, 50, 60, 70, 80, 90])
    >>> # Find the matching
    >>> idx = librosa.util.match_events(s_from, s_to)
    >>> idx
    array([0, 1, 1, 2, 3, 3, 4, 5, 6, 6, 7, 8, 8, 9, 9])
    >>> # Print each source value to its matching target
    >>> zip(s_from, s_to[idx])
    [(0, 0), (7, 10), (14, 10), (21, 20), (28, 30), (35, 30),
     (42, 40), (49, 50), (56, 60), (63, 60), (70, 70), (77, 80),
     (84, 80), (91, 90), (98, 90)]

    Parameters
    ----------
    events_from : ndarray [shape=(n,)]
      Array of events (eg, times, sample or frame indices) to match from.

    events_to : ndarray [shape=(m,)]
      Array of events (eg, times, sample or frame indices) to
      match against.

    left : bool
    right : bool
        If `False`, then matched events cannot be to the left (or right)
        of source events.

    Returns
    -------
    event_mapping : np.ndarray [shape=(n,)]
        For each event in `events_from`, the corresponding event
        index in `events_to`.

        `event_mapping[i] == arg min |events_from[i] - events_to[:]|`

    See Also
    --------
    match_intervals

    Raises
    ------
    ParameterError
        If either array of input events is not the correct shape
    '''
    if len(events_from) == 0 or len(events_to) == 0:
        raise ParameterError('Attempting to match empty event list')

    # If we can't match left or right, then only strict equivalence
    # counts as a match.
    if not (left or right) and not np.all(np.in1d(events_from, events_to)):
        raise ParameterError('Cannot match events with left=right=False '
                             'and events_from is not contained '
                             'in events_to')

    # If we can't match to the left, then there should be at least one
    # target event greater-equal to every source event
    if (not left) and max(events_to) < max(events_from):
        raise ParameterError('Cannot match events with left=False '
                             'and max(events_to) < max(events_from)')

    # If we can't match to the right, then there should be at least one
    # target event less-equal to every source event
    if (not right) and min(events_to) > min(events_from):
        raise ParameterError('Cannot match events with right=False '
                             'and min(events_to) > min(events_from)')

    # array of matched items
    output = np.empty_like(events_from, dtype=np.int)

    return __match_events_helper(output, events_from, events_to, left, right)

@numba.jit(nopython=True, cache=True)
def __match_events_helper(output, events_from, events_to, left=True, right=True):  # pragma: no cover
    # mock dictionary for events
    from_idx = np.argsort(events_from)
    sorted_from = events_from[from_idx]

    to_idx = np.argsort(events_to)
    sorted_to = events_to[to_idx]

    # find the matching indices
    matching_indices = np.searchsorted(sorted_to, sorted_from)

    # iterate over indices in matching_indices
    for ind, middle_ind in enumerate(matching_indices):
        left_flag = False
        right_flag = False

        left_ind = -1
        right_ind = len(matching_indices)

        left_diff = 0
        right_diff = 0
        mid_diff = 0

        middle_ind = matching_indices[ind]
        sorted_from_num = sorted_from[ind]

        # Prevent oob from chosen index
        if middle_ind == len(sorted_to):
            middle_ind -= 1

        # Permitted to look to the left
        if left and middle_ind > 0:
            left_ind = middle_ind - 1
            left_flag = True

        # Permitted to look to right
        if right and middle_ind < len(sorted_to) - 1:
            right_ind = middle_ind + 1
            right_flag = True

        mid_diff = abs(sorted_to[middle_ind] - sorted_from_num)
        if left and left_flag:
            left_diff = abs(sorted_to[left_ind] - sorted_from_num)
        if right and right_flag:
            right_diff = abs(sorted_to[right_ind] - sorted_from_num)

        if left_flag and (not right and (sorted_to[middle_ind] > sorted_from_num) or
                          (not right_flag and left_diff < mid_diff) or
                          (left_diff < right_diff and left_diff < mid_diff)):
            output[ind] = to_idx[left_ind]

        # Check if right should be chosen
        elif right_flag and (right_diff < mid_diff):
            output[ind] = to_idx[right_ind]

        # Selected index wins
        else:
            output[ind] = to_idx[middle_ind]

    # Undo sorting
    solutions = np.empty_like(output)
    solutions[from_idx] = output

    return solutions


#!/usr/bin/env python
# -*- coding: utf-8 -*-
'''Harmonic calculations for frequency representations'''

import numpy as np
import scipy.interpolate
import scipy.signal
from ..util.exceptions import ParameterError

__all__ = ['salience', 'interp_harmonics']

def salience(S, freqs, h_range, weights=None, aggregate=None,
             filter_peaks=True, fill_value=np.nan,  kind='linear', axis=0):
    """Harmonic salience function.

    Parameters
    ----------
    S : np.ndarray [shape=(d, n)]
        input time frequency magnitude representation (stft, ifgram, etc).
        Must be real-valued and non-negative.
    freqs : np.ndarray, shape=(S.shape[axis])
        The frequency values corresponding to S's elements along the
        chosen axis.
    h_range : list-like, non-negative
        Harmonics to include in salience computation.  The first harmonic (1)
        corresponds to `S` itself. Values less than one (e.g., 1/2) correspond
        to sub-harmonics.
    weights : list-like
        The weight to apply to each harmonic in the summation. (default:
        uniform weights). Must be the same length as `harmonics`.
    aggregate : function
        aggregation function (default: `np.average`)
        If `aggregate=np.average`, then a weighted average is
        computed per-harmonic according to the specified weights.
        For all other aggregation functions, all harmonics
        are treated equally.
    filter_peaks : bool
        If true, returns harmonic summation only on frequencies of peak
        magnitude. Otherwise returns harmonic summation over the full spectrum.
        Defaults to True.
    fill_value : float
        The value to fill non-peaks in the output representation. (default:
        np.nan) Only used if `filter_peaks == True`.
    kind : str
        Interpolation type for harmonic estimation.
        See `scipy.interpolate.interp1d`.
    axis : int
        The axis along which to compute harmonics

    Returns
    -------
    S_sal : np.ndarray, shape=(len(h_range), [x.shape])
        `S_sal` will have the same shape as `S`, and measure
        the overal harmonic energy at each frequency.

    See Also
    --------
    interp_harmonics

    Examples
    --------
    >>> y, sr = librosa.load(librosa.util.example_audio_file(),
    ...                      duration=15, offset=30)
    >>> S = np.abs(librosa.stft(y))
    >>> freqs = librosa.core.fft_frequencies(sr)
    >>> harms = [1, 2, 3, 4]
    >>> weights = [1.0, 0.5, 0.33, 0.25]
    >>> S_sal = librosa.salience(S, freqs, harms, weights, fill_value=0)
    >>> print(S_sal.shape)
    (1025, 646)
    >>> import matplotlib.pyplot as plt
    >>> plt.figure()
    >>> librosa.display.specshow(librosa.amplitude_to_db(S_sal,
    ...                                                  ref=np.max),
    ...                          sr=sr, y_axis='log', x_axis='time')
    >>> plt.colorbar()
    >>> plt.title('Salience spectrogram')
    >>> plt.tight_layout()
    """
    if aggregate is None:
        aggregate = np.average

    if weights is None:
        weights = np.ones((len(h_range), ))
    else:
        weights = np.array(weights, dtype=float)

    S_harm = interp_harmonics(S, freqs, h_range, kind=kind, axis=axis)

    if aggregate is np.average:
        S_sal = aggregate(S_harm, axis=0, weights=weights)
    else:
        S_sal = aggregate(S_harm, axis=0)

    if filter_peaks:
        S_peaks = scipy.signal.argrelmax(S, axis=0)
        S_out = np.empty(S.shape)
        S_out.fill(fill_value)
        S_out[S_peaks[0], S_peaks[1]] = S_sal[S_peaks[0], S_peaks[1]]

        S_sal = S_out

    return S_sal

def interp_harmonics(x, freqs, h_range, kind='linear', fill_value=0, axis=0):
    '''Compute the energy at harmonics of time-frequency representation.

    Given a frequency-based energy representation such as a spectrogram
    or tempogram, this function computes the energy at the chosen harmonics
    of the frequency axis.  (See examples below.)
    The resulting harmonic array can then be used as input to a salience
    computation.

    Parameters
    ----------
    x : np.ndarray
        The input energy

    freqs : np.ndarray, shape=(X.shape[axis])
        The frequency values corresponding to X's elements along the
        chosen axis.

    h_range : list-like, non-negative
        Harmonics to compute.  The first harmonic (1) corresponds to `x`
        itself.
        Values less than one (e.g., 1/2) correspond to sub-harmonics.

    kind : str
        Interpolation type.  See `scipy.interpolate.interp1d`.

    fill_value : float
        The value to fill when extrapolating beyond the observed
        frequency range.

    axis : int
        The axis along which to compute harmonics

    Returns
    -------
    x_harm : np.ndarray, shape=(len(h_range), [x.shape])
        `x_harm[i]` will have the same shape as `x`, and measure
        the energy at the `h_range[i]` harmonic of each frequency.

    See Also
    --------
    scipy.interpolate.interp1d

    Examples
    --------
    Estimate the harmonics of a time-averaged tempogram

    >>> y, sr = librosa.load(librosa.util.example_audio_file(),
    ...                      duration=15, offset=30)
    >>> # Compute the time-varying tempogram and average over time
    >>> tempi = np.mean(librosa.feature.tempogram(y=y, sr=sr), axis=1)
    >>> # We'll measure the first five harmonics
    >>> h_range = [1, 2, 3, 4, 5]
    >>> f_tempo = librosa.tempo_frequencies(len(tempi), sr=sr)
    >>> # Build the harmonic tensor
    >>> t_harmonics = librosa.interp_harmonics(tempi, f_tempo, h_range)
    >>> print(t_harmonics.shape)
    (5, 384)

    >>> # And plot the results
    >>> import matplotlib.pyplot as plt
    >>> plt.figure()
    >>> librosa.display.specshow(t_harmonics, x_axis='tempo', sr=sr)
    >>> plt.yticks(0.5 + np.arange(len(h_range)),
    ...            ['{:.3g}'.format(_) for _ in h_range])
    >>> plt.ylabel('Harmonic')
    >>> plt.xlabel('Tempo (BPM)')
    >>> plt.tight_layout()

    We can also compute frequency harmonics for spectrograms.
    To calculate sub-harmonic energy, use values < 1.

    >>> h_range = [1./3, 1./2, 1, 2, 3, 4]
    >>> S = np.abs(librosa.stft(y))
    >>> fft_freqs = librosa.fft_frequencies(sr=sr)
    >>> S_harm = librosa.interp_harmonics(S, fft_freqs, h_range, axis=0)
    >>> print(S_harm.shape)
    (6, 1025, 646)

    >>> plt.figure()
    >>> for i, _sh in enumerate(S_harm, 1):
    ...     plt.subplot(3, 2, i)
    ...     librosa.display.specshow(librosa.amplitude_to_db(_sh,
    ...                                                      ref=S.max()),
    ...                              sr=sr, y_axis='log')
    ...     plt.title('h={:.3g}'.format(h_range[i-1]))
    ...     plt.yticks([])
    >>> plt.tight_layout()
    '''

    # X_out will be the same shape as X, plus a leading
    # axis that has length = len(h_range)
    out_shape = [len(h_range)]
    out_shape.extend(x.shape)

    x_out = np.zeros(out_shape, dtype=x.dtype)

    if freqs.ndim == 1 and len(freqs) == x.shape[axis]:
        harmonics_1d(x_out, x, freqs, h_range,
                     kind=kind, fill_value=fill_value,
                     axis=axis)

    elif freqs.ndim == 2 and freqs.shape == x.shape:
        harmonics_2d(x_out, x, freqs, h_range,
                     kind=kind, fill_value=fill_value,
                     axis=axis)
    else:
        raise ParameterError('freqs.shape={} does not match '
                             'input shape={}'.format(freqs.shape, x.shape))

    return x_out

def harmonics_1d(harmonic_out, x, freqs, h_range, kind='linear',
                 fill_value=0, axis=0):
    '''Populate a harmonic tensor from a time-frequency representation.

    Parameters
    ----------
    harmonic_out : np.ndarray, shape=(len(h_range), X.shape)
        The output array to store harmonics

    X : np.ndarray
        The input energy

    freqs : np.ndarray, shape=(x.shape[axis])
        The frequency values corresponding to x's elements along the
        chosen axis.

    h_range : list-like, non-negative
        Harmonics to compute.  The first harmonic (1) corresponds to `x`
        itself.
        Values less than one (e.g., 1/2) correspond to sub-harmonics.

    kind : str
        Interpolation type.  See `scipy.interpolate.interp1d`.

    fill_value : float
        The value to fill when extrapolating beyond the observed
        frequency range.

    axis : int
        The axis along which to compute harmonics

    See Also
    --------
    harmonics
    scipy.interpolate.interp1d

    Examples
    --------
    Estimate the harmonics of a time-averaged tempogram

    >>> y, sr = librosa.load(librosa.util.example_audio_file(),
    ...                      duration=15, offset=30)
    >>> # Compute the time-varying tempogram and average over time
    >>> tempi = np.mean(librosa.feature.tempogram(y=y, sr=sr), axis=1)
    >>> # We'll measure the first five harmonics
    >>> h_range = [1, 2, 3, 4, 5]
    >>> f_tempo = librosa.tempo_frequencies(len(tempi), sr=sr)
    >>> # Build the harmonic tensor
    >>> t_harmonics = librosa.interp_harmonics(tempi, f_tempo, h_range)
    >>> print(t_harmonics.shape)
    (5, 384)

    >>> # And plot the results
    >>> import matplotlib.pyplot as plt
    >>> plt.figure()
    >>> librosa.display.specshow(t_harmonics, x_axis='tempo', sr=sr)
    >>> plt.yticks(0.5 + np.arange(len(h_range)),
    ...            ['{:.3g}'.format(_) for _ in h_range])
    >>> plt.ylabel('Harmonic')
    >>> plt.xlabel('Tempo (BPM)')
    >>> plt.tight_layout()

    We can also compute frequency harmonics for spectrograms.
    To calculate subharmonic energy, use values < 1.

    >>> h_range = [1./3, 1./2, 1, 2, 3, 4]
    >>> S = np.abs(librosa.stft(y))
    >>> fft_freqs = librosa.fft_frequencies(sr=sr)
    >>> S_harm = librosa.interp_harmonics(S, fft_freqs, h_range, axis=0)
    >>> print(S_harm.shape)
    (6, 1025, 646)

    >>> plt.figure()
    >>> for i, _sh in enumerate(S_harm, 1):
    ...     plt.subplot(3,2,i)
    ...     librosa.display.specshow(librosa.amplitude_to_db(_sh,
    ...                                                      ref=S.max()),
    ...                              sr=sr, y_axis='log')
    ...     plt.title('h={:.3g}'.format(h_range[i-1]))
    ...     plt.yticks([])
    >>> plt.tight_layout()
    '''

    # Note: this only works for fixed-grid, 1d interpolation
    f_interp = scipy.interpolate.interp1d(freqs, x,
                                          kind=kind,
                                          axis=axis,
                                          copy=False,
                                          bounds_error=False,
                                          fill_value=fill_value)

    idx_out = [slice(None)] * harmonic_out.ndim

    # Compute the output index of the interpolated values
    interp_axis = 1 + (axis % x.ndim)

    # Iterate over the harmonics range
    for h_index, harmonic in enumerate(h_range):
        idx_out[0] = h_index

        # Iterate over frequencies
        for f_index, frequency in enumerate(freqs):
            # Offset the output axis by 1 to account for the harmonic index
            idx_out[interp_axis] = f_index

            # Estimate the harmonic energy at this frequency across time
            harmonic_out[tuple(idx_out)] = f_interp(harmonic * frequency)

def harmonics_2d(harmonic_out, x, freqs, h_range, kind='linear', fill_value=0,
                 axis=0):
    '''Populate a harmonic tensor from a time-frequency representation with
    time-varying frequencies.

    Parameters
    ----------
    harmonic_out : np.ndarray
        The output array to store harmonics

    x : np.ndarray
        The input energy

    freqs : np.ndarray, shape=x.shape
        The frequency values corresponding to each element of `x`

    h_range : list-like, non-negative
        Harmonics to compute.  The first harmonic (1) corresponds to `x`
        itself.  Values less than one (e.g., 1/2) correspond to
        sub-harmonics.

    kind : str
        Interpolation type.  See `scipy.interpolate.interp1d`.

    fill_value : float
        The value to fill when extrapolating beyond the observed
        frequency range.

    axis : int
        The axis along which to compute harmonics

    See Also
    --------
    harmonics
    harmonics_1d
    '''
    idx_in = [slice(None)] * x.ndim
    idx_freq = [slice(None)] * x.ndim
    idx_out = [slice(None)] * harmonic_out.ndim

    # This is the non-interpolation axis
    ni_axis = (1 + axis) % x.ndim

    # For each value in the non-interpolated axis, compute its harmonics
    for i in range(x.shape[ni_axis]):
        idx_in[ni_axis] = slice(i, i + 1)
        idx_freq[ni_axis] = i
        idx_out[1 + ni_axis] = idx_in[ni_axis]

        harmonics_1d(harmonic_out[tuple(idx_out)], x[tuple(idx_in)], freqs[tuple(idx_freq)],
                     h_range, kind=kind, fill_value=fill_value,
                     axis=axis)


#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""Core IO, DSP and utility functions."""

import os
import six

import soundfile as sf
import audioread
import numpy as np
import scipy.signal
import resampy

from numba import jit
from .fft import get_fftlib
from .time_frequency import frames_to_samples, time_to_samples
from .._cache import cache
from .. import util
from ..util.exceptions import ParameterError

__all__ = ['load', 'to_mono', 'resample', 'get_duration', 'autocorrelate',
           'lpc', 'zero_crossings', 'clicks', 'tone', 'chirp']

# Resampling bandwidths as percentage of Nyquist
BW_BEST = resampy.filters.get_filter('kaiser_best')[2]
BW_FASTEST = resampy.filters.get_filter('kaiser_fast')[2]

# -- CORE ROUTINES --#
# Load should never be cached, since we cannot verify that the contents of
# 'path' are unchanged across calls.
def load(path, sr=22050, mono=True, offset=0.0, duration=None,
         dtype=np.float32, res_type='kaiser_best'):
    """Load an audio file as a floating point time series.

    Audio will be automatically resampled to the given rate
    (default `sr=22050`).

    To preserve the native sampling rate of the file, use `sr=None`.

    Parameters
    ----------
    path : string, int, or file-like object
        path to the input file.

        Any codec supported by `soundfile` or `audioread` will work.

        If the codec is supported by `soundfile`, then `path` can also be
        an open file descriptor (int), or any object implementing Python's
        file interface.

        If the codec is not supported by `soundfile` (e.g., MP3), then only
        string file paths are supported.

    sr   : number > 0 [scalar]
        target sampling rate

        'None' uses the native sampling rate

    mono : bool
        convert signal to mono

    offset : float
        start reading after this time (in seconds)

    duration : float
        only load up to this much audio (in seconds)

    dtype : numeric type
        data type of `y`

    res_type : str
        resample type (see note)

        .. note::
            By default, this uses `resampy`'s high-quality mode ('kaiser_best').

            For alternative resampling modes, see `resample`

        .. note::
           `audioread` may truncate the precision of the audio data to 16 bits.

           See https://librosa.github.io/librosa/ioformats.html for alternate
           loading methods.

    Returns
    -------
    y    : np.ndarray [shape=(n,) or (2, n)]
        audio time series

    sr   : number > 0 [scalar]
        sampling rate of `y`

    Examples
    --------
    >>> # Load an ogg vorbis file
    >>> filename = librosa.util.example_audio_file()
    >>> y, sr = librosa.load(filename)
    >>> y
    array([ -4.756e-06,  -6.020e-06, ...,  -1.040e-06,   0.000e+00], dtype=float32)
    >>> sr
    22050

    >>> # Load a file and resample to 11 KHz
    >>> filename = librosa.util.example_audio_file()
    >>> y, sr = librosa.load(filename, sr=11025)
    >>> y
    array([ -2.077e-06,  -2.928e-06, ...,  -4.395e-06,   0.000e+00], dtype=float32)
    >>> sr
    11025

    >>> # Load 5 seconds of a file, starting 15 seconds in
    >>> filename = librosa.util.example_audio_file()
    >>> y, sr = librosa.load(filename, offset=15.0, duration=5.0)
    >>> y
    array([ 0.069,  0.1  , ..., -0.101,  0.   ], dtype=float32)
    >>> sr
    22050

    """

    try:
        with sf.SoundFile(path) as sf_desc:
            sr_native = sf_desc.samplerate
            if offset:
                # Seek to the start of the target read
                sf_desc.seek(int(offset * sr_native))
            if duration is not None:
                frame_duration = int(duration * sr_native)
            else:
                frame_duration = -1

            # Load the target number of frames, and transpose to match librosa form
            y = sf_desc.read(frames=frame_duration, dtype=dtype, always_2d=False).T

    except RuntimeError as exc:
        # If soundfile failed, fall back to the audioread loader
        y, sr_native = __audioread_load(path, offset, duration, dtype)

    # Final cleanup for dtype and contiguity
    if mono:
        y = to_mono(y)

    if sr is not None:
        y = resample(y, sr_native, sr, res_type=res_type)

    else:
        sr = sr_native

    return y, sr

def __audioread_load(path, offset, duration, dtype):
    '''Load an audio buffer using audioread.

    This loads one block at a time, and then concatenates the results.
    '''

    y = []
    with audioread.audio_open(path) as input_file:
        sr_native = input_file.samplerate
        n_channels = input_file.channels

        s_start = int(np.round(sr_native * offset)) * n_channels

        if duration is None:
            s_end = np.inf
        else:
            s_end = s_start + (int(np.round(sr_native * duration))
                               * n_channels)

        n = 0

        for frame in input_file:
            frame = util.buf_to_float(frame, dtype=dtype)
            n_prev = n
            n = n + len(frame)

            if n < s_start:
                # offset is after the current frame
                # keep reading
                continue

            if s_end < n_prev:
                # we're off the end.  stop reading
                break

            if s_end < n:
                # the end is in this frame.  crop.
                frame = frame[:s_end - n_prev]

            if n_prev <= s_start <= n:
                # beginning is in this frame
                frame = frame[(s_start - n_prev):]

            # tack on the current frame
            y.append(frame)

    if y:
        y = np.concatenate(y)
        if n_channels > 1:
            y = y.reshape((-1, n_channels)).T
    else:
        y = np.empty(0, dtype=dtype)

    return y, sr_native

@cache(level=20)
def to_mono(y):
    '''Force an audio signal down to mono.

    Parameters
    ----------
    y : np.ndarray [shape=(2,n) or shape=(n,)]
        audio time series, either stereo or mono

    Returns
    -------
    y_mono : np.ndarray [shape=(n,)]
        `y` as a monophonic time-series

    Notes
    -----
    This function caches at level 20.

    Examples
    --------
    >>> y, sr = librosa.load(librosa.util.example_audio_file(), mono=False)
    >>> y.shape
    (2, 1355168)
    >>> y_mono = librosa.to_mono(y)
    >>> y_mono.shape
    (1355168,)

    '''

    # Validate the buffer.  Stereo is ok here.
    util.valid_audio(y, mono=False)

    if y.ndim > 1:
        y = np.mean(y, axis=0)

    return y

@cache(level=20)
def resample(y, orig_sr, target_sr, res_type='kaiser_best', fix=True, scale=False, **kwargs):
    """Resample a time series from orig_sr to target_sr

    Parameters
    ----------
    y : np.ndarray [shape=(n,) or shape=(2, n)]
        audio time series.  Can be mono or stereo.

    orig_sr : number > 0 [scalar]
        original sampling rate of `y`

    target_sr : number > 0 [scalar]
        target sampling rate

    res_type : str
        resample type (see note)

        .. note::
            By default, this uses `resampy`'s high-quality mode ('kaiser_best').

            To use a faster method, set `res_type='kaiser_fast'`.

            To use `scipy.signal.resample`, set `res_type='fft'` or `res_type='scipy'`.

            To use `scipy.signal.resample_poly`, set `res_type='polyphase'`.

        .. note::
            When using `res_type='polyphase'`, only integer sampling rates are
            supported.

    fix : bool
        adjust the length of the resampled signal to be of size exactly
        `ceil(target_sr * len(y) / orig_sr)`

    scale : bool
        Scale the resampled signal so that `y` and `y_hat` have approximately
        equal total energy.

    kwargs : additional keyword arguments
        If `fix==True`, additional keyword arguments to pass to
        `librosa.util.fix_length`.

    Returns
    -------
    y_hat : np.ndarray [shape=(n * target_sr / orig_sr,)]
        `y` resampled from `orig_sr` to `target_sr`

    Raises
    ------
    ParameterError
        If `res_type='polyphase'` and `orig_sr` or `target_sr` are not both
        integer-valued.

    See Also
    --------
    librosa.util.fix_length
    scipy.signal.resample
    resampy.resample

    Notes
    -----
    This function caches at level 20.

    Examples
    --------
    Downsample from 22 KHz to 8 KHz

    >>> y, sr = librosa.load(librosa.util.example_audio_file(), sr=22050)
    >>> y_8k = librosa.resample(y, sr, 8000)
    >>> y.shape, y_8k.shape
    ((1355168,), (491671,))
    """

    # First, validate the audio buffer
    util.valid_audio(y, mono=False)

    if orig_sr == target_sr:
        return y

    ratio = float(target_sr) / orig_sr

    n_samples = int(np.ceil(y.shape[-1] * ratio))

    if res_type in ('scipy', 'fft'):
        y_hat = scipy.signal.resample(y, n_samples, axis=-1)
    elif res_type == 'polyphase':
        if int(orig_sr) != orig_sr or int(target_sr) != target_sr:
            raise ParameterError('polyphase resampling is only supported for integer-valued sampling rates.')

        # For polyphase resampling, we need up- and down-sampling ratios
        # We can get those from the greatest common divisor of the rates
        # as long as the rates are integrable
        orig_sr = int(orig_sr)
        target_sr = int(target_sr)
        gcd = np.gcd(orig_sr, target_sr)
        y_hat = scipy.signal.resample_poly(y, target_sr // gcd, orig_sr // gcd, axis=-1)
    else:
        y_hat = resampy.resample(y, orig_sr, target_sr, filter=res_type, axis=-1)

    if fix:
        y_hat = util.fix_length(y_hat, n_samples, **kwargs)

    if scale:
        y_hat /= np.sqrt(ratio)

    return np.ascontiguousarray(y_hat, dtype=y.dtype)

def get_duration(y=None, sr=22050, S=None, n_fft=2048, hop_length=512,
                 center=True, filename=None):
    """Compute the duration (in seconds) of an audio time series,
    feature matrix, or filename.

    Examples
    --------
    >>> # Load the example audio file
    >>> y, sr = librosa.load(librosa.util.example_audio_file())
    >>> librosa.get_duration(y=y, sr=sr)
    61.45886621315193

    >>> # Or directly from an audio file
    >>> librosa.get_duration(filename=librosa.util.example_audio_file())
    61.4

    >>> # Or compute duration from an STFT matrix
    >>> y, sr = librosa.load(librosa.util.example_audio_file())
    >>> S = librosa.stft(y)
    >>> librosa.get_duration(S=S, sr=sr)
    61.44

    >>> # Or a non-centered STFT matrix
    >>> S_left = librosa.stft(y, center=False)
    >>> librosa.get_duration(S=S_left, sr=sr)
    61.3471201814059

    Parameters
    ----------
    y : np.ndarray [shape=(n,), (2, n)] or None
        audio time series

    sr : number > 0 [scalar]
        audio sampling rate of `y`

    S : np.ndarray [shape=(d, t)] or None
        STFT matrix, or any STFT-derived matrix (e.g., chromagram
        or mel spectrogram).
        Durations calculated from spectrogram inputs are only accurate
        up to the frame resolution. If high precision is required,
        it is better to use the audio time series directly.

    n_fft       : int > 0 [scalar]
        FFT window size for `S`

    hop_length  : int > 0 [ scalar]
        number of audio samples between columns of `S`

    center  : boolean
        - If `True`, `S[:, t]` is centered at `y[t * hop_length]`
        - If `False`, then `S[:, t]` begins at `y[t * hop_length]`

    filename : str
        If provided, all other parameters are ignored, and the
        duration is calculated directly from the audio file.
        Note that this avoids loading the contents into memory,
        and is therefore useful for querying the duration of
        long files.

    Returns
    -------
    d : float >= 0
        Duration (in seconds) of the input time series or spectrogram.

    Raises
    ------
    ParameterError
        if none of `y`, `S`, or `filename` are provided.

    Notes
    -----
    `get_duration` can be applied to a file (`filename`), a spectrogram (`S`),
    or audio buffer (`y, sr`).  Only one of these three options should be
    provided.  If you do provide multiple options (e.g., `filename` and `S`),
    then `filename` takes precedence over `S`, and `S` takes precedence over
    `(y, sr)`.
    """

    if filename is not None:
        try:
            return sf.info(filename).duration
        except:
            with audioread.audio_open(filename) as fdesc:
                return fdesc.duration

    if y is None:
        if S is None:
            raise ParameterError('At least one of (y, sr), S, or filename must be provided')

        n_frames = S.shape[1]
        n_samples = n_fft + hop_length * (n_frames - 1)

        # If centered, we lose half a window from each end of S
        if center:
            n_samples = n_samples - 2 * int(n_fft / 2)

    else:
        # Validate the audio buffer.  Stereo is okay here.
        util.valid_audio(y, mono=False)
        if y.ndim == 1:
            n_samples = len(y)
        else:
            n_samples = y.shape[-1]

    return float(n_samples) / sr

@cache(level=20)
def autocorrelate(y, max_size=None, axis=-1):
    """Bounded auto-correlation

    Parameters
    ----------
    y : np.ndarray
        array to autocorrelate

    max_size  : int > 0 or None
        maximum correlation lag.
        If unspecified, defaults to `y.shape[axis]` (unbounded)

    axis : int
        The axis along which to autocorrelate.
        By default, the last axis (-1) is taken.

    Returns
    -------
    z : np.ndarray
        truncated autocorrelation `y*y` along the specified axis.
        If `max_size` is specified, then `z.shape[axis]` is bounded
        to `max_size`.

    Notes
    -----
    This function caches at level 20.

    Examples
    --------
    Compute full autocorrelation of y

    >>> y, sr = librosa.load(librosa.util.example_audio_file(), offset=20, duration=10)
    >>> librosa.autocorrelate(y)
    array([  3.226e+03,   3.217e+03, ...,   8.277e-04,   3.575e-04], dtype=float32)

    Compute onset strength auto-correlation up to 4 seconds

    >>> import matplotlib.pyplot as plt
    >>> odf = librosa.onset.onset_strength(y=y, sr=sr, hop_length=512)
    >>> ac = librosa.autocorrelate(odf, max_size=4* sr / 512)
    >>> plt.plot(ac)
    >>> plt.title('Auto-correlation')
    >>> plt.xlabel('Lag (frames)')

    """

    if max_size is None:
        max_size = y.shape[axis]

    max_size = int(min(max_size, y.shape[axis]))

    # Compute the power spectrum along the chosen axis
    # Pad out the signal to support full-length auto-correlation.
    fft = get_fftlib()
    powspec = np.abs(fft.fft(y, n=2 * y.shape[axis] + 1, axis=axis))**2

    # Convert back to time domain
    autocorr = fft.ifft(powspec, axis=axis)

    # Slice down to max_size
    subslice = [slice(None)] * autocorr.ndim
    subslice[axis] = slice(max_size)

    autocorr = autocorr[tuple(subslice)]

    if not np.iscomplexobj(y):
        autocorr = autocorr.real

    return autocorr

def lpc(y, order):
    """Linear Prediction Coefficients via Burg's method

    This function applies Burg's method to estimate coefficients of a linear
    filter on `y` of order `order`.  Burg's method is an extension to the
    Yule-Walker approach, which are both sometimes referred to as LPC parameter
    estimation by autocorrelation.

    It follows the description and implementation approach described in the
    introduction in [1]_.  N.B. This paper describes a different method, which
    is not implemented here, but has been chosen for its clear explanation of
    Burg's technique in its introduction.

    .. [1] Larry Marple
           A New Autoregressive Spectrum Analysis Algorithm
           IEEE Transactions on Accoustics, Speech, and Signal Processing
           vol 28, no. 4, 1980

    Parameters
    ----------
    y : np.ndarray
        Time series to fit

    order : int > 0
        Order of the linear filter

    Returns
    -------
    a : np.ndarray of length order + 1
        LP prediction error coefficients, i.e. filter denominator polynomial

    Raises
    ------
    ParameterError
        - If y is not valid audio as per `util.valid_audio`
        - If order < 1 or not integer
    FloatingPointError
        - If y is ill-conditioned

    See also
    --------
    scipy.signal.lfilter

    Examples
    --------
    Compute LP coefficients of y at order 16 on entire series

    >>> y, sr = librosa.load(librosa.util.example_audio_file(), offset=30,
    ...                      duration=10)
    >>> librosa.lpc(y, 16)

    Compute LP coefficients, and plot LP estimate of original series

    >>> import matplotlib.pyplot as plt
    >>> import scipy
    >>> y, sr = librosa.load(librosa.util.example_audio_file(), offset=30,
    ...                      duration=0.020)
    >>> a = librosa.lpc(y, 2)
    >>> y_hat = scipy.signal.lfilter([0] + -1*a[1:], [1], y)
    >>> plt.figure()
    >>> plt.plot(y)
    >>> plt.plot(y_hat)
    >>> plt.legend(['y', 'y_hat'])
    >>> plt.title('LP Model Forward Prediction')

    """
    if not isinstance(order, int) or order < 1:
        raise ParameterError("order must be an integer > 0")

    util.valid_audio(y, mono=True)

    return __lpc(y, order)

@jit(nopython=True)
def __lpc(y, order):
    # This implementation follows the description of Burg's algorithm given in
    # section III of Marple's paper referenced in the docstring.
    #
    # We use the Levinson-Durbin recursion to compute AR coefficients for each
    # increasing model order by using those from the last. We maintain two
    # arrays and then flip them each time we increase the model order so that
    # we may use all the coefficients from the previous order while we compute
    # those for the new one. These two arrays hold ar_coeffs for order M and
    # order M-1.  (Corresponding to a_{M,k} and a_{M-1,k} in eqn 5)
    ar_coeffs = np.zeros(order+1, dtype=y.dtype)
    ar_coeffs[0] = 1
    ar_coeffs_prev = np.zeros(order+1, dtype=y.dtype)
    ar_coeffs_prev[0] = 1

    # These two arrays hold the forward and backward prediction error. They
    # correspond to f_{M-1,k} and b_{M-1,k} in eqns 10, 11, 13 and 14 of
    # Marple. First they are used to compute the reflection coefficient at
    # order M from M-1 then are re-used as f_{M,k} and b_{M,k} for each
    # iteration of the below loop
    fwd_pred_error = y[1:]
    bwd_pred_error = y[:-1]

    # DEN_{M} from eqn 16 of Marple.
    den = np.dot(fwd_pred_error, fwd_pred_error) \
        + np.dot(bwd_pred_error, bwd_pred_error)

    for i in range(order):
        if den <= 0:
            raise FloatingPointError('numerical error, input ill-conditioned?')

        # Eqn 15 of Marple, with fwd_pred_error and bwd_pred_error
        # corresponding to f_{M-1,k+1} and b{M-1,k} and the result as a_{M,M}
        reflect_coeff = -2 * np.dot(bwd_pred_error, fwd_pred_error) / den

        # Now we use the reflection coefficient and the AR coefficients from
        # the last model order to compute all of the AR coefficients for the
        # current one.  This is the Levinson-Durbin recursion described in
        # eqn 5.
        # Note 1: We don't have to care about complex conjugates as our signals
        # are all real-valued
        # Note 2: j counts 1..order+1, i-j+1 counts order..0
        # Note 3: The first element of ar_coeffs* is always 1, which copies in
        # the reflection coefficient at the end of the new AR coefficient array
        # after the preceding coefficients
        ar_coeffs_prev, ar_coeffs = ar_coeffs, ar_coeffs_prev
        for j in range(1, i+2):
            ar_coeffs[j] = ar_coeffs_prev[j] + reflect_coeff*ar_coeffs_prev[i - j + 1]

        # Update the forward and backward prediction errors corresponding to
        # eqns 13 and 14.  We start with f_{M-1,k+1} and b_{M-1,k} and use them
        # to compute f_{M,k} and b_{M,k}
        fwd_pred_error_tmp = fwd_pred_error
        fwd_pred_error = fwd_pred_error + reflect_coeff*bwd_pred_error
        bwd_pred_error = bwd_pred_error + reflect_coeff*fwd_pred_error_tmp

        # SNIP - we are now done with order M and advance. M-1 <- M

        # Compute DEN_{M} using the recursion from eqn 17.
        #
        # reflect_coeff = a_{M-1,M-1}      (we have advanced M)
        # den =  DEN_{M-1}                 (rhs)
        # bwd_pred_error = b_{M-1,N-M+1}   (we have advanced M)
        # fwd_pred_error = f_{M-1,k}       (we have advanced M)
        # den <- DEN_{M}                   (lhs)
        #
        q = 1 - reflect_coeff**2
        den = q*den - bwd_pred_error[-1]**2 - fwd_pred_error[0]**2

        # Shift up forward error.
        #
        # fwd_pred_error <- f_{M-1,k+1}
        # bwd_pred_error <- b_{M-1,k}
        #
        # N.B. We do this after computing the denominator using eqn 17 but
        # before using it in the numerator in eqn 15.
        fwd_pred_error = fwd_pred_error[1:]
        bwd_pred_error = bwd_pred_error[:-1]

    return ar_coeffs

@cache(level=20)
def zero_crossings(y, threshold=1e-10, ref_magnitude=None, pad=True,
                   zero_pos=True, axis=-1):
    '''Find the zero-crossings of a signal `y`: indices `i` such that
    `sign(y[i]) != sign(y[j])`.

    If `y` is multi-dimensional, then zero-crossings are computed along
    the specified `axis`.

    Parameters
    ----------
    y : np.ndarray
        The input array

    threshold : float > 0 or None
        If specified, values where `-threshold <= y <= threshold` are
        clipped to 0.

    ref_magnitude : float > 0 or callable
        If numeric, the threshold is scaled relative to `ref_magnitude`.

        If callable, the threshold is scaled relative to
        `ref_magnitude(np.abs(y))`.

    pad : boolean
        If `True`, then `y[0]` is considered a valid zero-crossing.

    zero_pos : boolean
        If `True` then the value 0 is interpreted as having positive sign.

        If `False`, then 0, -1, and +1 all have distinct signs.

    axis : int
        Axis along which to compute zero-crossings.

    Returns
    -------
    zero_crossings : np.ndarray [shape=y.shape, dtype=boolean]
        Indicator array of zero-crossings in `y` along the selected axis.

    Notes
    -----
    This function caches at level 20.

    Examples
    --------
    >>> # Generate a time-series
    >>> y = np.sin(np.linspace(0, 4 * 2 * np.pi, 20))
    >>> y
    array([  0.000e+00,   9.694e-01,   4.759e-01,  -7.357e-01,
            -8.372e-01,   3.247e-01,   9.966e-01,   1.646e-01,
            -9.158e-01,  -6.142e-01,   6.142e-01,   9.158e-01,
            -1.646e-01,  -9.966e-01,  -3.247e-01,   8.372e-01,
             7.357e-01,  -4.759e-01,  -9.694e-01,  -9.797e-16])
    >>> # Compute zero-crossings
    >>> z = librosa.zero_crossings(y)
    >>> z
    array([ True, False, False,  True, False,  True, False, False,
            True, False,  True, False,  True, False, False,  True,
           False,  True, False,  True], dtype=bool)
    >>> # Stack y against the zero-crossing indicator
    >>> np.vstack([y, z]).T
    array([[  0.000e+00,   1.000e+00],
           [  9.694e-01,   0.000e+00],
           [  4.759e-01,   0.000e+00],
           [ -7.357e-01,   1.000e+00],
           [ -8.372e-01,   0.000e+00],
           [  3.247e-01,   1.000e+00],
           [  9.966e-01,   0.000e+00],
           [  1.646e-01,   0.000e+00],
           [ -9.158e-01,   1.000e+00],
           [ -6.142e-01,   0.000e+00],
           [  6.142e-01,   1.000e+00],
           [  9.158e-01,   0.000e+00],
           [ -1.646e-01,   1.000e+00],
           [ -9.966e-01,   0.000e+00],
           [ -3.247e-01,   0.000e+00],
           [  8.372e-01,   1.000e+00],
           [  7.357e-01,   0.000e+00],
           [ -4.759e-01,   1.000e+00],
           [ -9.694e-01,   0.000e+00],
           [ -9.797e-16,   1.000e+00]])
    >>> # Find the indices of zero-crossings
    >>> np.nonzero(z)
    (array([ 0,  3,  5,  8, 10, 12, 15, 17, 19]),)
    '''

    # Clip within the threshold
    if threshold is None:
        threshold = 0.0

    if six.callable(ref_magnitude):
        threshold = threshold * ref_magnitude(np.abs(y))

    elif ref_magnitude is not None:
        threshold = threshold * ref_magnitude

    if threshold > 0:
        y = y.copy()
        y[np.abs(y) <= threshold] = 0

    # Extract the sign bit
    if zero_pos:
        y_sign = np.signbit(y)
    else:
        y_sign = np.sign(y)

    # Find the change-points by slicing
    slice_pre = [slice(None)] * y.ndim
    slice_pre[axis] = slice(1, None)

    slice_post = [slice(None)] * y.ndim
    slice_post[axis] = slice(-1)

    # Since we've offset the input by one, pad back onto the front
    padding = [(0, 0)] * y.ndim
    padding[axis] = (1, 0)

    return np.pad((y_sign[tuple(slice_post)] != y_sign[tuple(slice_pre)]),
                  padding,
                  mode='constant',
                  constant_values=pad)

def clicks(times=None, frames=None, sr=22050, hop_length=512,
           click_freq=1000.0, click_duration=0.1, click=None, length=None):
    """Returns a signal with the signal `click` placed at each specified time

    Parameters
    ----------
    times : np.ndarray or None
        times to place clicks, in seconds

    frames : np.ndarray or None
        frame indices to place clicks

    sr : number > 0
        desired sampling rate of the output signal

    hop_length : int > 0
        if positions are specified by `frames`, the number of samples between frames.

    click_freq : float > 0
        frequency (in Hz) of the default click signal.  Default is 1KHz.

    click_duration : float > 0
        duration (in seconds) of the default click signal.  Default is 100ms.

    click : np.ndarray or None
        optional click signal sample to use instead of the default blip.

    length : int > 0
        desired number of samples in the output signal

    Returns
    -------
    click_signal : np.ndarray
        Synthesized click signal

    Raises
    ------
    ParameterError
        - If neither `times` nor `frames` are provided.
        - If any of `click_freq`, `click_duration`, or `length` are out of range.

    Examples
    --------
    >>> # Sonify detected beat events
    >>> y, sr = librosa.load(librosa.util.example_audio_file())
    >>> tempo, beats = librosa.beat.beat_track(y=y, sr=sr)
    >>> y_beats = librosa.clicks(frames=beats, sr=sr)

    >>> # Or generate a signal of the same length as y
    >>> y_beats = librosa.clicks(frames=beats, sr=sr, length=len(y))

    >>> # Or use timing instead of frame indices
    >>> times = librosa.frames_to_time(beats, sr=sr)
    >>> y_beat_times = librosa.clicks(times=times, sr=sr)

    >>> # Or with a click frequency of 880Hz and a 500ms sample
    >>> y_beat_times880 = librosa.clicks(times=times, sr=sr,
    ...                                  click_freq=880, click_duration=0.5)

    Display click waveform next to the spectrogram

    >>> import matplotlib.pyplot as plt
    >>> plt.figure()
    >>> S = librosa.feature.melspectrogram(y=y, sr=sr)
    >>> ax = plt.subplot(2,1,2)
    >>> librosa.display.specshow(librosa.power_to_db(S, ref=np.max),
    ...                          x_axis='time', y_axis='mel')
    >>> plt.subplot(2,1,1, sharex=ax)
    >>> librosa.display.waveplot(y_beat_times, sr=sr, label='Beat clicks')
    >>> plt.legend()
    >>> plt.xlim(15, 30)
    >>> plt.tight_layout()
    """

    # Compute sample positions from time or frames
    if times is None:
        if frames is None:
            raise ParameterError('either "times" or "frames" must be provided')

        positions = frames_to_samples(frames, hop_length=hop_length)
    else:
        # Convert times to positions
        positions = time_to_samples(times, sr=sr)

    if click is not None:
        # Check that we have a well-formed audio buffer
        util.valid_audio(click, mono=True)

    else:
        # Create default click signal
        if click_duration <= 0:
            raise ParameterError('click_duration must be strictly positive')

        if click_freq <= 0:
            raise ParameterError('click_freq must be strictly positive')

        angular_freq = 2 * np.pi * click_freq / float(sr)

        click = np.logspace(0, -10,
                            num=int(np.round(sr * click_duration)),
                            base=2.0)

        click *= np.sin(angular_freq * np.arange(len(click)))

    # Set default length
    if length is None:
        length = positions.max() + click.shape[0]
    else:
        if length < 1:
            raise ParameterError('length must be a positive integer')

        # Filter out any positions past the length boundary
        positions = positions[positions < length]

    # Pre-allocate click signal
    click_signal = np.zeros(length, dtype=np.float32)

    # Place clicks
    for start in positions:
        # Compute the end-point of this click
        end = start + click.shape[0]

        if end >= length:
            click_signal[start:] += click[:length - start]
        else:
            # Normally, just add a click here
            click_signal[start:end] += click

    return click_signal

def tone(frequency, sr=22050, length=None, duration=None, phi=None):
    """Returns a pure tone signal. The signal generated is a cosine wave.

    Parameters
    ----------
    frequency : float > 0
        frequency

    sr : number > 0
        desired sampling rate of the output signal

    length : int > 0
        desired number of samples in the output signal. When both `duration` and `length` are defined,
        `length` would take priority.

    duration : float > 0
        desired duration in seconds. When both `duration` and `length` are defined, `length` would take priority.

    phi : float or None
        phase offset, in radians. If unspecified, defaults to `-np.pi * 0.5`.

    Returns
    -------
    tone_signal : np.ndarray [shape=(length,), dtype=float64]
        Synthesized pure sine tone signal

    Raises
    ------
    ParameterError
        - If `frequency` is not provided.
        - If neither `length` nor `duration` are provided.

    Examples
    --------
    >>> # Generate a pure sine tone A4
    >>> tone = librosa.tone(440, duration=1)

    >>> # Or generate the same signal using `length`
    >>> tone = librosa.tone(440, sr=22050, length=22050)

    Display spectrogram

    >>> import matplotlib.pyplot as plt
    >>> plt.figure()
    >>> S = librosa.feature.melspectrogram(y=tone)
    >>> librosa.display.specshow(librosa.power_to_db(S, ref=np.max),
    ...                          x_axis='time', y_axis='mel')
    """

    if frequency is None:
        raise ParameterError('"frequency" must be provided')

    # Compute signal length
    if length is None:
        if duration is None:
            raise ParameterError('either "length" or "duration" must be provided')
        length = duration * sr

    if phi is None:
        phi = -np.pi * 0.5

    step = 1.0 / sr
    return np.cos(2 * np.pi * frequency * (np.arange(step * length, step=step)) + phi)

def chirp(fmin, fmax, sr=22050, length=None, duration=None, linear=False, phi=None):
    """Returns a chirp signal that goes from frequency `fmin` to frequency `fmax`

    Parameters
    ----------
    fmin : float > 0
        initial frequency

    fmax : float > 0
        final frequency

    sr : number > 0
        desired sampling rate of the output signal

    length : int > 0
        desired number of samples in the output signal.
        When both `duration` and `length` are defined, `length` would take priority.

    duration : float > 0
        desired duration in seconds.
        When both `duration` and `length` are defined, `length` would take priority.

    linear : boolean
        - If `True`, use a linear sweep, i.e., frequency changes linearly with time
        - If `False`, use a exponential sweep.
        Default is `False`.

    phi : float or None
        phase offset, in radians.
        If unspecified, defaults to `-np.pi * 0.5`.

    Returns
    -------
    chirp_signal : np.ndarray [shape=(length,), dtype=float64]
        Synthesized chirp signal

    Raises
    ------
    ParameterError
        - If either `fmin` or `fmax` are not provided.
        - If neither `length` nor `duration` are provided.

    See Also
    --------
    scipy.signal.chirp

    Examples
    --------
    >>> # Generate a exponential chirp from A4 to A5
    >>> exponential_chirp = librosa.chirp(440, 880, duration=1)

    >>> # Or generate the same signal using `length`
    >>> exponential_chirp = librosa.chirp(440, 880, sr=22050, length=22050)

    >>> # Or generate a linear chirp instead
    >>> linear_chirp = librosa.chirp(440, 880, duration=1, linear=True)

    Display spectrogram for both exponential and linear chirps

    >>> import matplotlib.pyplot as plt
    >>> plt.figure()
    >>> S_exponential = librosa.feature.melspectrogram(y=exponential_chirp)
    >>> ax = plt.subplot(2,1,1)
    >>> librosa.display.specshow(librosa.power_to_db(S_exponential, ref=np.max),
    ...                          x_axis='time', y_axis='mel')
    >>> plt.subplot(2,1,2, sharex=ax)
    >>> S_linear = librosa.feature.melspectrogram(y=linear_chirp)
    >>> librosa.display.specshow(librosa.power_to_db(S_linear, ref=np.max),
    ...                          x_axis='time', y_axis='mel')
    >>> plt.tight_layout()
    """

    if fmin is None or fmax is None:
        raise ParameterError('both "fmin" and "fmax" must be provided')

    # Compute signal duration
    period = 1.0 / sr
    if length is None:
        if duration is None:
            raise ParameterError('either "length" or "duration" must be provided')
    else:
        duration = period * length

    if phi is None:
        phi = -np.pi * 0.5

    method = 'linear' if linear else 'logarithmic'
    return scipy.signal.chirp(
        np.arange(duration, step=period),
        fmin,
        duration,
        fmax,
        method=method,
        phi=phi / np.pi * 180,  # scipy.signal.chirp uses degrees for phase offset
    )


#!/usr/bin/env python
# -*- coding: utf-8 -*-
'''Rhythmic feature extraction'''

import numpy as np

from .. import util

from ..core.audio import autocorrelate
from ..util.exceptions import ParameterError
from ..filters import get_window

__all__ = ['tempogram']

# -- Rhythmic features -- #
def tempogram(y=None, sr=22050, onset_envelope=None, hop_length=512,
              win_length=384, center=True, window='hann', norm=np.inf):
    '''Compute the tempogram: local autocorrelation of the onset strength envelope. [1]_

    .. [1] Grosche, Peter, Meinard Müller, and Frank Kurth.
        "Cyclic tempogram - A mid-level tempo representation for music signals."
        ICASSP, 2010.

    Parameters
    ----------
    y : np.ndarray [shape=(n,)] or None
        Audio time series.

    sr : number > 0 [scalar]
        sampling rate of `y`

    onset_envelope : np.ndarray [shape=(n,) or (m, n)] or None
        Optional pre-computed onset strength envelope as provided by
        `onset.onset_strength`.

        If multi-dimensional, tempograms are computed independently for each
        band (first dimension).

    hop_length : int > 0
        number of audio samples between successive onset measurements

    win_length : int > 0
        length of the onset autocorrelation window (in frames/onset measurements)
        The default settings (384) corresponds to `384 * hop_length / sr ~= 8.9s`.

    center : bool
        If `True`, onset autocorrelation windows are centered.
        If `False`, windows are left-aligned.

    window : string, function, number, tuple, or np.ndarray [shape=(win_length,)]
        A window specification as in `core.stft`.

    norm : {np.inf, -np.inf, 0, float > 0, None}
        Normalization mode.  Set to `None` to disable normalization.

    Returns
    -------
    tempogram : np.ndarray [shape=(win_length, n) or (m, win_length, n)]
        Localized autocorrelation of the onset strength envelope.

        If given multi-band input (`onset_envelope.shape==(m,n)`) then
        `tempogram[i]` is the tempogram of `onset_envelope[i]`.

    Raises
    ------
    ParameterError
        if neither `y` nor `onset_envelope` are provided

        if `win_length < 1`

    See Also
    --------
    librosa.onset.onset_strength
    librosa.util.normalize
    librosa.core.stft

    Examples
    --------
    >>> # Compute local onset autocorrelation
    >>> y, sr = librosa.load(librosa.util.example_audio_file())
    >>> hop_length = 512
    >>> oenv = librosa.onset.onset_strength(y=y, sr=sr, hop_length=hop_length)
    >>> tempogram = librosa.feature.tempogram(onset_envelope=oenv, sr=sr,
    ...                                       hop_length=hop_length)
    >>> # Compute global onset autocorrelation
    >>> ac_global = librosa.autocorrelate(oenv, max_size=tempogram.shape[0])
    >>> ac_global = librosa.util.normalize(ac_global)
    >>> # Estimate the global tempo for display purposes
    >>> tempo = librosa.beat.tempo(onset_envelope=oenv, sr=sr,
    ...                            hop_length=hop_length)[0]

    >>> import matplotlib.pyplot as plt
    >>> plt.figure(figsize=(8, 8))
    >>> plt.subplot(4, 1, 1)
    >>> plt.plot(oenv, label='Onset strength')
    >>> plt.xticks([])
    >>> plt.legend(frameon=True)
    >>> plt.axis('tight')
    >>> plt.subplot(4, 1, 2)
    >>> # We'll truncate the display to a narrower range of tempi
    >>> librosa.display.specshow(tempogram, sr=sr, hop_length=hop_length,
    >>>                          x_axis='time', y_axis='tempo')
    >>> plt.axhline(tempo, color='w', linestyle='--', alpha=1,
    ...             label='Estimated tempo={:g}'.format(tempo))
    >>> plt.legend(frameon=True, framealpha=0.75)
    >>> plt.subplot(4, 1, 3)
    >>> x = np.linspace(0, tempogram.shape[0] * float(hop_length) / sr,
    ...                 num=tempogram.shape[0])
    >>> plt.plot(x, np.mean(tempogram, axis=1), label='Mean local autocorrelation')
    >>> plt.plot(x, ac_global, '--', alpha=0.75, label='Global autocorrelation')
    >>> plt.xlabel('Lag (seconds)')
    >>> plt.axis('tight')
    >>> plt.legend(frameon=True)
    >>> plt.subplot(4,1,4)
    >>> # We can also plot on a BPM axis
    >>> freqs = librosa.tempo_frequencies(tempogram.shape[0], hop_length=hop_length, sr=sr)
    >>> plt.semilogx(freqs[1:], np.mean(tempogram[1:], axis=1),
    ...              label='Mean local autocorrelation', basex=2)
    >>> plt.semilogx(freqs[1:], ac_global[1:], '--', alpha=0.75,
    ...              label='Global autocorrelation', basex=2)
    >>> plt.axvline(tempo, color='black', linestyle='--', alpha=.8,
    ...             label='Estimated tempo={:g}'.format(tempo))
    >>> plt.legend(frameon=True)
    >>> plt.xlabel('BPM')
    >>> plt.axis('tight')
    >>> plt.grid()
    >>> plt.tight_layout()
    '''

    from ..onset import onset_strength

    if win_length < 1:
        raise ParameterError('win_length must be a positive integer')

    ac_window = get_window(window, win_length, fftbins=True)

    if onset_envelope is None:
        if y is None:
            raise ParameterError('Either y or onset_envelope must be provided')

        onset_envelope = onset_strength(y=y, sr=sr, hop_length=hop_length)

    else:
        # Force row-contiguity to avoid framing errors below
        onset_envelope = np.ascontiguousarray(onset_envelope)

    if onset_envelope.ndim > 1:
        # If we have multi-band input, iterate over rows
        return np.asarray([tempogram(onset_envelope=oe_subband,
                                     hop_length=hop_length,
                                     win_length=win_length,
                                     center=center,
                                     window=window,
                                     norm=norm) for oe_subband in onset_envelope])

    # Center the autocorrelation windows
    n = len(onset_envelope)

    if center:
        onset_envelope = np.pad(onset_envelope, int(win_length // 2),
                                mode='linear_ramp', end_values=[0, 0])

    # Carve onset envelope into frames
    odf_frame = util.frame(onset_envelope,
                           frame_length=win_length,
                           hop_length=1)

    # Truncate to the length of the original signal
    if center:
        odf_frame = odf_frame[:, :n]

    # Window, autocorrelate, and normalize
    return util.normalize(autocorrelate(odf_frame * ac_window[:, np.newaxis],
                                        axis=0),
                          norm=norm, axis=0)


#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""Utility functions for dealing with files"""

import os
import glob
import pkg_resources

import six

EXAMPLE_AUDIO = 'example_data/Kevin_MacLeod_-_Vibe_Ace.ogg'

__all__ = ['example_audio_file', 'find_files']

def example_audio_file():
    '''Get the path to an included audio example file.

    Examples
    --------
    >>> # Load the waveform from the example track
    >>> y, sr = librosa.load(librosa.util.example_audio_file())

    Returns
    -------
    filename : str
        Path to the audio example file included with librosa

    .. raw:: html

        Vibe Ace
        (Kevin MacLeod)
        / CC BY 3.0

    '''

    return pkg_resources.resource_filename(__name__, EXAMPLE_AUDIO)

def find_files(directory, ext=None, recurse=True, case_sensitive=False,
               limit=None, offset=0):
    '''Get a sorted list of (audio) files in a directory or directory sub-tree.

    Examples
    --------
    >>> # Get all audio files in a directory sub-tree
    >>> files = librosa.util.find_files('~/Music')

    >>> # Look only within a specific directory, not the sub-tree
    >>> files = librosa.util.find_files('~/Music', recurse=False)

    >>> # Only look for mp3 files
    >>> files = librosa.util.find_files('~/Music', ext='mp3')

    >>> # Or just mp3 and ogg
    >>> files = librosa.util.find_files('~/Music', ext=['mp3', 'ogg'])

    >>> # Only get the first 10 files
    >>> files = librosa.util.find_files('~/Music', limit=10)

    >>> # Or last 10 files
    >>> files = librosa.util.find_files('~/Music', offset=-10)

    Parameters
    ----------
    directory : str
        Path to look for files

    ext : str or list of str
        A file extension or list of file extensions to include in the search.

        Default: `['aac', 'au', 'flac', 'm4a', 'mp3', 'ogg', 'wav']`

    recurse : boolean
        If `True`, then all subfolders of `directory` will be searched.

        Otherwise, only `directory` will be searched.

    case_sensitive : boolean
        If `False`, files matching upper-case version of
        extensions will be included.

    limit : int > 0 or None
        Return at most `limit` files. If `None`, all files are returned.

    offset : int
        Return files starting at `offset` within the list.

        Use negative values to offset from the end of the list.

    Returns
    -------
    files : list of str
        The list of audio files.
    '''

    if ext is None:
        ext = ['aac', 'au', 'flac', 'm4a', 'mp3', 'ogg', 'wav']

    elif isinstance(ext, six.string_types):
        ext = [ext]

    # Cast into a set
    ext = set(ext)

    # Generate upper-case versions
    if not case_sensitive:
        # Force to lower-case
        ext = set([e.lower() for e in ext])
        # Add in upper-case versions
        ext |= set([e.upper() for e in ext])

    files = set()

    if recurse:
        for walk in os.walk(directory):
            files |= __get_files(walk[0], ext)
    else:
        files = __get_files(directory, ext)

    files = list(files)
    files.sort()
    files = files[offset:]
    if limit is not None:
        files = files[:limit]

    return files

def __get_files(dir_name, extensions):
    '''Helper function to get files in a single directory'''

    # Expand out the directory
    dir_name = os.path.abspath(os.path.expanduser(dir_name))

    myfiles = set()

    for sub_ext in extensions:
        globstr = os.path.join(dir_name, '*' + os.path.extsep + sub_ext)
        myfiles |= set(glob.glob(globstr))

    return myfiles


#!/usr/bin/env python
'''CREATED:2013-12-08 14:28:34 by Brian McFee <brm2132@columbia.edu>

Demonstration of phase vocoder time stretching.
'''
from __future__ import print_function

import argparse
import sys
import librosa

def stretch_demo(input_file, output_file, speed):
    '''Phase-vocoder time stretch demo function.

    :parameters:
      - input_file : str
          path to input audio
      - output_file : str
          path to save output (wav)
      - speed : float > 0
          speed up by this factor
    '''

    # 1. Load the wav file, resample
    print('Loading ', input_file)

    y, sr = librosa.load(input_file)

    # 2. Time-stretch through effects module
    print('Playing back at {:3.0f}% speed'.format(speed * 100))

    y_stretch = librosa.effects.time_stretch(y, speed)

    print('Saving stretched audio to: ', output_file)
    librosa.output.write_wav(output_file, y_stretch, sr)

def process_arguments(args):
    '''Argparse function to get the program parameters'''

    parser = argparse.ArgumentParser(description='Time stretching example')

    parser.add_argument('input_file',
                        action='store',
                        help='path to the input file (wav, mp3, etc)')

    parser.add_argument('output_file',
                        action='store',
                        help='path to the stretched output (wav)')

    parser.add_argument('-s', '--speed',
                        action='store',
                        type=float,
                        default=2.0,
                        required=False,
                        help='speed')

    return vars(parser.parse_args(args))

if __name__ == '__main__':
    # get the parameters
    parameters = process_arguments(sys.argv[1:])

    # Run the HPSS code
    stretch_demo(parameters['input_file'],
                 parameters['output_file'],
                 parameters['speed'])


#!/usr/bin/env python
'''CREATED:2013-12-08 14:28:34 by Brian McFee <brm2132@columbia.edu>

Demonstration of harmonic-percussive source separation
'''
from __future__ import print_function

import argparse
import sys
import librosa

def hpss_demo(input_file, output_harmonic, output_percussive):
    '''HPSS demo function.

    :parameters:
      - input_file : str
          path to input audio
      - output_harmonic : str
          path to save output harmonic (wav)
      - output_percussive : str
          path to save output harmonic (wav)
    '''

    # 1. Load the wav file, resample
    print('Loading ', input_file)

    y, sr = librosa.load(input_file)

    # Separate components with the effects module
    print('Separating harmonics and percussives... ')
    y_harmonic, y_percussive = librosa.effects.hpss(y)

    # 5. Save the results
    print('Saving harmonic audio to: ', output_harmonic)
    librosa.output.write_wav(output_harmonic, y_harmonic, sr)

    print('Saving percussive audio to: ', output_percussive)
    librosa.output.write_wav(output_percussive, y_percussive, sr)

def process_arguments(args):
    '''Argparse function to get the program parameters'''

    parser = argparse.ArgumentParser(description='harmonic-percussive example')

    parser.add_argument('input_file',
                        action='store',
                        help='path to the input file (wav, mp3, etc)')

    parser.add_argument('output_harmonic',
                        action='store',
                        help='path to the harmonic output (wav)')

    parser.add_argument('output_percussive',
                        action='store',
                        help='path to the percussive output (wav)')

    return vars(parser.parse_args(args))

if __name__ == '__main__':
    # get the parameters
    parameters = process_arguments(sys.argv[1:])

    # Run the HPSS code
    hpss_demo(parameters['input_file'],
              parameters['output_harmonic'],
              parameters['output_percussive'])


#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
Beat and tempo
==============
.. autosummary::
   :toctree: generated/

   beat_track
   tempo
"""

import numpy as np
import scipy

from ._cache import cache
from . import core
from . import onset
from . import util
from .feature import tempogram
from .util.exceptions import ParameterError

__all__ = ['beat_track', 'tempo']

def beat_track(y=None, sr=22050, onset_envelope=None, hop_length=512,
               start_bpm=120.0, tightness=100, trim=True, bpm=None,
               units='frames'):
    r'''Dynamic programming beat tracker.

    Beats are detected in three stages, following the method of [1]_:
      1. Measure onset strength
      2. Estimate tempo from onset correlation
      3. Pick peaks in onset strength approximately consistent with estimated
         tempo

    .. [1] Ellis, Daniel PW. "Beat tracking by dynamic programming."
           Journal of New Music Research 36.1 (2007): 51-60.
           http://labrosa.ee.columbia.edu/projects/beattrack/

    Parameters
    ----------

    y : np.ndarray [shape=(n,)] or None
        audio time series

    sr : number > 0 [scalar]
        sampling rate of `y`

    onset_envelope : np.ndarray [shape=(n,)] or None
        (optional) pre-computed onset strength envelope.

    hop_length : int > 0 [scalar]
        number of audio samples between successive `onset_envelope` values

    start_bpm  : float > 0 [scalar]
        initial guess for the tempo estimator (in beats per minute)

    tightness  : float [scalar]
        tightness of beat distribution around tempo

    trim       : bool [scalar]
        trim leading/trailing beats with weak onsets

    bpm        : float [scalar]
        (optional) If provided, use `bpm` as the tempo instead of
        estimating it from `onsets`.

    units : {'frames', 'samples', 'time'}
        The units to encode detected beat events in.
        By default, 'frames' are used.

    Returns
    -------

    tempo : float [scalar, non-negative]
        estimated global tempo (in beats per minute)

    beats : np.ndarray [shape=(m,)]
        estimated beat event locations in the specified units
        (default is frame indices)

    .. note::
        If no onset strength could be detected, beat_tracker estimates 0 BPM
        and returns an empty list.

    Raises
    ------
    ParameterError
        if neither `y` nor `onset_envelope` are provided

        or if `units` is not one of 'frames', 'samples', or 'time'

    See Also
    --------
    librosa.onset.onset_strength

    Examples
    --------
    Track beats using time series input

    >>> y, sr = librosa.load(librosa.util.example_audio_file())

    >>> tempo, beats = librosa.beat.beat_track(y=y, sr=sr)
    >>> tempo
    64.599609375

    Print the first 20 beat frames

    >>> beats[:20]
    array([ 320,  357,  397,  436,  480,  525,  569,  609,  658,
            698,  737,  777,  817,  857,  896,  936,  976, 1016,
           1055, 1095])

    Or print them as timestamps

    >>> librosa.frames_to_time(beats[:20], sr=sr)
    array([  7.43 ,   8.29 ,   9.218,  10.124,  11.146,  12.19 ,
            13.212,  14.141,  15.279,  16.208,  17.113,  18.042,
            18.971,  19.9  ,  20.805,  21.734,  22.663,  23.591,
            24.497,  25.426])

    Track beats using a pre-computed onset envelope

    >>> onset_env = librosa.onset.onset_strength(y, sr=sr,
    ...                                          aggregate=np.median)
    >>> tempo, beats = librosa.beat.beat_track(onset_envelope=onset_env,
    ...                                        sr=sr)
    >>> tempo
    64.599609375
    >>> beats[:20]
    array([ 320,  357,  397,  436,  480,  525,  569,  609,  658,
            698,  737,  777,  817,  857,  896,  936,  976, 1016,
           1055, 1095])

    Plot the beat events against the onset strength envelope

    >>> import matplotlib.pyplot as plt
    >>> hop_length = 512
    >>> plt.figure(figsize=(8, 4))
    >>> times = librosa.frames_to_time(np.arange(len(onset_env)),
    ...                                sr=sr, hop_length=hop_length)
    >>> plt.plot(times, librosa.util.normalize(onset_env),
    ...          label='Onset strength')
    >>> plt.vlines(times[beats], 0, 1, alpha=0.5, color='r',
    ...            linestyle='--', label='Beats')
    >>> plt.legend(frameon=True, framealpha=0.75)
    >>> # Limit the plot to a 15-second window
    >>> plt.xlim(15, 30)
    >>> plt.gca().xaxis.set_major_formatter(librosa.display.TimeFormatter())
    >>> plt.tight_layout()
    '''

    # First, get the frame->beat strength profile if we don't already have one
    if onset_envelope is None:
        if y is None:
            raise ParameterError('y or onset_envelope must be provided')

        onset_envelope = onset.onset_strength(y=y,
                                              sr=sr,
                                              hop_length=hop_length,
                                              aggregate=np.median)

    # Do we have any onsets to grab?
    if not onset_envelope.any():
        return (0, np.array([], dtype=int))

    # Estimate BPM if one was not provided
    if bpm is None:
        bpm = tempo(onset_envelope=onset_envelope,
                    sr=sr,
                    hop_length=hop_length,
                    start_bpm=start_bpm)[0]

    # Then, run the tracker
    beats = __beat_tracker(onset_envelope,
                           bpm,
                           float(sr) / hop_length,
                           tightness,
                           trim)

    if units == 'frames':
        pass
    elif units == 'samples':
        beats = core.frames_to_samples(beats, hop_length=hop_length)
    elif units == 'time':
        beats = core.frames_to_time(beats, hop_length=hop_length, sr=sr)
    else:
        raise ParameterError('Invalid unit type: {}'.format(units))

    return (bpm, beats)

@cache(level=30)
def tempo(y=None, sr=22050, onset_envelope=None, hop_length=512, start_bpm=120,
          std_bpm=1.0, ac_size=8.0, max_tempo=320.0, aggregate=np.mean):
    """Estimate the tempo (beats per minute)

    Parameters
    ----------
    y : np.ndarray [shape=(n,)] or None
        audio time series

    sr : number > 0 [scalar]
        sampling rate of the time series

    onset_envelope    : np.ndarray [shape=(n,)]
        pre-computed onset strength envelope

    hop_length : int > 0 [scalar]
        hop length of the time series

    start_bpm : float [scalar]
        initial guess of the BPM

    std_bpm : float > 0 [scalar]
        standard deviation of tempo distribution

    ac_size : float > 0 [scalar]
        length (in seconds) of the auto-correlation window

    max_tempo : float > 0 [scalar, optional]
        If provided, only estimate tempo below this threshold

    aggregate : callable [optional]
        Aggregation function for estimating global tempo.
        If `None`, then tempo is estimated independently for each frame.

    Returns
    -------
    tempo : np.ndarray [scalar]
        estimated tempo (beats per minute)

    See Also
    --------
    librosa.onset.onset_strength
    librosa.feature.tempogram

    Notes
    -----
    This function caches at level 30.

    Examples
    --------
    >>> # Estimate a static tempo
    >>> y, sr = librosa.load(librosa.util.example_audio_file())
    >>> onset_env = librosa.onset.onset_strength(y, sr=sr)
    >>> tempo = librosa.beat.tempo(onset_envelope=onset_env, sr=sr)
    >>> tempo
    array([129.199])

    >>> # Or a dynamic tempo
    >>> dtempo = librosa.beat.tempo(onset_envelope=onset_env, sr=sr,
    ...                             aggregate=None)
    >>> dtempo
    array([ 143.555,  143.555,  143.555, ...,  161.499,  161.499,
            172.266])

    Plot the estimated tempo against the onset autocorrelation

    >>> import matplotlib.pyplot as plt
    >>> # Convert to scalar
    >>> tempo = np.asscalar(tempo)
    >>> # Compute 2-second windowed autocorrelation
    >>> hop_length = 512
    >>> ac = librosa.autocorrelate(onset_env, 2 * sr // hop_length)
    >>> freqs = librosa.tempo_frequencies(len(ac), sr=sr,
    ...                                   hop_length=hop_length)
    >>> # Plot on a BPM axis.  We skip the first (0-lag) bin.
    >>> plt.figure(figsize=(8,4))
    >>> plt.semilogx(freqs[1:], librosa.util.normalize(ac)[1:],
    ...              label='Onset autocorrelation', basex=2)
    >>> plt.axvline(tempo, 0, 1, color='r', alpha=0.75, linestyle='--',
    ...            label='Tempo: {:.2f} BPM'.format(tempo))
    >>> plt.xlabel('Tempo (BPM)')
    >>> plt.grid()
    >>> plt.title('Static tempo estimation')
    >>> plt.legend(frameon=True)
    >>> plt.axis('tight')

    Plot dynamic tempo estimates over a tempogram

    >>> plt.figure()
    >>> tg = librosa.feature.tempogram(onset_envelope=onset_env, sr=sr,
    ...                                hop_length=hop_length)
    >>> librosa.display.specshow(tg, x_axis='time', y_axis='tempo')
    >>> plt.plot(librosa.frames_to_time(np.arange(len(dtempo))), dtempo,
    ...          color='w', linewidth=1.5, label='Tempo estimate')
    >>> plt.title('Dynamic tempo estimation')
    >>> plt.legend(frameon=True, framealpha=0.75)
    """

    if start_bpm <= 0:
        raise ParameterError('start_bpm must be strictly positive')

    win_length = np.asscalar(core.time_to_frames(ac_size, sr=sr,
                                                 hop_length=hop_length))

    tg = tempogram(y=y, sr=sr,
                   onset_envelope=onset_envelope,
                   hop_length=hop_length,
                   win_length=win_length)

    # Eventually, we want this to work for time-varying tempo
    if aggregate is not None:
        tg = aggregate(tg, axis=1, keepdims=True)

    # Get the BPM values for each bin, skipping the 0-lag bin
    bpms = core.tempo_frequencies(tg.shape[0], hop_length=hop_length, sr=sr)

    # Weight the autocorrelation by a log-normal distribution
    prior = np.exp(-0.5 * ((np.log2(bpms) - np.log2(start_bpm)) / std_bpm)**2)

    # Kill everything above the max tempo
    if max_tempo is not None:
        max_idx = np.argmax(bpms < max_tempo)
        prior[:max_idx] = 0

    # Really, instead of multiplying by the prior, we should set up a
    # probabilistic model for tempo and add log-probabilities.
    # This would give us a chance to recover from null signals and
    # rely on the prior.
    # it would also make time aggregation much more natural

    # Get the maximum, weighted by the prior
    best_period = np.argmax(tg * prior[:, np.newaxis], axis=0)

    tempi = bpms[best_period]
    # Wherever the best tempo is index 0, return start_bpm
    tempi[best_period == 0] = start_bpm
    return tempi

def __beat_tracker(onset_envelope, bpm, fft_res, tightness, trim):
    """Internal function that tracks beats in an onset strength envelope.

    Parameters
    ----------
    onset_envelope : np.ndarray [shape=(n,)]
        onset strength envelope

    bpm : float [scalar]
        tempo estimate

    fft_res  : float [scalar]
        resolution of the fft (sr / hop_length)

    tightness: float [scalar]
        how closely do we adhere to bpm?

    trim : bool [scalar]
        trim leading/trailing beats with weak onsets?

    Returns
    -------
    beats : np.ndarray [shape=(n,)]
        frame numbers of beat events
    """

    if bpm <= 0:
        raise ParameterError('bpm must be strictly positive')

    # convert bpm to a sample period for searching
    period = round(60.0 * fft_res / bpm)

    # localscore is a smoothed version of AGC'd onset envelope
    localscore = __beat_local_score(onset_envelope, period)

    # run the DP
    backlink, cumscore = __beat_track_dp(localscore, period, tightness)

    # get the position of the last beat
    beats = [__last_beat(cumscore)]

    # Reconstruct the beat path from backlinks
    while backlink[beats[-1]] >= 0:
        beats.append(backlink[beats[-1]])

    # Put the beats in ascending order
    # Convert into an array of frame numbers
    beats = np.array(beats[::-1], dtype=int)

    # Discard spurious trailing beats
    beats = __trim_beats(localscore, beats, trim)

    return beats

# -- Helper functions for beat tracking
def __normalize_onsets(onsets):
    '''Maps onset strength function into the range [0, 1]'''

    norm = onsets.std(ddof=1)
    if norm > 0:
        onsets = onsets / norm
    return onsets

def __beat_local_score(onset_envelope, period):
    '''Construct the local score for an onset envlope and given period'''

    window = np.exp(-0.5 * (np.arange(-period, period+1)*32.0/period)**2)
    return scipy.signal.convolve(__normalize_onsets(onset_envelope),
                                 window,
                                 'same')

def __beat_track_dp(localscore, period, tightness):
    """Core dynamic program for beat tracking"""

    backlink = np.zeros_like(localscore, dtype=int)
    cumscore = np.zeros_like(localscore)

    # Search range for previous beat
    window = np.arange(-2 * period, -np.round(period / 2) + 1, dtype=int)

    # Make a score window, which begins biased toward start_bpm and skewed
    if tightness <= 0:
        raise ParameterError('tightness must be strictly positive')

    txwt = -tightness * (np.log(-window / period) ** 2)

    # Are we on the first beat?
    first_beat = True
    for i, score_i in enumerate(localscore):

        # Are we reaching back before time 0?
        z_pad = np.maximum(0, min(- window[0], len(window)))

        # Search over all possible predecessors
        candidates = txwt.copy()
        candidates[z_pad:] = candidates[z_pad:] + cumscore[window[z_pad:]]

        # Find the best preceding beat
        beat_location = np.argmax(candidates)

        # Add the local score
        cumscore[i] = score_i + candidates[beat_location]

        # Special case the first onset.  Stop if the localscore is small
        if first_beat and score_i < 0.01 * localscore.max():
            backlink[i] = -1
        else:
            backlink[i] = window[beat_location]
            first_beat = False

        # Update the time range
        window = window + 1

    return backlink, cumscore

def __last_beat(cumscore):
    """Get the last beat from the cumulative score array"""

    maxes = util.localmax(cumscore)
    med_score = np.median(cumscore[np.argwhere(maxes)])

    # The last of these is the last beat (since score generally increases)
    return np.argwhere((cumscore * maxes * 2 > med_score)).max()

def __trim_beats(localscore, beats, trim):
    """Final post-processing: throw out spurious leading/trailing beats"""

    smooth_boe = scipy.signal.convolve(localscore[beats],
                                       scipy.signal.hann(5),
                                       'same')

    if trim:
        threshold = 0.5 * ((smooth_boe**2).mean()**0.5)
    else:
        threshold = 0.0

    valid = np.argwhere(smooth_boe > threshold)

    return beats[valid.min():valid.max()]