#!/usr/bin/env python
# coding: utf-8
import librosa
import logging
import numpy as np
from scipy.spatial import distance
from scipy import signal
from scipy.ndimage import filters
import pylab as plt

import msaf
from msaf.algorithms.interface import SegmenterInterface

def median_filter(X, M=8):
    """Median filter along the first axis of the feature matrix X."""
    for i in range(X.shape[1]):
        X[:, i] = filters.median_filter(X[:, i], size=M)
    return X

def compute_gaussian_krnl(M):
    """Creates a gaussian kernel following Foote's paper."""
    g = signal.gaussian(M, M // 3., sym=True)
    G = np.dot(g.reshape(-1, 1), g.reshape(1, -1))
    G[M // 2:, :M // 2] = -G[M // 2:, :M // 2]
    G[:M // 2, M // 2:] = -G[:M // 2, M // 2:]
    return G

def compute_ssm(X, metric="seuclidean"):
    """Computes the self-similarity matrix of X."""
    D = distance.pdist(X, metric=metric)
    D = distance.squareform(D)
    D /= D.max()
    return 1 - D

def compute_nc(X, G):
    """Computes the novelty curve from the self-similarity matrix X and
        the gaussian kernel G."""
    N = X.shape[0]
    M = G.shape[0]
    nc = np.zeros(N)

    for i in range(M // 2, N - M // 2 + 1):
        nc[i] = np.sum(X[i - M // 2:i + M // 2, i - M // 2:i + M // 2] * G)

    # Normalize
    nc += nc.min()
    nc /= nc.max()
    return nc

def pick_peaks(nc, L=16):
    """Obtain peaks from a novelty curve using an adaptive threshold."""
    offset = nc.mean() / 20.

    nc = filters.gaussian_filter1d(nc, sigma=4)  # Smooth out nc

    th = filters.median_filter(nc, size=L) + offset
    #th = filters.gaussian_filter(nc, sigma=L/2., mode="nearest") + offset

    peaks = []
    for i in range(1, nc.shape[0] - 1):
        # is it a peak?
        if nc[i - 1] < nc[i] and nc[i] > nc[i + 1]:
            # is it above the threshold?
            if nc[i] > th[i]:
                peaks.append(i)
    #plt.plot(nc)
    #plt.plot(th)
    #for peak in peaks:
        #plt.axvline(peak)
    #plt.show()

    return peaks

class Segmenter(SegmenterInterface):
    """
    This script identifies the boundaries of a given track using the Foote
    method:

    Foote, J. (2000). Automatic Audio Segmentation Using a Measure Of Audio
    Novelty. In Proc. of the IEEE International Conference of Multimedia and
    Expo (pp. 452–455). New York City, NY, USA.
    """
    def processFlat(self):
        """Main process.
        Returns
        -------
        est_idxs : np.array(N)
            Estimated indeces the segment boundaries in frames.
        est_labels : np.array(N-1)
            Estimated labels for the segments.
        """
        # Preprocess to obtain features
        F = self._preprocess()

        # Normalize
        F = msaf.utils.normalize(F, norm_type=self.config["bound_norm_feats"])

        # Make sure that the M_gaussian is even
        if self.config["M_gaussian"] % 2 == 1:
            self.config["M_gaussian"] += 1

        # Median filter
        F = median_filter(F, M=self.config["m_median"])
        #plt.imshow(F.T, interpolation="nearest", aspect="auto"); plt.show()

        # Self similarity matrix
        S = compute_ssm(F)

        # Compute gaussian kernel
        G = compute_gaussian_krnl(self.config["M_gaussian"])
        #plt.imshow(S, interpolation="nearest", aspect="auto"); plt.show()

        # Compute the novelty curve
        nc = compute_nc(S, G)

        # Find peaks in the novelty curve
        est_idxs = pick_peaks(nc, L=self.config["L_peaks"])

        # Add first and last frames
        est_idxs = np.concatenate(([0], est_idxs, [F.shape[0] - 1]))

        # Empty labels
        est_labels = np.ones(len(est_idxs) - 1) * -1

        # Post process estimations
        est_idxs, est_labels = self._postprocess(est_idxs, est_labels)

        return est_idxs, est_labels
        # plt.figure(1)
        # plt.plot(nc);
        # [plt.axvline(p, color="m") for p in est_bounds]
        # [plt.axvline(b, color="g") for b in ann_bounds]
        # plt.figure(2)
        # plt.imshow(S, interpolation="nearest", aspect="auto")
        # [plt.axvline(b, color="g") for b in ann_bounds]
        # plt.show()


#!/usr/bin/python2.6
#
# Copyright (C) Christian Thurau, 2010.
# Licensed under the GNU General Public License (GPL).
# http://www.gnu.org/licenses/gpl.txt
"""
PyMF Simplex Volume Maximization [1]

    SIVM_GSAT: class for gsat-SiVM

[1] C. Thurau, K. Kersting, and C. Bauckhage. Yes We Can - Simplex Volume
Maximization for Descriptive Web-Scale Matrix Factorization. In Proc. Int.
Conf. on Information and Knowledge Management. ACM. 2010.
"""

import logging
import numpy as np
from .dist import *
from .vol import cmdet
from .sivm import SIVM

__all__ = ["SIVM_GSAT"]

class SIVM_GSAT(SIVM):
    """
    SIVM(data, num_bases=4, dist_measure='l2')

    Simplex Volume Maximization. Factorize a data matrix into two matrices s.t.
    F = | data - W*H | is minimal. H is restricted to convexity. W is iteratively
    found by maximizing the volume of the resulting simplex (see [1]). Can be
    applied to data streams using the .online_update_w(vec) function which decides
    on adding data sample "vec" to the already selected basis vectors.

    Parameters
    ----------
    data : array_like, shape (_data_dimension, _num_samples)
        the input data
    num_bases: int, optional
        Number of bases to compute (column rank of W and row rank of H).
        4 (default)
    dist_measure : one of 'l2' ,'cosine', 'l1', 'kl'
        Standard is 'l2' which maximizes the volume of the simplex. In contrast,
        'cosine' maximizes the volume of a cone (see [1] for details).
     init : string (default: 'fastmap')
        'fastmap' or 'origin'. Sets the method used for finding the very first
        basis vector. 'Origin' assumes the zero vector, 'Fastmap' picks one of
        the two vectors that have the largest pairwise distance.
    Attributes
    ----------
    W : "data_dimension x num_bases" matrix of basis vectors
    H : "num bases x num_samples" matrix of coefficients
    ferr : frobenius norm (after calling .factorize())

    Example
    -------
    Applying SIVM to some rather stupid data set:

    >>> import numpy as np
    >>> data = np.array([[1.0, 0.0, 2.0], [0.0, 1.0, 1.0]])
    >>> sivm_mdl = SIVM_GSAT(data, num_bases=2)
    >>> sivm_mdl.factorize()

    The basis vectors are now stored in sivm_mdl.W, the coefficients in sivm_mdl.H.
    To compute coefficients for an existing set of basis vectors simply    copy W
    to sivm_mdl.W, and set compute_w to False:

    >>> data = np.array([[1.5, 1.3], [1.2, 0.3]])
    >>> W = np.array([[1.0, 0.0], [0.0, 1.0]])
    >>> sivm_mdl = SIVM_GSAT(data, num_bases=2)
    >>> sivm_mdl.W = W
    >>> sivm_mdl.factorize(compute_w=False)

    The result is a set of coefficients sivm_mdl.H, s.t. data = W * sivm_mdl.H.
    """

    def init_w(self):
        self.select = range(self._num_bases)
        self.W = self.data[:, self.select]

    def online_update_w(self, vec):
        # update D if it does not exist
        k = self._num_bases
        if not hasattr(self, 'D'):
            self.D = np.zeros((k + 1, k + 1))
            self.D[:k, :k] = pdist(self.W, self.W)
            self.V = cmdet(self.D[:k, :k])

        tmp_d = self._distfunc(self.W, vec.reshape((-1,1)))
        self.D[k, :-1] = tmp_d
        self.D[:-1, k] = tmp_d

        v = np.zeros((self._num_bases + 1))

        for i in range(self._num_bases):
                # compute volume for each combination...
                s = np.setdiff1d(range(self._num_bases + 1), [i])
                v[i] = cmdet((self.D[s,:])[:,s])

        # select index that maximizes the volume
        v[-1] = self.V
        s = np.argmax(v)

        if s < self._num_bases:
            self.W[:,s] = vec
            self.D[:self._num_bases, :self._num_bases] = pdist(self.W, self.W)

            if not hasattr(self, '_v'):
                self._v = [self.V]
            self.V = v[s]
            self._v.append(v[s])

            self._logger.info('Volume increased:' + str(self.V))
            return True, s

        return False,-1

    def update_w(self):
        n = np.int(np.floor(np.random.random() * self._num_samples))
        if n not in self.select:
            updated, s = self.online_update_w(self.data[:,n])
            if updated:
                self.select[s] = n
                self._logger.info('Current selection:' + str(self.select))

    def factorize(self, show_progress=False, compute_w=True, compute_h=True,
                  compute_err=True, niter=1):
        """ Factorize s.t. WH = data

            Parameters
            ----------
            show_progress : bool
                    print some extra information to stdout.
            niter : int
                    number of iterations.
            compute_h : bool
                    iteratively update values for H.
            compute_w : bool
                    iteratively update values for W.
            compute_err : bool
                    compute Frobenius norm |data-WH| after each update and store
                    it to .ferr[k].

            Updated Values
            --------------
            .W : updated values for W.
            .H : updated values for H.
            .ferr : Frobenius norm |data-WH|.
        """
        if show_progress:
            self._logger.setLevel(logging.INFO)
        else:
            self._logger.setLevel(logging.ERROR)

        # create W and H if they don't already exist
        # -> any custom initialization to W,H should be done before
        if not hasattr(self,'W'):
               self.init_w()

        if not hasattr(self,'H'):
                self.init_h()

        if compute_err:
            self.ferr = np.zeros(niter)

        for i in range(niter):
            if compute_w:
                self.update_w()

            if compute_h:
                self.update_h()

            if compute_err:
                self.ferr[i] = self.frobenius_norm()
                self._logger.info('Iteration ' + str(i+1) + '/' + str(niter) +
                    ' FN:' + str(self.ferr[i]))
            else:
                self._logger.info('Iteration ' + str(i+1) + '/' + str(niter))

if __name__ == "__main__":
    import doctest
    doctest.testmod()


#!/usr/bin/env python
# coding: utf-8
import librosa
import logging
import numpy as np
from scipy.spatial import distance
from scipy import signal
from scipy.ndimage import filters

from msaf.algorithms.interface import SegmenterInterface
import msaf.utils as U

def median_filter(X, M=8):
    """Median filter along the first axis of the feature matrix X."""
    for i in range(X.shape[1]):
        X[:, i] = filters.median_filter(X[:, i], size=M)
    return X

def gaussian_filter(X, M=8, axis=0):
    """Gaussian filter along the first axis of the feature matrix X."""
    for i in range(X.shape[axis]):
        if axis == 1:
            X[:, i] = filters.gaussian_filter(X[:, i], sigma=M / 2.)
        elif axis == 0:
            X[i, :] = filters.gaussian_filter(X[i, :], sigma=M / 2.)
    return X

def compute_gaussian_krnl(M):
    """Creates a gaussian kernel following Serra's paper."""
    g = signal.gaussian(M, M / 3., sym=True)
    G = np.dot(g.reshape(-1, 1), g.reshape(1, -1))
    G[M // 2:, :M // 2] = -G[M // 2:, :M // 2]
    G[:M // 2, M // 1:] = -G[:M // 2, M // 1:]
    return G

def compute_ssm(X, metric="seuclidean"):
    """Computes the self-similarity matrix of X."""
    D = distance.pdist(X, metric=metric)
    D = distance.squareform(D)
    D /= float(D.max())
    return 1 - D

def compute_nc(X):
    """Computes the novelty curve from the structural features."""
    N = X.shape[0]
    # nc = np.sum(np.diff(X, axis=0), axis=1) # Difference between SF's

    nc = np.zeros(N)
    for i in range(N - 1):
        nc[i] = distance.euclidean(X[i, :], X[i + 1, :])

    # Normalize
    nc += np.abs(nc.min())
    nc /= float(nc.max())
    return nc

def pick_peaks(nc, L=16, offset_denom=0.1):
    """Obtain peaks from a novelty curve using an adaptive threshold."""
    offset = nc.mean() * float(offset_denom)
    th = filters.median_filter(nc, size=L) + offset
    #th = filters.gaussian_filter(nc, sigma=L/2., mode="nearest") + offset
    #import pylab as plt
    #plt.plot(nc)
    #plt.plot(th)
    #plt.show()
    # th = np.ones(nc.shape[0]) * nc.mean() - 0.08
    peaks = []
    for i in range(1, nc.shape[0] - 1):
        # is it a peak?
        if nc[i - 1] < nc[i] and nc[i] > nc[i + 1]:
            # is it above the threshold?
            if nc[i] > th[i]:
                peaks.append(i)
    return peaks

def circular_shift(X):
    """Shifts circularly the X squre matrix in order to get a
        time-lag matrix."""
    N = X.shape[0]
    L = np.zeros(X.shape)
    for i in range(N):
        L[i, :] = np.asarray([X[(i + j) % N, j] for j in range(N)])
    return L

def embedded_space(X, m, tau=1):
    """Time-delay embedding with m dimensions and tau delays."""
    N = X.shape[0] - int(np.ceil(m))
    Y = np.zeros((N, int(np.ceil(X.shape[1] * m))))
    for i in range(N):
        # print X[i:i+m,:].flatten().shape, w, X.shape
        # print Y[i,:].shape
        rem = int((m % 1) * X.shape[1])  # Reminder for float m
        Y[i, :] = np.concatenate((X[i:i + int(m), :].flatten(),
                                 X[i + int(m), :rem]))
    return Y

class Segmenter(SegmenterInterface):
    """
    This script identifies the boundaries of a given track using the Serrà
    method:

    Serrà, J., Müller, M., Grosche, P., & Arcos, J. L. (2012). Unsupervised
    Detection of Music Boundaries by Time Series Structure Features.
    In Proc. of the 26th AAAI Conference on Artificial Intelligence
    (pp. 1613–1619).Toronto, Canada.
    """
    def processFlat(self):
        """Main process.
        Returns
        -------
        est_idxs : np.array(N)
            Estimated times for the segment boundaries in frame indeces.
        est_labels : np.array(N-1)
            Estimated labels for the segments.
        """
        # Structural Features params
        Mp = self.config["Mp_adaptive"]   # Size of the adaptive threshold for
                                          # peak picking
        od = self.config["offset_thres"]  # Offset coefficient for adaptive
                                          # thresholding

        M = self.config["M_gaussian"]     # Size of gaussian kernel in beats
        m = self.config["m_embedded"]     # Number of embedded dimensions
        k = self.config["k_nearest"]      # k*N-nearest neighbors for the
                                          # recurrence plot

        # Preprocess to obtain features, times, and input boundary indeces
        F = self._preprocess()

        # Normalize
        F = U.normalize(F, norm_type=self.config["bound_norm_feats"])

        # Check size in case the track is too short
        if F.shape[0] > 20:

            if self.framesync:
                red = 0.1
                F_copy = np.copy(F)
                F = librosa.util.utils.sync(
                    F.T, np.linspace(0, F.shape[0], num=F.shape[0] * red),
                    pad=False).T

            # Emedding the feature space (i.e. shingle)
            E = embedded_space(F, m)
            # plt.imshow(E.T, interpolation="nearest", aspect="auto"); plt.show()

            # Recurrence matrix
            R = librosa.segment.recurrence_matrix(
                E.T,
                k=k * int(F.shape[0]),
                width=1,  # zeros from the diagonal
                metric="euclidean",
                sym=True).astype(np.float32)

            # Circular shift
            L = circular_shift(R)
            #plt.imshow(L, interpolation="nearest", cmap=plt.get_cmap("binary"))
            #plt.show()

            # Obtain structural features by filtering the lag matrix
            SF = gaussian_filter(L.T, M=M, axis=1)
            SF = gaussian_filter(L.T, M=1, axis=0)
            # plt.imshow(SF.T, interpolation="nearest", aspect="auto")
            #plt.show()

            # Compute the novelty curve
            nc = compute_nc(SF)

            # Find peaks in the novelty curve
            est_bounds = pick_peaks(nc, L=Mp, offset_denom=od)

            # Re-align embedded space
            est_bounds = np.asarray(est_bounds) + int(np.ceil(m / 2.))

            if self.framesync:
                est_bounds /= red
                F = F_copy
        else:
            est_bounds = []

        # Add first and last frames
        est_idxs = np.concatenate(([0], est_bounds, [F.shape[0] - 1]))
        est_idxs = np.unique(est_idxs)

        assert est_idxs[0] == 0 and est_idxs[-1] == F.shape[0] - 1

        # Empty labels
        est_labels = np.ones(len(est_idxs) - 1) * - 1

        # Post process estimations
        est_idxs, est_labels = self._postprocess(est_idxs, est_labels)

        # plt.figure(1)
        # plt.plot(nc);
        # [plt.axvline(p, color="m", ymin=.6) for p in est_bounds]
        # [plt.axvline(b, color="b", ymax=.6, ymin=.3) for b in brian_bounds]
        # [plt.axvline(b, color="g", ymax=.3) for b in ann_bounds]
        # plt.show()

        return est_idxs, est_labels


"""This script contains methods to plot multiple aspects of the results
of MSAF.
"""

import jams
import logging
import mir_eval
import numpy as np
import os

# Local stuff
from msaf import io
from msaf import utils

translate_ids = {
    "2dfmc": "2D-FMC",
    "cnmf3": "C-NMF",
    "foote": "Ckboard",
    "levy": "CC",
    "cc": "CC",
    "olda": "OLDA",
    "serra": "SF",
    "sf": "SF",
    "siplca": "SI-PLCA"
}

def _plot_formatting(title, est_file, algo_ids, last_bound, N, output_file):
    """Formats the plot with the correct axis labels, title, ticks, and
    so on."""
    import matplotlib.pyplot as plt
    if title is None:
        title = os.path.basename(est_file).split(".")[0]
    plt.title(title)
    plt.yticks(np.arange(0, 1, 1 / float(N)) + 1 / (float(N) * 2))
    plt.gcf().subplots_adjust(bottom=0.22)
    plt.gca().set_yticklabels(algo_ids)
    plt.xlabel("Time (seconds)")
    plt.xlim((0, last_bound))
    plt.tight_layout()
    if output_file is not None:
        plt.savefig(output_file)
    plt.show()

def plot_boundaries(all_boundaries, est_file, algo_ids=None, title=None,
                    output_file=None):
    """Plots all the boundaries.

    Parameters
    ----------
    all_boundaries: list
        A list of np.arrays containing the times of the boundaries, one array
        for each algorithm.
    est_file: str
        Path to the estimated file (JSON file)
    algo_ids : list
        List of algorithm ids to to read boundaries from.
        If None, all algorithm ids are read.
    title : str
        Title of the plot. If None, the name of the file is printed instead.
    """
    import matplotlib.pyplot as plt
    N = len(all_boundaries)  # Number of lists of boundaries
    if algo_ids is None:
        algo_ids = io.get_algo_ids(est_file)

    # Translate ids
    for i, algo_id in enumerate(algo_ids):
        algo_ids[i] = translate_ids[algo_id]
    algo_ids = ["GT"] + algo_ids

    figsize = (6, 4)
    plt.figure(1, figsize=figsize, dpi=120, facecolor='w', edgecolor='k')
    for i, boundaries in enumerate(all_boundaries):
        color = "b"
        if i == 0:
            color = "g"
        for b in boundaries:
            plt.axvline(b, i / float(N), (i + 1) / float(N), color=color)
        plt.axhline(i / float(N), color="k", linewidth=1)

    # Format plot
    _plot_formatting(title, est_file, algo_ids, all_boundaries[0][-1], N,
                     output_file)

def plot_labels(all_labels, gt_times, est_file, algo_ids=None, title=None,
                output_file=None):
    """Plots all the labels.

    Parameters
    ----------
    all_labels: list
        A list of np.arrays containing the labels of the boundaries, one array
        for each algorithm.
    gt_times: np.array
        Array with the ground truth boundaries.
    est_file: str
        Path to the estimated file (JSON file)
    algo_ids : list
        List of algorithm ids to to read boundaries from.
        If None, all algorithm ids are read.
    title : str
        Title of the plot. If None, the name of the file is printed instead.
    """
    import matplotlib.pyplot as plt
    N = len(all_labels)  # Number of lists of labels
    if algo_ids is None:
        algo_ids = io.get_algo_ids(est_file)

    # Translate ids
    for i, algo_id in enumerate(algo_ids):
        algo_ids[i] = translate_ids[algo_id]
    algo_ids = ["GT"] + algo_ids

    # Index the labels to normalize them
    for i, labels in enumerate(all_labels):
        all_labels[i] = mir_eval.util.index_labels(labels)[0]

    # Get color map
    cm = plt.get_cmap('gist_rainbow')
    max_label = max(max(labels) for labels in all_labels)

    # To intervals
    gt_inters = utils.times_to_intervals(gt_times)

    # Plot labels
    figsize = (6, 4)
    plt.figure(1, figsize=figsize, dpi=120, facecolor='w', edgecolor='k')
    for i, labels in enumerate(all_labels):
        for label, inter in zip(labels, gt_inters):
            plt.axvspan(inter[0], inter[1], ymin=i / float(N),
                        ymax=(i + 1) / float(N), alpha=0.6,
                        color=cm(label / float(max_label)))
        plt.axhline(i / float(N), color="k", linewidth=1)

    # Draw the boundary lines
    for bound in gt_times:
        plt.axvline(bound, color="g")

    # Format plot
    _plot_formatting(title, est_file, algo_ids, gt_times[-1], N,
                     output_file)

def plot_one_track(file_struct, est_times, est_labels, boundaries_id, labels_id,
                   title=None):
    """Plots the results of one track, with ground truth if it exists."""
    import matplotlib.pyplot as plt
    # Set up the boundaries id
    bid_lid = boundaries_id
    if labels_id is not None:
        bid_lid += " + " + labels_id
    try:
        # Read file
        jam = jams.load(file_struct.ref_file)
        ann = jam.search(namespace='segment_.*')[0]
        ref_inters, ref_labels = ann.to_interval_values()

        # To times
        ref_times = utils.intervals_to_times(ref_inters)
        all_boundaries = [ref_times, est_times]
        all_labels = [ref_labels, est_labels]
        algo_ids = ["GT", bid_lid]
    except:
        logging.warning("No references found in %s. Not plotting groundtruth"
                        % file_struct.ref_file)
        all_boundaries = [est_times]
        all_labels = [est_labels]
        algo_ids = [bid_lid]

    N = len(all_boundaries)

    # Index the labels to normalize them
    for i, labels in enumerate(all_labels):
        all_labels[i] = mir_eval.util.index_labels(labels)[0]

    # Get color map
    cm = plt.get_cmap('gist_rainbow')
    max_label = max(max(labels) for labels in all_labels)

    figsize = (8, 4)
    plt.figure(1, figsize=figsize, dpi=120, facecolor='w', edgecolor='k')
    for i, boundaries in enumerate(all_boundaries):
        color = "b"
        if i == 0:
            color = "g"
        for b in boundaries:
            plt.axvline(b, i / float(N), (i + 1) / float(N), color=color)
        if labels_id is not None:
            labels = all_labels[i]
            inters = utils.times_to_intervals(boundaries)
            for label, inter in zip(labels, inters):
                plt.axvspan(inter[0], inter[1], ymin=i / float(N),
                            ymax=(i + 1) / float(N), alpha=0.6,
                            color=cm(label / float(max_label)))
        plt.axhline(i / float(N), color="k", linewidth=1)

    # Format plot
    _plot_formatting(title, os.path.basename(file_struct.audio_file), algo_ids,
                     all_boundaries[0][-1], N, None)

def plot_tree(T, res=None, title=None, cmap_id="Pastel2"):
    """Plots a given tree, containing hierarchical segmentation.

    Parameters
    ----------
    T: mir_eval.segment.tree
        A tree object containing the hierarchical segmentation.
    res: float
        Frame-rate resolution of the tree (None to use seconds).
    title: str
        Title for the plot. `None` for no title.
    cmap_id: str
        Color Map ID
    """
    import matplotlib.pyplot as plt
    def round_time(t, res=0.1):
        v = int(t / float(res)) * res
        return v

    # Get color map
    cmap = plt.get_cmap(cmap_id)

    # Get segments by level
    level_bounds = []
    for level in T.levels:
        if level == "root":
            continue
        segments = T.get_segments_in_level(level)
        level_bounds.append(segments)

    # Plot axvspans for each segment
    B = float(len(level_bounds))
    #plt.figure(figsize=figsize)
    for i, segments in enumerate(level_bounds):
        labels = utils.segment_labels_to_floats(segments)
        for segment, label in zip(segments, labels):
            #print i, label, cmap(label)
            if res is None:
                start = segment.start
                end = segment.end
                xlabel = "Time (seconds)"
            else:
                start = int(round_time(segment.start, res=res) / res)
                end = int(round_time(segment.end, res=res) / res)
                xlabel = "Time (frames)"
            plt.axvspan(start, end,
                        ymax=(len(level_bounds) - i) / B,
                        ymin=(len(level_bounds) - i - 1) / B,
                        facecolor=cmap(label))

    # Plot labels
    L = float(len(T.levels) - 1)
    plt.yticks(np.linspace(0, (L - 1) / L, num=L) + 1 / L / 2.,
               T.levels[1:][::-1])
    plt.xlabel(xlabel)
    if title is not None:
        plt.title(title)
    plt.gca().set_xlim([0, end])


#!/usr/bin/env python
# coding: utf-8
import logging
import numpy as np
import scipy.cluster.vq as vq
from sklearn import mixture
from sklearn.cluster import KMeans

# Local stuff
from . import utils_2dfmc as utils2d
from .xmeans import XMeans

import msaf.utils as U
from msaf.algorithms.interface import SegmenterInterface

def get_feat_segments(F, bound_idxs):
    """Returns a set of segments defined by the bound_idxs.

    Parameters
    ----------
    F: np.ndarray
        Matrix containing the features, one feature vector per row.
    bound_idxs: np.ndarray
        Array with boundary indeces.

    Returns
    -------
    feat_segments: list
        List of segments, one for each boundary interval.
    """
    # Make sure bound_idxs are not empty
    assert len(bound_idxs) > 0, "Boundaries can't be empty"

    # Make sure that boundaries are sorted
    bound_idxs = np.sort(bound_idxs)

    # Make sure we're not out of bounds
    assert bound_idxs[0] >= 0 and bound_idxs[-1] < F.shape[0], \
        "Boundaries are not correct for the given feature dimensions."

    # Obtain the segments
    feat_segments = []
    for i in range(len(bound_idxs) - 1):
        feat_segments.append(F[bound_idxs[i]:bound_idxs[i + 1], :])
    return feat_segments

def feat_segments_to_2dfmc_max(feat_segments, offset=4):
    """From a list of feature segments, return a list of 2D-Fourier Magnitude
    Coefs using the maximum segment size as main size and zero pad the rest.

    Parameters
    ----------
    feat_segments: list
        List of segments, one for each boundary interval.
    offset: int >= 0
        Number of frames to ignore from beginning and end of each segment.

    Returns
    -------
    fmcs: np.ndarray
        Tensor containing the 2D-FMC matrices, one matrix per segment.
    """
    if len(feat_segments) == 0:
        return []

    # Get maximum segment size
    max_len = max([feat_segment.shape[0] for feat_segment in feat_segments])

    fmcs = []
    for feat_segment in feat_segments:
        # Zero pad if needed
        X = np.zeros((max_len, feat_segment.shape[1]))

        # Remove a set of frames in the beginning an end of the segment
        if feat_segment.shape[0] <= offset or offset == 0:
            X[:feat_segment.shape[0], :] = feat_segment
        else:
            X[:feat_segment.shape[0] - offset, :] = \
                feat_segment[offset // 2:-offset // 2, :]

        # Compute the 2D-FMC
        try:
            fmcs.append(utils2d.compute_ffmc2d(X))
        except:
            logging.warning("Couldn't compute the 2D Fourier Transform")
            fmcs.append(np.zeros((X.shape[0] * X.shape[1]) // 2 + 1))

        # Normalize
        # fmcs[-1] = fmcs[-1] / float(fmcs[-1].max())

    return np.asarray(fmcs)

def compute_labels_kmeans(fmcs, k):
    # Removing the higher frequencies seem to yield better results
    fmcs = fmcs[:, fmcs.shape[1] // 2:]

    # Pre-process
    fmcs = np.log1p(fmcs)
    wfmcs = vq.whiten(fmcs)

    # Make sure we are not using more clusters than existing segments
    if k > fmcs.shape[0]:
        k = fmcs.shape[0]

    # K-means
    kmeans = KMeans(n_clusters=k, n_init=100)
    kmeans.fit(wfmcs)

    return kmeans.labels_

def compute_similarity(F, bound_idxs, dirichlet=False, xmeans=False, k=5,
                       offset=4):
    """Main function to compute the segment similarity of file file_struct.

    Parameters
    ----------
    F: np.ndarray
        Matrix containing one feature vector per row.
    bound_idxs: np.ndarray
        Array with the indeces of the segment boundaries.
    dirichlet: boolean
        Whether to use the dirichlet estimator of the number of unique labels.
    xmeans: boolean
        Whether to use the xmeans estimator of the number of unique labels.
    k: int > 0
        If the other two predictors are `False`, use fixed number of labels.
    offset: int >= 0
        Number of frames to ignore from beginning and end of each segment.

    Returns
    -------
    labels_est: np.ndarray
        Estimated labels, containing integer identifiers.
    """
    # Get the feature segments
    feat_segments = get_feat_segments(F, bound_idxs)

    # Get the 2D-FMCs segments
    fmcs = feat_segments_to_2dfmc_max(feat_segments, offset)
    if len(fmcs) == 0:
        return np.arange(len(bound_idxs) - 1)

    # Compute the labels using kmeans
    if dirichlet:
        k_init = np.min([fmcs.shape[0], k])
        # Only compute the dirichlet method if the fmc shape is small enough
        if fmcs.shape[1] > 500:
            labels_est = compute_labels_kmeans(fmcs, k=k)
        else:
            dpgmm = mixture.DPGMM(n_components=k_init, covariance_type='full')
            # dpgmm = mixture.VBGMM(n_components=k_init, covariance_type='full')
            dpgmm.fit(fmcs)
            k = len(dpgmm.means_)
            labels_est = dpgmm.predict(fmcs)
            # print("Estimated with Dirichlet Process:", k)
    if xmeans:
        xm = XMeans(fmcs, plot=False)
        k = xm.estimate_K_knee(th=0.01, maxK=8)
        labels_est = compute_labels_kmeans(fmcs, k=k)
        # print("Estimated with Xmeans:", k)
    else:
        labels_est = compute_labels_kmeans(fmcs, k=k)

    return labels_est

class Segmenter(SegmenterInterface):
    """
    This method labels segments using the 2D-FMC method described here:

    Nieto, O., Bello, J.P., Music Segment Similarity Using 2D-Fourier Magnitude
    Coefficients. Proc. of the 39th IEEE International Conference on Acoustics,
    Speech, and Signal Processing (ICASSP). Florence, Italy, 2014 (`PDF`_).

    .. _PDF: http://marl.smusic.nyu.edu/nieto/publications/NietoBello-ICASSP14.pdf
    """
    def processFlat(self):
        """Main process.
        Returns
        -------
        est_idx : np.array(N)
            Estimated indeces for the segment boundaries in frames.
        est_labels : np.array(N-1)
            Estimated labels for the segments.
        """
        # Preprocess to obtain features, times, and input boundary indeces
        F = self._preprocess()

        # Normalize
        F = U.normalize(F, norm_type=self.config["label_norm_feats"],
                        floor=self.config["label_norm_floor"],
                        min_db=self.config["label_norm_min_db"])

        # Find the labels using 2D-FMCs
        est_labels = compute_similarity(F, self.in_bound_idxs,
                                        dirichlet=self.config["dirichlet"],
                                        xmeans=self.config["xmeans"],
                                        k=self.config["k"],
                                        offset=self.config["2dfmc_offset"])

        # Post process estimations
        self.in_bound_idxs, est_labels = self._postprocess(self.in_bound_idxs,
                                                           est_labels)

        return self.in_bound_idxs, est_labels


# CREATED:2013-08-19 16:12:01 by Brian McFee <brm2132@columbia.edu>
#
# Ordinal LDA

import itertools
import numpy as np
import scipy.linalg
from sklearn.base import BaseEstimator, TransformerMixin

class OLDA(BaseEstimator, TransformerMixin):

    def __init__(self, sigma=1e-4):
        '''Ordinal linear discriminant analysis

        Arguments:
        ----------
        sigma : float
            Regularization parameter
        '''

        self.sigma = sigma
        self.scatter_ordinal_ = None
        self.scatter_within_ = None

    def fit(self, X, Y):
        '''Fit the OLDA model

        Parameters
        ----------
        X : array-like, shape [n_samples]
            Training data: each example is an n_features-by-* data array

        Y : array-like, shape [n_samples]
            Training labels: each label is an array of change-points
                             (eg, a list of segment boundaries)

        Returns
        -------
        self : object
        '''

        # Re-initialize the scatter matrices
        self.scatter_ordinal_ = None
        self.scatter_within_  = None

        # Reduce to partial-fit
        self.partial_fit(X, Y)

        return self

    def partial_fit(self, X, Y):
        '''Partial-fit the OLDA model

        Parameters
        ----------
        X : array-like, shape [n_samples]
            Training data: each example is an n_features-by-* data array

        Y : array-like, shape [n_samples]
            Training labels: each label is an array of change-points
                             (eg, a list of segment boundaries)

        Returns
        -------
        self : object
        '''

        for (xi, yi) in itertools.izip(X, Y):

            prev_mean       = None
            prev_length     = None

            if self.scatter_within_ is None:
                # First round: initialize
                d, n = xi.shape

                if yi[0] > 0:
                    yi = np.concatenate([np.array([0]), yi])
                if yi[-1] < n:
                    yi = np.concatenate([yi, np.array([n])])

                self.scatter_within_  = self.sigma * np.eye(d)
                self.scatter_ordinal_ = np.zeros(d)

            # iterate over segments
            for (seg_start, seg_end) in zip(yi[:-1], yi[1:]):

                seg_length = seg_end - seg_start

                if seg_length < 2:
                    continue

                seg_mean = np.mean(xi[:, seg_start:seg_end], axis=1, keepdims=True)
                seg_cov  = np.cov(xi[:, seg_start:seg_end])
                self.scatter_within_ = self.scatter_within_ + seg_length * seg_cov

                if prev_mean is not None:
                    diff_ord = seg_mean - (prev_length * prev_mean + seg_length * seg_mean) / (prev_length + seg_length)
                    self.scatter_ordinal_ = self.scatter_ordinal_ + seg_length * np.dot(diff_ord, diff_ord.T)

                    diff_ord = prev_mean - (prev_length * prev_mean + seg_length * seg_mean) / (prev_length + seg_length)
                    self.scatter_ordinal_ = self.scatter_ordinal_ + prev_length * np.dot(diff_ord, diff_ord.T)

                prev_mean = seg_mean
                prev_length = seg_length

        e_vals, e_vecs = scipy.linalg.eig(self.scatter_ordinal_, self.scatter_within_)
        self.e_vals_ = e_vals
        self.e_vecs_ = e_vecs
        self.components_ = e_vecs.T
        return self

    def transform(self, X):
        '''Transform data by FDA

        Parameters
        ----------
        X : array-like, shape [n_samples]
            Data to be transformed. Each example is a d-by-* feature matrix

        Returns
        -------
        X_new : array, shape (n_samples)
        '''

        return [self.components_.dot(xi) for xi in X]


"""
Each feature must inherit from the base class :class:`msaf.base.Features` to be
included in the whole framework.

Here is a list of all the available features:

.. autosummary::
    :toctree: generated/

    CQT
    MFCC
    PCP
    Tonnetz
    Tempogram
    Features
"""

from builtins import super
import librosa
import numpy as np

# Local stuff
from msaf import config
from msaf.base import Features
from msaf.exceptions import FeatureParamsError

class CQT(Features):
    """This class contains the implementation of the Constant-Q Transform.

    These features contain both harmonic and timbral content of the given
    audio signal.
    """
    def __init__(self, file_struct, feat_type, sr=config.sample_rate,
                 hop_length=config.hop_size, n_bins=config.cqt.bins,
                 norm=config.cqt.norm, filter_scale=config.cqt.filter_scale,
                 ref_power=config.cqt.ref_power):
        """Constructor of the class.

        Parameters
        ----------
        file_struct: `msaf.input_output.FileStruct`
            Object containing the file paths from where to extract/read
            the features.
        feat_type: `FeatureTypes`
            Enum containing the type of features.
        sr: int > 0
            Sampling rate for the analysis.
        hop_length: int > 0
            Hop size in frames for the analysis.
        n_bins: int > 0
            Number of frequency bins for the CQT.
        norm: float
            Type of norm to use for basis function normalization.
        filter_scale: float
            The scale of the filter for the CQT.
        ref_power: str
            The reference power for logarithmic scaling.
            See `configdefaults.py` for the possible values.
        """
        # Init the parent
        super().__init__(file_struct=file_struct, sr=sr, hop_length=hop_length,
                         feat_type=feat_type)
        # Init the CQT parameters
        self.n_bins = n_bins
        self.norm = norm
        self.filter_scale = filter_scale
        if ref_power == "max":
            self.ref_power = np.max
        elif ref_power == "min":
            self.ref_power = np.min
        elif ref_power == "median":
            self.ref_power = np.median
        else:
            raise FeatureParamsError("Wrong value for ref_power")

    @classmethod
    def get_id(self):
        """Identifier of these features."""
        return "cqt"

    def compute_features(self):
        """Actual implementation of the features.

        Returns
        -------
        cqt: np.array(N, F)
            The features, each row representing a feature vector for a give
            time frame/beat.
        """
        linear_cqt = np.abs(librosa.cqt(
            self._audio, sr=self.sr, hop_length=self.hop_length,
            n_bins=self.n_bins, norm=self.norm, filter_scale=self.filter_scale)
                            ) ** 2
        cqt = librosa.amplitude_to_db(linear_cqt, ref=self.ref_power).T
        return cqt

class MFCC(Features):
    """This class contains the implementation of the MFCC Features.

    The Mel-Frequency Cepstral Coefficients contain timbral content of a
    given audio signal.
    """
    def __init__(self, file_struct, feat_type, sr=config.sample_rate,
                 hop_length=config.hop_size, n_fft=config.n_fft,
                 n_mels=config.mfcc.n_mels, n_mfcc=config.mfcc.n_mfcc,
                 ref_power=config.mfcc.ref_power):
        """Constructor of the class.

        Parameters
        ----------
        file_struct: `msaf.input_output.FileStruct`
            Object containing the file paths from where to extract/read
            the features.
        feat_type: `FeatureTypes`
            Enum containing the type of features.
        sr: int > 0
            Sampling rate for the analysis.
        hop_length: int > 0
            Hop size in frames for the analysis.
        n_fft: int > 0
            Number of frames for the FFT.
        n_mels: int > 0
            Number of mel filters.
        n_mfcc: int > 0
            Number of mel coefficients.
        ref_power: function
            The reference power for logarithmic scaling.
        """
        # Init the parent
        super().__init__(file_struct=file_struct, sr=sr, hop_length=hop_length,
                         feat_type=feat_type)
        # Init the MFCC parameters
        self.n_fft = n_fft
        self.n_mels = n_mels
        self.n_mfcc = n_mfcc
        if ref_power == "max":
            self.ref_power = np.max
        elif ref_power == "min":
            self.ref_power = np.min
        elif ref_power == "median":
            self.ref_power = np.median
        else:
            raise FeatureParamsError("Wrong value for ref_power")

    @classmethod
    def get_id(self):
        """Identifier of these features."""
        return "mfcc"

    def compute_features(self):
        """Actual implementation of the features.

        Returns
        -------
        mfcc: np.array(N, F)
            The features, each row representing a feature vector for a give
            time frame/beat.
        """
        S = librosa.feature.melspectrogram(self._audio,
                                           sr=self.sr,
                                           n_fft=self.n_fft,
                                           hop_length=self.hop_length,
                                           n_mels=self.n_mels)
        log_S = librosa.amplitude_to_db(S, ref=self.ref_power)
        mfcc = librosa.feature.mfcc(S=log_S, n_mfcc=self.n_mfcc).T
        return mfcc

class PCP(Features):
    """This class contains the implementation of the Pitch Class Profiles.

    The PCPs contain harmonic content of a given audio signal.
    """
    def __init__(self, file_struct, feat_type, sr=config.sample_rate,
                 hop_length=config.hop_size, n_bins=config.pcp.bins,
                 norm=config.pcp.norm, f_min=config.pcp.f_min,
                 n_octaves=config.pcp.n_octaves):
        """Constructor of the class.

        Parameters
        ----------
        file_struct: `msaf.input_output.FileStruct`
            Object containing the file paths from where to extract/read
            the features.
        feat_type: `FeatureTypes`
            Enum containing the type of features.
        sr: int > 0
            Sampling rate for the analysis.
        hop_length: int > 0
            Hop size in frames for the analysis.
        n_bins: int > 0
            Number of bins for the CQT computation.
        norm: int > 0
            Normalization parameter.
        f_min: float > 0
            Minimum frequency.
        n_octaves: int > 0
            Number of octaves.
        """
        # Init the parent
        super().__init__(file_struct=file_struct, sr=sr, hop_length=hop_length,
                         feat_type=feat_type)
        # Init the PCP parameters
        self.n_bins = n_bins
        self.norm = norm
        self.f_min = f_min
        self.n_octaves = n_octaves

    @classmethod
    def get_id(self):
        """Identifier of these features."""
        return "pcp"

    def compute_features(self):
        """Actual implementation of the features.

        Returns
        -------
        pcp: np.array(N, F)
            The features, each row representing a feature vector for a give
            time frame/beat.
        """
        audio_harmonic, _ = self.compute_HPSS()
        pcp_cqt = np.abs(librosa.hybrid_cqt(audio_harmonic,
                                            sr=self.sr,
                                            hop_length=self.hop_length,
                                            n_bins=self.n_bins,
                                            norm=self.norm,
                                            fmin=self.f_min)) ** 2
        pcp = librosa.feature.chroma_cqt(C=pcp_cqt,
                                         sr=self.sr,
                                         hop_length=self.hop_length,
                                         n_octaves=self.n_octaves,
                                         fmin=self.f_min).T
        return pcp

class Tonnetz(Features):
    """This class contains the implementation of the Tonal Centroids.

    The Tonal Centroids (or Tonnetz) contain harmonic content of a given audio
    signal.
    """
    def __init__(self, file_struct, feat_type, sr=config.sample_rate,
                 hop_length=config.hop_size, n_bins=config.tonnetz.bins,
                 norm=config.tonnetz.norm, f_min=config.tonnetz.f_min,
                 n_octaves=config.tonnetz.n_octaves):
        """Constructor of the class.

        Parameters
        ----------
        file_struct: `msaf.input_output.FileStruct`
            Object containing the file paths from where to extract/read
            the features.
        feat_type: `FeatureTypes`
            Enum containing the type of features.
        sr: int > 0
            Sampling rate for the analysis.
        hop_length: int > 0
            Hop size in frames for the analysis.
        n_bins: int > 0
            Number of bins for the CQT computation.
        norm: int > 0
            Normalization parameter.
        f_min: float > 0
            Minimum frequency.
        n_octaves: int > 0
            Number of octaves.
        """
        # Init the parent
        super().__init__(file_struct=file_struct, sr=sr, hop_length=hop_length,
                         feat_type=feat_type)
        # Init the local parameters
        self.n_bins = n_bins
        self.norm = norm
        self.f_min = f_min
        self.n_octaves = n_octaves

    @classmethod
    def get_id(self):
        """Identifier of these features."""
        return "tonnetz"

    def compute_features(self):
        """Actual implementation of the features.

        Returns
        -------
        tonnetz: np.array(N, F)
            The features, each row representing a feature vector for a give
            time frame/beat.
        """
        pcp = PCP(self.file_struct, self.feat_type, self.sr, self.hop_length,
                  self.n_bins, self.norm, self.f_min, self.n_octaves).features
        tonnetz = librosa.feature.tonnetz(chroma=pcp.T).T
        return tonnetz

class Tempogram(Features):
    """This class contains the implementation of the Tempogram feature.

    The Tempogram contains rhythmic content of a given audio signal.
    """
    def __init__(self, file_struct, feat_type, sr=config.sample_rate,
                 hop_length=config.hop_size,
                 win_length=config.tempogram.win_length):
        """Constructor of the class.

        Parameters
        ----------
        file_struct: `msaf.input_output.FileStruct`
            Object containing the file paths from where to extract/read
            the features.
        feat_type: `FeatureTypes`
            Enum containing the type of features.
        sr: int > 0
            Sampling rate for the analysis.
        hop_length: int > 0
            Hop size in frames for the analysis.
        win_length: int > 0
            The size of the window for the tempogram.
        """
        # Init the parent
        super().__init__(file_struct=file_struct, sr=sr, hop_length=hop_length,
                         feat_type=feat_type)
        # Init the local parameters
        self.win_length = win_length

    @classmethod
    def get_id(self):
        """Identifier of these features."""
        return "tempogram"

    def compute_features(self):
        """Actual implementation of the features.

        Returns
        -------
        tempogram: np.array(N, F)
            The features, each row representing a feature vector for a give
            time frame/beat.
        """
        return librosa.feature.tempogram(self._audio, sr=self.sr,
                                         hop_length=self.hop_length,
                                         win_length=self.win_length).T


"""
These set of functions help the algorithms of MSAF to read and write files
of the Segmentation Dataset.
"""
import datetime
import glob
import json
import logging
import os
import re
from collections import defaultdict

import jams
import numpy as np
import six

# Local stuff
import msaf
from msaf import utils
from msaf.exceptions import NoEstimationsError

# Put dataset config in a global var
ds_config = msaf.config.dataset

class FileStruct:
    def __init__(self, audio_file):
        """Creates the entire file structure given the audio file."""
        self.ds_path = os.path.dirname(os.path.dirname(audio_file))
        self.audio_file = audio_file
        self.est_file = self._get_dataset_file(ds_config.estimations_dir,
                                               ds_config.estimations_ext)
        self.features_file = self._get_dataset_file(ds_config.features_dir,
                                                    ds_config.features_ext)
        self.ref_file = self._get_dataset_file(ds_config.references_dir,
                                               ds_config.references_ext)

    def _get_dataset_file(self, dir, ext):
        """Gets the desired dataset file."""
        audio_file_ext = "." + self.audio_file.split(".")[-1]
        base_file = os.path.basename(self.audio_file).replace(
            audio_file_ext, ext)
        return os.path.join(self.ds_path, dir, base_file)

    def __repr__(self):
        """Prints the file structure."""
        return "FileStruct(\n\tds_path=%s,\n\taudio_file=%s,\n\test_file=%s," \
            "\n\tfeatures_file=%s,\n\tref_file=%s\n)" % (
                self.ds_path, self.audio_file, self.est_file,
                self.features_file, self.ref_file)

def read_estimations(est_file, boundaries_id, labels_id=None, **params):
    """Reads the estimations (boundaries and/or labels) from a jams file
    containing the estimations of an algorithm.

    Parameters
    ----------
    est_file : str
        Path to the estimated file (JAMS file).
    boundaries_id : str
        Identifier of the algorithm used to compute the boundaries.
    labels_id : str
        Identifier of the algorithm used to compute the labels.
    params : dict
        Additional search parameters. E.g. {"feature" : "pcp"}.

    Returns
    -------
    boundaries : np.array((N,2))
        Array containing the estimated boundaries in intervals.
    labels : np.array(N)
        Array containing the estimated labels.
        Empty array if labels_id is None.
    """
    # Open file and read jams
    jam = jams.load(est_file)

    # Find correct estimation
    est = find_estimation(jam, boundaries_id, labels_id, params)
    if est is None:
        raise NoEstimationsError("No estimations for file: %s" % est_file)

    # Get data values
    all_boundaries, all_labels = est.to_interval_values()

    if params["hier"]:
        hier_bounds = defaultdict(list)
        hier_labels = defaultdict(list)
        for bounds, labels in zip(all_boundaries, all_labels):
            level = labels["level"]
            hier_bounds[level].append(bounds)
            hier_labels[level].append(labels["label"])
        # Order
        all_boundaries = []
        all_labels = []
        for key in sorted(list(hier_bounds.keys())):
            all_boundaries.append(np.asarray(hier_bounds[key]))
            all_labels.append(np.asarray(hier_labels[key]))

    return all_boundaries, all_labels

def read_references(audio_path, annotator_id=0):
    """Reads the boundary times and the labels.

    Parameters
    ----------
    audio_path : str
        Path to the audio file

    Returns
    -------
    ref_times : list
        List of boundary times
    ref_labels : list
        List of labels

    Raises
    ------
    IOError: if `audio_path` doesn't exist.
    """
    # Dataset path
    ds_path = os.path.dirname(os.path.dirname(audio_path))

    # Read references
    jam_path = os.path.join(ds_path, ds_config.references_dir,
                            os.path.basename(audio_path)[:-4] +
                            ds_config.references_ext)

    jam = jams.load(jam_path, validate=False)
    ann = jam.search(namespace='segment_.*')[annotator_id]
    ref_inters, ref_labels = ann.to_interval_values()

    # Intervals to times
    ref_times = utils.intervals_to_times(ref_inters)

    return ref_times, ref_labels

def align_times(times, frames):
    """Aligns the times to the closest frame times (e.g. beats).

    Parameters
    ----------
    times: np.ndarray
        Times in seconds to be aligned.
    frames: np.ndarray
        Frame times in seconds.

    Returns
    -------
    aligned_times: np.ndarray
        Aligned times.
    """
    dist = np.minimum.outer(times, frames)
    bound_frames = np.argmax(np.maximum(0, dist), axis=1)
    aligned_times = np.unique(bound_frames)
    return aligned_times

def find_estimation(jam, boundaries_id, labels_id, params):
    """Finds the correct estimation from all the estimations contained in a
    JAMS file given the specified arguments.

    Parameters
    ----------
    jam : jams.JAMS
        JAMS object.
    boundaries_id : str
        Identifier of the algorithm used to compute the boundaries.
    labels_id : str
        Identifier of the algorithm used to compute the labels.
    params : dict
        Additional search parameters. E.g. {"feature" : "pcp"}.

    Returns
    -------
    ann : jams.Annotation
        Found estimation.
        `None` if it couldn't be found.
    """
    # Use handy JAMS search interface
    namespace = "multi_segment" if params["hier"] else "segment_open"
    # TODO: This is a workaround to issue in JAMS. Should be
    # resolved in JAMS 0.2.3, but for now, this works too.
    ann = jam.search(namespace=namespace).\
        search(**{"Sandbox.boundaries_id": boundaries_id}).\
        search(**{"Sandbox.labels_id": lambda x:
                  (isinstance(x, six.string_types) and
                   re.match(labels_id, x) is not None) or x is None})
    for key, val in zip(params.keys(), params.values()):
        if isinstance(val, six.string_types):
            ann = ann.search(**{"Sandbox.%s" % key: val})
        else:
            ann = ann.search(**{"Sandbox.%s" % key: lambda x: x == val})

    # Check estimations found
    if len(ann) > 1:
        logging.warning("More than one estimation with same parameters.")

    if len(ann) > 0:
        ann = ann[0]

    # If we couldn't find anything, let's return None
    if not ann:
        ann = None

    return ann

def save_estimations(file_struct, times, labels, boundaries_id, labels_id,
                     **params):
    """Saves the segment estimations in a JAMS file.

    Parameters
    ----------
    file_struct : FileStruct
        Object with the different file paths of the current file.
    times : np.array or list
        Estimated boundary times.
        If `list`, estimated hierarchical boundaries.
    labels : np.array(N, 2)
        Estimated labels (None in case we are only storing boundary
        evaluations).
    boundaries_id : str
        Boundary algorithm identifier.
    labels_id : str
        Labels algorithm identifier.
    params : dict
        Dictionary with additional parameters for both algorithms.
    """
    # Remove features if they exist
    params.pop("features", None)

    # Get duration
    dur = get_duration(file_struct.features_file)

    # Convert to intervals and sanity check
    if 'numpy' in str(type(times)):
        # Flat check
        inters = utils.times_to_intervals(times)
        assert len(inters) == len(labels), "Number of boundary intervals " \
            "(%d) and labels (%d) do not match" % (len(inters), len(labels))
        # Put into lists to simplify the writing process later
        inters = [inters]
        labels = [labels]
    else:
        # Hierarchical check
        inters = []
        for level in range(len(times)):
            est_inters = utils.times_to_intervals(times[level])
            inters.append(est_inters)
            assert len(inters[level]) == len(labels[level]), \
                "Number of boundary intervals (%d) and labels (%d) do not " \
                "match in level %d" % (len(inters[level]), len(labels[level]),
                                       level)

    # Create new estimation
    namespace = "multi_segment" if params["hier"] else "segment_open"
    ann = jams.Annotation(namespace=namespace)

    # Find estimation in file
    if os.path.isfile(file_struct.est_file):
        jam = jams.load(file_struct.est_file, validate=False)
        curr_ann = find_estimation(jam, boundaries_id, labels_id, params)
        if curr_ann is not None:
            curr_ann.data = ann.data  # cleanup all data
            ann = curr_ann  # This will overwrite the existing estimation
        else:
            jam.annotations.append(ann)
    else:
        # Create new JAMS if it doesn't exist
        jam = jams.JAMS()
        jam.file_metadata.duration = dur
        jam.annotations.append(ann)

    # Save metadata and parameters
    ann.annotation_metadata.version = msaf.__version__
    ann.annotation_metadata.data_source = "MSAF"
    sandbox = {}
    sandbox["boundaries_id"] = boundaries_id
    sandbox["labels_id"] = labels_id
    sandbox["timestamp"] = \
        datetime.datetime.today().strftime("%Y/%m/%d %H:%M:%S")
    for key in params:
        sandbox[key] = params[key]
    ann.sandbox = sandbox

    # Save actual data
    for i, (level_inters, level_labels) in enumerate(zip(inters, labels)):
        for bound_inter, label in zip(level_inters, level_labels):
            dur = float(bound_inter[1]) - float(bound_inter[0])
            label = chr(int(label) + 65)
            if params["hier"]:
                value = {"label": label, "level": i}
            else:
                value = label
            ann.append(time=bound_inter[0], duration=dur,
                       value=value)

    # Write results
    jam.save(file_struct.est_file)

def get_all_boundary_algorithms():
    """Gets all the possible boundary algorithms in MSAF.

    Returns
    -------
    algo_ids : list
        List of all the IDs of boundary algorithms (strings).
    """
    algo_ids = []
    for name in msaf.algorithms.__all__:
        module = eval(msaf.algorithms.__name__ + "." + name)
        if module.is_boundary_type:
            algo_ids.append(module.algo_id)
    return algo_ids

def get_all_label_algorithms():
    """Gets all the possible label (structural grouping) algorithms in MSAF.

    Returns
    -------
    algo_ids : list
        List of all the IDs of label algorithms (strings).
    """
    algo_ids = []
    for name in msaf.algorithms.__all__:
        module = eval(msaf.algorithms.__name__ + "." + name)
        if module.is_label_type:
            algo_ids.append(module.algo_id)
    return algo_ids

def get_configuration(feature, annot_beats, framesync, boundaries_id,
                      labels_id):
    """Gets the configuration dictionary from the current parameters of the
    algorithms to be evaluated."""
    config = {}
    config["annot_beats"] = annot_beats
    config["feature"] = feature
    config["framesync"] = framesync
    bound_config = {}
    if boundaries_id != "gt":
        bound_config = \
            eval(msaf.algorithms.__name__ + "." + boundaries_id).config
        config.update(bound_config)
    if labels_id is not None:
        label_config = \
            eval(msaf.algorithms.__name__ + "." + labels_id).config

        # Make sure we don't have parameter name duplicates
        if labels_id != boundaries_id:
            overlap = set(bound_config.keys()). \
                intersection(set(label_config.keys()))
            assert len(overlap) == 0, \
                "Parameter %s must not exist both in %s and %s algorithms" % \
                (overlap, boundaries_id, labels_id)
        config.update(label_config)
    return config

def get_dataset_files(in_path):
    """Gets the files of the given dataset."""
    # Get audio files
    audio_files = []
    for ext in ds_config.audio_exts:
        audio_files += glob.glob(
            os.path.join(in_path, ds_config.audio_dir, "*" + ext))

    # Make sure directories exist
    utils.ensure_dir(os.path.join(in_path, ds_config.features_dir))
    utils.ensure_dir(os.path.join(in_path, ds_config.estimations_dir))
    utils.ensure_dir(os.path.join(in_path, ds_config.references_dir))

    # Get the file structs
    file_structs = []
    for audio_file in audio_files:
        file_structs.append(FileStruct(audio_file))

    # Sort by audio file name
    file_structs = sorted(file_structs,
                          key=lambda file_struct: file_struct.audio_file)

    return file_structs

def read_hier_references(jams_file, annotation_id=0, exclude_levels=[]):
    """Reads hierarchical references from a jams file.

    Parameters
    ----------
    jams_file : str
        Path to the jams file.
    annotation_id : int > 0
        Identifier of the annotator to read from.
    exclude_levels: list
        List of levels to exclude. Empty list to include all levels.

    Returns
    -------
    hier_bounds : list
        List of the segment boundary times in seconds for each level.
    hier_labels : list
        List of the segment labels for each level.
    hier_levels : list
        List of strings for the level identifiers.
    """
    hier_bounds = []
    hier_labels = []
    hier_levels = []
    jam = jams.load(jams_file)
    namespaces = ["segment_salami_upper", "segment_salami_function",
                  "segment_open", "segment_tut", "segment_salami_lower"]

    # Remove levels if needed
    for exclude in exclude_levels:
        if exclude in namespaces:
            namespaces.remove(exclude)

    # Build hierarchy references
    for ns in namespaces:
        ann = jam.search(namespace=ns)
        if not ann:
            continue
        ref_inters, ref_labels = ann[annotation_id].to_interval_values()
        hier_bounds.append(utils.intervals_to_times(ref_inters))
        hier_labels.append(ref_labels)
        hier_levels.append(ns)

    return hier_bounds, hier_labels, hier_levels

def get_duration(features_file):
    """Reads the duration of a given features file.

    Parameters
    ----------
    features_file: str
        Path to the JSON file containing the features.

    Returns
    -------
    dur: float
        Duration of the analyzed file.
    """
    with open(features_file) as f:
        feats = json.load(f)
    return float(feats["globals"]["dur"])

def write_mirex(times, labels, out_file):
    """Writes results to file using the standard MIREX format.

    Parameters
    ----------
    times: np.array
        Times in seconds of the boundaries.
    labels: np.array
        Labels associated to the segments defined by the boundaries.
    out_file: str
        Output file path to save the results.
    """
    inters = msaf.utils.times_to_intervals(times)
    assert len(inters) == len(labels)
    out_str = ""
    for inter, label in zip(inters, labels):
        out_str += "%.3f\t%.3f\t%s\n" % (inter[0], inter[1], label)
    with open(out_file, "w") as f:
        f.write(out_str[:-1])


"""
Example of algorithm for MSAF
"""
from msaf.algorithms.interface import SegmenterInterface
import numpy as np

class Segmenter(SegmenterInterface):
    def processFlat(self):
        """Main process.
        Returns
        -------
        est_idxs : np.array(N)
            Estimated indeces the segment boundaries in frame indeces.
        est_labels : np.array(N-1)
            Estimated labels for the segments.
        """
        # Preprocess to obtain features (array(n_frames, n_features))
        F = self._preprocess()

        # Do something with the default parameters
        # (these are defined in the in the config.py file).
        assert self.config["my_param1"] == 1.0

        # Identify boundaries in frame indeces with the new algorithm
        my_bounds = np.array([0, F.shape[0] - 1])

        # Label the segments (use -1 to have empty segments)
        my_labels = np.ones(len(my_bounds) - 1) * -1

        # Post process estimations
        est_idxs, est_labels = self._postprocess(my_bounds, my_labels)

        # We're done!
        return est_idxs, est_labels


#!/usr/bin/env python

import argparse
import numpy as np
import librosa
import os
import time

from joblib import Parallel, delayed
import cPickle as pickle

import msaf
from msaf.base import Features

from segmenter import features

def align_segmentation(beat_times, song):
    '''Load a ground-truth segmentation, and align times to the nearest
    detected beats.

    Arguments:
        beat_times -- array
        song -- path to the audio file

    Returns:
        segment_beats -- array
            beat-aligned segment boundaries

        segment_times -- array
            true segment times

        segment_labels -- array
            list of segment labels
    '''
    try:
        segment_times, segment_labels = msaf.io.read_references(song)
    except:
        return None, None, None
    segment_times = np.asarray(segment_times)

    # Map to intervals
    segment_intervals = msaf.utils.times_to_intervals(segment_times)

    # Map beats to intervals
    beat_intervals = np.asarray(zip(beat_times[:-1], beat_times[1:]))

    # Map beats to segments
    beat_segment_ids = librosa.util.match_intervals(beat_intervals,
                                                    segment_intervals)

    segment_beats = []
    segment_times_out = []
    segment_labels_out = []

    # print segment_times, beat_segment_ids, len(beat_times),
    # len(beat_segment_ids)
    for i in range(segment_times.shape[0]):
        hits = np.argwhere(beat_segment_ids == i)
        if len(hits) > 0 and i < len(segment_intervals) and \
                i < len(segment_labels):
            segment_beats.extend(hits[0])
            segment_times_out.append(segment_intervals[i, :])
            segment_labels_out.append(segment_labels[i])

    # Pull out the segment start times
    segment_beats = list(segment_beats)
    # segment_times_out = np.asarray(
    #  segment_times_out)[:, 0].squeeze().reshape((-1, 1))

    # if segment_times_out.ndim == 0:
    #    segment_times_out = segment_times_out[np.newaxis]
    segment_times_out = segment_times

    return segment_beats, segment_times_out, segment_labels_out

def get_annotation(song, rootpath):
    return '%s/annotations/%s.jams' % (rootpath, os.path.basename(song)[:-4])

def import_data(file_struct, rootpath, output_path, annot_beats):
    msaf.utils.ensure_dir(output_path)
    msaf.utils.ensure_dir(os.path.join(output_path, "features"))
    data_file = '%s/features/%s_annotbeatsE%d.pickle' % \
        (output_path, os.path.splitext(
            os.path.basename(file_struct.audio_file))[0], annot_beats)

    if os.path.exists(data_file):
        with open(data_file, 'r') as f:
            Data = pickle.load(f)
            print(file_struct.audio_file, 'cached!')
    else:
        X, _ = features(file_struct, annot_beats)
        pcp_obj = Features.select_features("pcp", file_struct, annot_beats,
                                           framesync=False)
        B = pcp_obj.frame_times[:pcp_obj.features.shape[0]]

        if X is None:
            return X

        Y, T, L = align_segmentation(
            get_annotation(file_struct.audio_file, rootpath), B,
            file_struct.audio_file)

        if Y is None:
            return Y

        Data = {'features': X,
                'beats': B,
                'filename': file_struct.audio_file,
                'segment_times': T,
                'segment_labels': L,
                'segments': Y}
        print(file_struct.audio_file, 'processed!')

        with open(data_file, 'w') as f:
            pickle.dump(Data, f)

    return Data

def make_dataset(n=None, n_jobs=1, rootpath='', output_path='',
                 annot_beats=False):

    audio_files = msaf.io.get_dataset_files(rootpath)

    if n is None:
        n = len(audio_files)

    data = Parallel(n_jobs=n_jobs)(delayed(import_data)(
        file_struct, rootpath, output_path, annot_beats)
        for file_struct in audio_files[:n])

    X, Y, B, T, F, L = [], [], [], [], [], []
    for d in data:
        if d is None:
            continue
        X.append(d['features'])
        Y.append(d['segments'])
        B.append(d['beats'])
        T.append(d['segment_times'])
        F.append(d['filename'])
        L.append(d['segment_labels'])

    return X, Y, B, T, F, L

if __name__ == '__main__':
    parser = argparse.ArgumentParser(
        description="Trains OLDA from the given dataset",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument("ds_path",
                        action="store",
                        help="Input dataset dir")
    parser.add_argument("output_path",
                        action="store",
                        help="Output dir to store the features")
    parser.add_argument("-n",
                        action="store",
                        dest="n",
                        type=int,
                        help="Number of files to process (`None` for all)",
                        default=None)
    parser.add_argument("-j",
                        action="store",
                        dest="n_jobs",
                        type=int,
                        help="Number of jobs (threads)",
                        default=1)
    parser.add_argument("-b",
                        action="store_true",
                        dest="annot_beats",
                        help="Use annotated beats",
                        default=False)
    args = parser.parse_args()
    start_time = time.time()
    X, Y, B, T, F, L = make_dataset(n=args.n,
                                    n_jobs=args.n_jobs,
                                    rootpath=args.ds_path,
                                    output_path=args.output_path,
                                    annot_beats=args.annot_beats)

    ds_path = args.ds_path
    if ds_path[-1] == "/":
        ds_path = ds_path[:-1]

    if args.annot_beats:
        out_path = '%s/AnnotBeats_%s_data.pickle' % (
            args.output_path, os.path.basename(ds_path))
    else:
        out_path = '%s/EstBeats_%s_data.pickle' % (
            args.output_path, os.path.basename(ds_path))
    with open(out_path, 'w') as f:
        pickle.dump((X, Y, B, T, F, L), f)


"""
Base module containing parent classes for the Features.
In following versions, base classes for algorithms should also be included
here.
"""

import collections
import datetime
from enum import Enum
import librosa
import logging
import jams
import json
import numpy as np
import os
import six

# Local stuff
import msaf
from msaf.exceptions import WrongFeaturesFormatError, NoFeaturesFileError,\
    FeaturesNotFound, FeatureTypeNotFound, FeatureParamsError, NoAudioFileError

# Three types of features at the moment:
#   - framesync: Frame-wise synchronous.
#   - est_beatsync: Beat-synchronous using estimated beats with librosa
#   - ann_beatsync: Beat-synchronous using annotated beats from ground-truth
FeatureTypes = Enum('FeatureTypes', 'framesync est_beatsync ann_beatsync')

# All available features
features_registry = {}

class MetaFeatures(type):
    """Meta-class to register the available features."""
    def __new__(meta, name, bases, class_dict):
        cls = type.__new__(meta, name, bases, class_dict)
        # Register classes that inherit from the base class Features
        if "Features" in [base.__name__ for base in bases]:
            features_registry[cls.get_id()] = cls
        return cls

class Features(six.with_metaclass(MetaFeatures)):
    """This is the base class for all the features in MSAF.

    It contains functions to automatically estimate beats, read annotated
    beats, compute beat-synchronous features, read and write features.

    It should be straightforward to add features in MSAF, simply by writing
    classes that inherit from this one.

    The `features` getter does the main job, and it returns a matrix `(N, F)`,
    where `N` is the number of frames an `F` is the number of features
    per frames.
    """
    def __init__(self, file_struct, sr, hop_length, feat_type):
        """Init function for the base class to make sure all features have
        at least these parameters as attributes.

        Parameters
        ----------
        file_struct: `msaf.input_output.FileStruct`
            Object containing the paths to the files.
        sr: int > 0
            Sampling rate of the audio file.
        hop_length: int > 0
            Hop in frames of the features to be computed.
        feat_type: `FeatureTypes`
            Enum containing the type of feature.
        """
        # Set the global parameters
        self.file_struct = file_struct
        self.sr = sr
        self.hop_length = hop_length
        self.feat_type = feat_type

        # The following attributes will be populated, if needed,
        # once the `features` getter is called
        self.dur = None  # The duration of the audio file in seconds
        self._features = None  # The actual features
        self._framesync_features = None  # Frame-sync features
        self._est_beatsync_features = None  # Estimated Beat-sync features
        self._ann_beatsync_features = None  # Annotated Beat-sync features
        self._audio = None  # Actual audio signal
        self._audio_harmonic = None  # Harmonic audio signal
        self._audio_percussive = None  # Percussive audio signal
        self._framesync_times = None  # The times of the framesync features
        self._est_beatsync_times = None  # Estimated beat-sync times
        self._est_beats_times = None  # Estimated beat times
        self._est_beats_frames = None  # Estimated beats in frames
        self._ann_beatsync_times = None  # Annotated beat-sync times
        self._ann_beats_times = None  # Annotated beat times
        self._ann_beats_frames = None  # Annotated beats in frames

        # Differentiate global params from sublcass attributes.
        # This is a bit hacky... I accept Pull Requests ^_^
        self._global_param_names = ["file_struct", "sr", "feat_type",
                                    "hop_length", "dur"]

    def compute_HPSS(self):
        """Computes harmonic-percussive source separation.

        Returns
        -------
        audio_harmonic: np.array
            The harmonic component of the audio signal
        audio_percussive: np.array
            The percussive component of the audio signal
        """
        return librosa.effects.hpss(self._audio)

    def estimate_beats(self):
        """Estimates the beats using librosa.

        Returns
        -------
        times: np.array
            Times of estimated beats in seconds.
        frames: np.array
            Frame indeces of estimated beats.
        """
        # Compute harmonic-percussive source separation if needed
        if self._audio_percussive is None:
            self._audio_harmonic, self._audio_percussive = self.compute_HPSS()

        # Compute beats
        tempo, frames = librosa.beat.beat_track(
            y=self._audio_percussive, sr=self.sr,
            hop_length=self.hop_length)

        # To times
        times = librosa.frames_to_time(frames, sr=self.sr,
                                       hop_length=self.hop_length)

        # TODO: Is this really necessary?
        if len(times) > 0 and times[0] == 0:
            times = times[1:]
            frames = frames[1:]

        return times, frames

    def read_ann_beats(self):
        """Reads the annotated beats if available.

        Returns
        -------
        times: np.array
            Times of annotated beats in seconds.
        frames: np.array
            Frame indeces of annotated beats.
        """
        times, frames = (None, None)

        # Read annotations if they exist in correct folder
        if os.path.isfile(self.file_struct.ref_file):
            try:
                jam = jams.load(self.file_struct.ref_file)
            except TypeError:
                logging.warning(
                    "Can't read JAMS file %s. Maybe it's not "
                    "compatible with current JAMS version?" %
                    self.file_struct.ref_file)
                return times, frames
            beat_annot = jam.search(namespace="beat.*")

            # If beat annotations exist, get times and frames
            if len(beat_annot) > 0:
                beats_inters, _ = beat_annot[0].to_interval_values()
                times = beats_inters[:, 0]
                frames = librosa.time_to_frames(times, sr=self.sr,
                                                hop_length=self.hop_length)
        return times, frames

    def compute_beat_sync_features(self, beat_frames, beat_times, pad):
        """Make the features beat-synchronous.

        Parameters
        ----------
        beat_frames: np.array
            The frame indeces of the beat positions.
        beat_times: np.array
            The time points of the beat positions (in seconds).
        pad: boolean
            If `True`, `beat_frames` is padded to span the full range.

        Returns
        -------
        beatsync_feats: np.array
            The beat-synchronized features.
            `None` if the beat_frames was `None`.
        beatsync_times: np.array
            The beat-synchronized times.
            `None` if the beat_frames was `None`.
        """
        if beat_frames is None:
            return None, None

        # Make beat synchronous
        beatsync_feats = librosa.util.utils.sync(self._framesync_features.T,
                                                 beat_frames, pad=pad).T

        # Assign times (and add last time if padded)
        beatsync_times = np.copy(beat_times)
        if beatsync_times.shape[0] != beatsync_feats.shape[0]:
            beatsync_times = np.concatenate((beatsync_times,
                                             [self._framesync_times[-1]]))
        return beatsync_feats, beatsync_times

    def read_features(self, tol=1e-3):
        """Reads the features from a file and stores them in the current
        object.

        Parameters
        ----------
        tol: float
            Tolerance level to detect duration of audio.
        """
        try:
            # Read JSON file
            with open(self.file_struct.features_file) as f:
                feats = json.load(f)

            # Store duration
            if self.dur is None:
                self.dur = float(feats["globals"]["dur"])

            # Check that we have the correct global parameters
            assert(np.isclose(
                self.dur, float(feats["globals"]["dur"]), rtol=tol))
            assert(self.sr == int(feats["globals"]["sample_rate"]))
            assert(self.hop_length == int(feats["globals"]["hop_length"]))
            assert(os.path.basename(self.file_struct.audio_file) ==
                   os.path.basename(feats["globals"]["audio_file"]))

            # Check for specific features params
            feat_params_err = FeatureParamsError(
                "Couldn't find features for %s id in file %s" %
                (self.get_id(), self.file_struct.features_file))
            if self.get_id() not in feats.keys():
                raise feat_params_err
            for param_name in self.get_param_names():
                value = getattr(self, param_name)
                if hasattr(value, '__call__'):
                    # Special case of functions
                    if value.__name__ != \
                            feats[self.get_id()]["params"][param_name]:
                        raise feat_params_err
                else:
                    if str(value) != \
                            feats[self.get_id()]["params"][param_name]:
                        raise feat_params_err

            # Store actual features
            self._est_beats_times = np.array(feats["est_beats"])
            self._est_beatsync_times = np.array(feats["est_beatsync_times"])
            self._est_beats_frames = librosa.core.time_to_frames(
                self._est_beats_times, sr=self.sr, hop_length=self.hop_length)
            self._framesync_features = \
                np.array(feats[self.get_id()]["framesync"])
            self._est_beatsync_features = \
                np.array(feats[self.get_id()]["est_beatsync"])

            # Read annotated beats if available
            if "ann_beats" in feats.keys():
                self._ann_beats_times = np.array(feats["ann_beats"])
                self._ann_beatsync_times = np.array(feats["ann_beatsync_times"])
                self._ann_beats_frames = librosa.core.time_to_frames(
                    self._ann_beats_times, sr=self.sr,
                    hop_length=self.hop_length)
                self._ann_beatsync_features = \
                    np.array(feats[self.get_id()]["ann_beatsync"])
        except KeyError:
            raise WrongFeaturesFormatError(
                "The features file %s is not correctly formatted" %
                self.file_struct.features_file)
        except AssertionError:
            raise FeaturesNotFound(
                "The features for the given parameters were not found in "
                "features file %s" % self.file_struct.features_file)
        except IOError:
            raise NoFeaturesFileError("Could not find features file %s",
                                      self.file_struct.features_file)

    def write_features(self):
        """Saves features to file."""
        out_json = collections.OrderedDict()
        try:
            # Only save the necessary information
            self.read_features()
        except (WrongFeaturesFormatError, FeaturesNotFound,
                NoFeaturesFileError):
            # We need to create the file or overwite it
            # Metadata
            out_json = collections.OrderedDict({"metadata": {
                "versions": {"librosa": librosa.__version__,
                             "msaf": msaf.__version__,
                             "numpy": np.__version__},
                "timestamp": datetime.datetime.today().strftime(
                    "%Y/%m/%d %H:%M:%S")}})

            # Global parameters
            out_json["globals"] = {
                "dur": self.dur,
                "sample_rate": self.sr,
                "hop_length": self.hop_length,
                "audio_file": self.file_struct.audio_file
            }

            # Beats
            out_json["est_beats"] = self._est_beats_times.tolist()
            out_json["est_beatsync_times"] = self._est_beatsync_times.tolist()
            if self._ann_beats_times is not None:
                out_json["ann_beats"] = self._ann_beats_times.tolist()
                out_json["ann_beatsync_times"] = self._ann_beatsync_times.tolist()
        except FeatureParamsError:
            # We have other features in the file, simply add these ones
            with open(self.file_struct.features_file) as f:
                out_json = json.load(f)
        finally:
            # Specific parameters of the current features
            out_json[self.get_id()] = {}
            out_json[self.get_id()]["params"] = {}
            for param_name in self.get_param_names():
                value = getattr(self, param_name)
                # Check for special case of functions
                if hasattr(value, '__call__'):
                    value = value.__name__
                else:
                    value = str(value)
                out_json[self.get_id()]["params"][param_name] = value

            # Actual features
            out_json[self.get_id()]["framesync"] = \
                self._framesync_features.tolist()
            out_json[self.get_id()]["est_beatsync"] = \
                self._est_beatsync_features.tolist()
            if self._ann_beatsync_features is not None:
                out_json[self.get_id()]["ann_beatsync"] = \
                    self._ann_beatsync_features.tolist()

            # Save it
            with open(self.file_struct.features_file, "w") as f:
                json.dump(out_json, f, indent=2)

    def get_param_names(self):
        """Returns the parameter names for these features, avoiding
        the global parameters."""
        return [name for name in vars(self) if not name.startswith('_') and
                name not in self._global_param_names]

    def _compute_framesync_times(self):
        """Computes the framesync times based on the framesync features."""
        self._framesync_times = librosa.core.frames_to_time(
            np.arange(self._framesync_features.shape[0]), self.sr,
            self.hop_length)

    def _compute_all_features(self):
        """Computes all the features (beatsync, framesync) from the audio."""
        # Read actual audio waveform
        self._audio, _ = librosa.load(self.file_struct.audio_file,
                                      sr=self.sr)

        # Get duration of audio file
        self.dur = len(self._audio) / float(self.sr)

        # Compute actual features
        self._framesync_features = self.compute_features()

        # Compute framesync times
        self._compute_framesync_times()

        # Compute/Read beats
        self._est_beats_times, self._est_beats_frames = self.estimate_beats()
        self._ann_beats_times, self._ann_beats_frames = self.read_ann_beats()

        # Beat-Synchronize
        pad = True  # Always append to the end of the features
        self._est_beatsync_features, self._est_beatsync_times = \
            self.compute_beat_sync_features(self._est_beats_frames,
                                            self._est_beats_times, pad)
        self._ann_beatsync_features, self._ann_beatsync_times = \
            self.compute_beat_sync_features(self._ann_beats_frames,
                                            self._ann_beats_times, pad)

    @property
    def frame_times(self):
        """This getter returns the frame times, for the corresponding type of
        features."""
        frame_times = None
        # Make sure we have already computed the features
        self.features
        if self.feat_type is FeatureTypes.framesync:
            self._compute_framesync_times()
            frame_times = self._framesync_times
        elif self.feat_type is FeatureTypes.est_beatsync:
            frame_times = self._est_beatsync_times
        elif self.feat_type is FeatureTypes.ann_beatsync:
            frame_times = self._ann_beatsync_times

        return frame_times

    @property
    def features(self):
        """This getter will compute the actual features if they haven't
        been computed yet.

        Returns
        -------
        features: np.array
            The actual features. Each row corresponds to a feature vector.
        """
        # Compute features if needed
        if self._features is None:
            try:
                self.read_features()
            except (NoFeaturesFileError, FeaturesNotFound,
                    WrongFeaturesFormatError, FeatureParamsError) as e:
                try:
                    self._compute_all_features()
                    self.write_features()
                except IOError:
                    if isinstance(e, FeaturesNotFound) or \
                            isinstance(e, FeatureParamsError):
                        msg = "Computation of the features is needed for " \
                            "current parameters but no audio file was found." \
                            "Please, change your parameters or add the audio" \
                            " file in %s"
                    else:
                        msg = "Couldn't find audio file in %s"
                    raise NoAudioFileError(msg % self.file_struct.audio_file)

        # Choose features based on type
        if self.feat_type is FeatureTypes.framesync:
            self._features = self._framesync_features
        elif self.feat_type is FeatureTypes.est_beatsync:
            self._features = self._est_beatsync_features
        elif self.feat_type is FeatureTypes.ann_beatsync:
            if self._ann_beatsync_features is None:
                raise FeatureTypeNotFound(
                    "Feature type %s is not valid because no annotated beats "
                    "were found" % self.feat_type)
            self._features = self._ann_beatsync_features
        else:
            raise FeatureTypeNotFound("Feature type %s is not valid." %
                                      self.feat_type)

        return self._features

    @classmethod
    def select_features(cls, features_id, file_struct, annot_beats, framesync):
        """Selects the features from the given parameters.

        Parameters
        ----------
        features_id: str
            The identifier of the features (it must be a key inside the
            `features_registry`)
        file_struct: msaf.io.FileStruct
            The file struct containing the files to extract the features from
        annot_beats: boolean
            Whether to use annotated (`True`) or estimated (`False`) beats
        framesync: boolean
            Whether to use framesync (`True`) or beatsync (`False`) features

        Returns
        -------
        features: obj
            The actual features object that inherits from `msaf.Features`
        """
        if not annot_beats and framesync:
            feat_type = FeatureTypes.framesync
        elif annot_beats and not framesync:
            feat_type = FeatureTypes.ann_beatsync
        elif not annot_beats and not framesync:
            feat_type = FeatureTypes.est_beatsync
        else:
            raise FeatureTypeNotFound("Type of features not valid.")

        # Select features with default parameters
        if features_id not in features_registry.keys():
            raise FeaturesNotFound(
                "The features '%s' are invalid (valid features are %s)"
                % (features_id, features_registry.keys()))

        return features_registry[features_id](file_struct, feat_type)

    def compute_features(self):
        raise NotImplementedError("This method must contain the actual "
                                  "implementation of the features")

    @classmethod
    def get_id(self):
        raise NotImplementedError("This method must return a string identifier"
                                  " of the features")


"""Interface for all the algorithms in MSAF."""
import numpy as np
import msaf.utils as U

class SegmenterInterface:
    """This class is an interface for all the segmenter algorithms included
    in MSAF. These segmenters must inherit from it and implement one of the
    following methods:
            processFlat()
            processHierarchical()

    Additionally, two private helper functions are provided:
        - preprocess
        - postprocess

    These are meant to do common tasks for all the segmenters and they
    should be called inside the process method if needed.

    All segmenters must return estimates times for the boundaries (est_times),
    and estimated labels (est_labels), **even if they can't compute them**.

    The following three types of algorithms with their behaviors are:
        - Computes boundaries and labels:
            If in_bound_times is None:
                Compute the est_times
            Else:
                Do not compute est_times, simply use in_bound_times instead
            If in_labels is None:
                Compute the est_labels
            Else:
                Do not compute est_labels, simply use in_labels instead

        - Computes boundaries only:
            Compute boundaries and return est_labels as None.

        - Computes labels only:
            Use in_bound_times in order to compute the labels.
            Return est_times as in_bound_times and the computed labels.

    In these cases, est_times or est_labels will be empty (None).
    """
    def __init__(self, file_struct, in_bound_idxs=None, feature="pcp",
                 annot_beats=False, framesync=False, features=None, **config):
        """Inits the Segmenter.

        Parameters
        ----------
        file_struct: `msaf.io.FileStruct`
            Object with the file paths.
        in_bound_idxs: np.array
            Array containing the frame indeces of the previously find
            boundaries. `None` for computing them.
        feature: str
            Identifier of the features (e.g., pcp, mfcc)
        annot_beats: boolean
            Whether to use annotated beats or estimated ones.
        framesync: boolean
            Whether to use frame-synchronous or beat-synchronous features.
        features: dict
            Previously computed features. `None` for reading them.
        config: dict
            Configuration for the given algorithm (see module's __config.py__).
        """
        self.file_struct = file_struct
        self.audio_file = file_struct.audio_file
        self.in_bound_idxs = in_bound_idxs
        self.feature_str = feature
        self.annot_beats = annot_beats
        self.framesync = framesync
        self.config = config
        self.features = features

    def processFlat(self):
        """Main process to obtain the flat segmentation of a given track."""
        raise NotImplementedError("This method does not return flat "
                                  "segmentations.")

    def processHierarchical(self):
        """Main process to obtian the hierarchical segmentation of a given
        track."""
        raise NotImplementedError("This method does not return hierarchical "
                                  "segmentations.")

    def _preprocess(self, valid_features=["pcp", "tonnetz", "mfcc",
                                          "cqt", "tempogram"]):
        """This method obtains the actual features."""
        # Use specific feature
        if self.feature_str not in valid_features:
            raise RuntimeError("Feature %s in not valid for algorithm: %s "
                               "(valid features are %s)." %
                               (self.feature_str, __name__, valid_features))
        else:
            try:
                F = self.features.features
            except KeyError:
                raise RuntimeError("Feature %s in not supported by MSAF" %
                                   (self.feature_str))

        return F

    def _postprocess(self, est_idxs, est_labels):
        """Post processes the estimations from the algorithm, removing empty
        segments and making sure the lenghts of the boundaries and labels
        match."""
        # Make sure we are using the previously input bounds, if any
        if self.in_bound_idxs is not None:
            F = self._preprocess()
            est_labels = U.synchronize_labels(self.in_bound_idxs, est_idxs,
                                              est_labels, F.shape[0])
            est_idxs = self.in_bound_idxs

        # Remove empty segments if needed
        est_idxs, est_labels = U.remove_empty_segments(est_idxs, est_labels)

        assert len(est_idxs) - 1 == len(est_labels), "Number of boundaries " \
            "(%d) and number of labels(%d) don't match" % (len(est_idxs),
                                                           len(est_labels))

        # Make sure the indeces are integers
        est_idxs = np.asarray(est_idxs, dtype=int)

        return est_idxs, est_labels