# LICENSE: Simplified BSD https://github.com/mmp2/megaman/blob/master/LICENSE

import warnings
import numpy as np
from scipy import sparse
from scipy.linalg import eigh, eig
from scipy.sparse.linalg import lobpcg, eigs, eigsh
from sklearn.utils.validation import check_random_state

from .validation import check_array


EIGEN_SOLVERS = ['auto', 'dense', 'arpack', 'lobpcg']
BAD_EIGEN_SOLVERS = {}
AMG_KWDS = ['strength', 'aggregate', 'smooth', 'max_levels', 'max_coarse']

try:
    from pyamg import smoothed_aggregation_solver
    PYAMG_LOADED = True
    EIGEN_SOLVERS.append('amg')
except ImportError:
    PYAMG_LOADED = False
    BAD_EIGEN_SOLVERS['amg'] = """The eigen_solver was set to 'amg',
    but pyamg is not available. Please either
    install pyamg or use another method."""


def check_eigen_solver(eigen_solver, solver_kwds, size=None, nvec=None):
    """Check that the selected eigensolver is valid

    Parameters
    ----------
    eigen_solver : string
        string value to validate
    size, nvec : int (optional)
        if both provided, use the specified problem size and number of vectors
        to determine the optimal method to use with eigen_solver='auto'

    Returns
    -------
    eigen_solver : string
        The eigen solver. This only differs from the input if
        eigen_solver == 'auto' and `size` is specified.
    """
    if eigen_solver in BAD_EIGEN_SOLVERS:
        raise ValueError(BAD_EIGEN_SOLVERS[eigen_solver])
    elif eigen_solver not in EIGEN_SOLVERS:
        raise ValueError("Unrecognized eigen_solver: '{0}'."
                         "Should be one of: {1}".format(eigen_solver,
                                                        EIGEN_SOLVERS))

    if size is not None and nvec is not None:
        # do some checks of the eigensolver
        if eigen_solver == 'lobpcg' and size < 5 * nvec + 1:
            warnings.warn("lobpcg does not perform well with small matrices or "
                          "with large numbers of vectors. Switching to 'dense'")
            eigen_solver = 'dense'
            solver_kwds = None

        elif eigen_solver == 'auto':
            if size > 200 and nvec < 10:
                if PYAMG_LOADED:
                    eigen_solver = 'amg'
                    solver_kwds = None
                else:
                    eigen_solver = 'arpack'
                    solver_kwds = None
            else:
                eigen_solver = 'dense'
                solver_kwds = None

    return eigen_solver, solver_kwds


def _is_symmetric(M, tol = 1e-8):
    if sparse.isspmatrix(M):
        conditions = np.abs((M - M.T).data) < tol
    else:
        conditions = np.abs((M - M.T)) < tol
    return(np.all(conditions))


def eigen_decomposition(G, n_components=8, eigen_solver='auto',
                        random_state=None,
                        drop_first=True, largest=True, solver_kwds=None):
    """
    Function to compute the eigendecomposition of a square matrix.

    Parameters
    ----------
    G : array_like or sparse matrix
        The square matrix for which to compute the eigen-decomposition.
    n_components : integer, optional
        The number of eigenvectors to return
    eigen_solver : {'auto', 'dense', 'arpack', 'lobpcg', or 'amg'}
        'auto' :
            attempt to choose the best method for input data (default)
        'dense' :
            use standard dense matrix operations for the eigenvalue decomposition.
            For this method, M must be an array or matrix type.
            This method should be avoided for large problems.
        'arpack' :
            use arnoldi iteration in shift-invert mode. For this method,
            M may be a dense matrix, sparse matrix, or general linear operator.
            Warning: ARPACK can be unstable for some problems.  It is best to
            try several random seeds in order to check results.
        'lobpcg' :
            Locally Optimal Block Preconditioned Conjugate Gradient Method.
            A preconditioned eigensolver for large symmetric positive definite
            (SPD) generalized eigenproblems.
        'amg' :
            Algebraic Multigrid solver (requires ``pyamg`` to be installed)
            It can be faster on very large, sparse problems, but may also lead
            to instabilities.
    random_state : int seed, RandomState instance, or None (default)
        A pseudo random number generator used for the initialization of the
        lobpcg eigen vectors decomposition when eigen_solver == 'amg'.
        By default, arpack is used.
    solver_kwds : any additional keyword arguments to pass to the selected eigen_solver

    Returns
    -------
    lambdas, diffusion_map : eigenvalues, eigenvectors
    """
    n_nodes = G.shape[0]
    if drop_first:
        n_components = n_components + 1

    eigen_solver, solver_kwds = check_eigen_solver(eigen_solver, solver_kwds,
                                                   size=n_nodes,
                                                   nvec=n_components)
    random_state = check_random_state(random_state)

    # Convert G to best type for eigendecomposition
    if sparse.issparse(G):
        if G.getformat() is not 'csr':
            G.tocsr()
    G = G.astype(np.float)

    # Check for symmetry
    is_symmetric = _is_symmetric(G)

    # Try Eigen Methods:
    if eigen_solver == 'arpack':
        # This matches the internal initial state used by ARPACK
        v0 = random_state.uniform(-1, 1, G.shape[0])
        if is_symmetric: