#!/usr/bin/env python
# -*- coding: utf-8 -*-
'''Non-negative least squares'''

# The scipy library provides an nnls solver, but it does
# not generalize efficiently to matrix-valued problems.
# We therefore provide an alternate solver here.
#
# The vectorized solver uses the L-BFGS-B over blocks of
# data to efficiently solve the constrained least-squares problem.

import numpy as np
import scipy.optimize
from .utils import MAX_MEM_BLOCK


__all__ = ['nnls']


def _nnls_obj(x, shape, A, B):
    '''Compute the objective and gradient for NNLS'''

    # Scipy's lbfgs flattens all arrays, so we first reshape
    # the iterate x
    x = x.reshape(shape)

    # Compute the difference matrix
    diff = np.dot(A, x) - B
    
    # Compute the objective value
    value = 0.5 * np.sum(diff**2)

    # And the gradient
    grad = np.dot(A.T, diff)

    # Flatten the gradient
    return value, grad.flatten()


def _nnls_lbfgs_block(A, B, x_init=None, **kwargs):
    '''Solve the constrained problem over a single block

    Parameters
    ----------
    A : np.ndarray [shape=(m, d)]
        The basis matrix

    B : np.ndarray [shape=(m, N)]
        The regression targets

    x_init : np.ndarray [shape=(d, N)]
        An initial guess

    kwargs
        Additional keyword arguments to `scipy.optimize.fmin_l_bfgs_b`

    Returns
    -------
    x : np.ndarray [shape=(d, N)]
        Non-negative matrix such that Ax ~= B
    '''

    # If we don't have an initial point, start at the projected
    # least squares solution
    if x_init is None:
        x_init = np.linalg.lstsq(A, B, rcond=None)[0]
        np.clip(x_init, 0, None, out=x_init)

    # Adapt the hessian approximation to the dimension of the problem
    kwargs.setdefault('m', A.shape[1])

    # Construct non-negative bounds
    bounds = [(0, None)] * x_init.size
    shape = x_init.shape

    # optimize
    x, obj_value, diagnostics = scipy.optimize.fmin_l_bfgs_b(_nnls_obj, x_init,
                                                             args=(shape, A, B),
                                                             bounds=bounds,
                                                             **kwargs)
    # reshape the solution
    return x.reshape(shape)


def nnls(A, B, **kwargs):
    '''Non-negative least squares.

    Given two matrices A and B, find a non-negative matrix X
    that minimizes the sum squared error:

        err(X) = sum_i,j ((AX)[i,j] - B[i, j])^2

    Parameters
    ----------
    A : np.ndarray [shape=(m, n)]
        The basis matrix

    B : np.ndarray [shape=(m, N)]
        The target matrix.

    kwargs
        Additional keyword arguments to `scipy.optimize.fmin_l_bfgs_b`

    Returns
    -------
    X : np.ndarray [shape=(n, N), non-negative]
        A minimizing solution to |AX - B|^2

    See Also
    --------
    scipy.optimize.nnls
    scipy.optimize.fmin_l_bfgs_b

    Examples
    --------
    Approximate a magnitude spectrum from its mel spectrogram

    >>> y, sr = librosa.load(librosa.util.example_audio_file(), offset=30, duration=10)
    >>> S = np.abs(librosa.stft(y, n_fft=2048))
    >>> M = librosa.feature.melspectrogram(S=S, sr=sr, power=1)
    >>> mel_basis = librosa.filters.mel(sr, n_fft=2048, n_mels=M.shape[0])
    >>> S_recover = librosa.util.nnls(mel_basis, M)

    Plot the results

    >>> import matplotlib.pyplot as plt
    >>> plt.figure()
    >>> plt.subplot(3,1,1)
    >>> librosa.display.specshow(librosa.amplitude_to_db(S, ref=np.max), y_axis='log')
    >>> plt.colorbar()
    >>> plt.title('Original spectrogram (1025 bins)')
    >>> plt.subplot(3,1,2)
    >>> librosa.display.specshow(librosa.amplitude_to_db(M, ref=np.max),
    ...                          y_axis='mel')
    >>> plt.title('Mel spectrogram (128 bins)')
    >>> plt.colorbar()
    >>> plt.subplot(3,1,3)
    >>> librosa.display.specshow(librosa.amplitude_to_db(S_recover, ref=np.max),
    ...                          y_axis='log')
    >>> plt.colorbar()
    >>> plt.title('Reconstructed spectrogram (1025 bins)')
    >>> plt.tight_layout()
    '''

    # If B is a single vector, punt up to the scipy method
    if B.ndim == 1:
        return scipy.optimize.nnls(A, B)[0]

    n_columns = int(MAX_MEM_BLOCK // (A.shape[-1] * A.itemsize))

    # Process in blocks:
    if B.shape[-1] <= n_columns:
        return _nnls_lbfgs_block(A, B, **kwargs).astype(A.dtype)

    x = np.linalg.lstsq(A, B, rcond=None)[0].astype(A.dtype)
    np.clip(x, 0, None, out=x)
    x_init = x

    for bl_s in range(0, x.shape[-1], n_columns):
        bl_t = min(bl_s + n_columns, B.shape[-1])
        x[:, bl_s:bl_t] = _nnls_lbfgs_block(A, B[:, bl_s:bl_t],
                                            x_init=x_init[:, bl_s:bl_t],
                                            **kwargs)
    return x

@@ -4064,16 +4064,35 @@
 ayout()%0A
+    %3E%3E%3E plt.show()%0A
     '''%0A

add show plot in example code