import numpy as np
from scipy.sparse import coo_matrix, csr_matrix, csc_matrix
from scipy.sparse.linalg import spsolve

from compmech.logger import *


def remove_null_cols(*args, **kwargs):
    """Remove null rows and cols of a symmetric, square sparse matrix.

    Parameters
    ----------
    args : list of sparse matrices
        The first matrix in this list will be used to extract the columns
        to be removed from all the other matrices. Use :class:`csr_matrix` to
        obtain a better performance.

    Returns
    -------
    out : list of sparse matrices and removed columns
        A list with the reduced matrices in the same order of ``args`` plus
        an array containing the removed columns at the last position.

    """
    silent = kwargs.get('silent', False)
    args = list(args)
    log('Removing null columns...', level=3, silent=silent)
    num_cols = args[0].shape[1]

    if isinstance(args[0], csr_matrix):
        m = args[0]
    else:
        m = csr_matrix(args[0])
    rows, cols = m.nonzero()
    used_cols = np.unique(cols)

    for i, arg in enumerate(args):
        if isinstance(arg, csr_matrix):
            m = arg
        else:
            m = csr_matrix(arg)
        m = m[used_cols, :]
        #NOTE below, converting to csc_matrix seems to require more time than
        #     the "slow" column slicing for csr_matrix
        m = m[:, used_cols]
        args[i] = m
    args.append(used_cols)
    log('{} columns removed'.format(num_cols - used_cols.shape[0]),
            level=4, silent=silent)
    log('finished!', level=3, silent=silent)

    return args


def solve(a, b, **kwargs):
    """Wrapper for spsolve removing null columns

    The null columns of matrix ``a`` is removed such and the linear system of
    equations is solved. The corresponding values of the solution ``x`` where
    the columns are null will also be null values.

    Parameters
    ----------
    a : ndarray or sparse matrix
        A square matrix that will be converted to CSR form in the solution.
    b : scipy sparse matrix
        The matrix or vector representing the right hand side of the equation.
    kwargs : keyword arguments, optional
        Other arguments directly passed to :func:`spsolve`.

    Returns
    -------
    x : ndarray or sparse matrix
        The solution of the sparse linear equation.
        If ``b`` is a vector, then ``x`` is a vector of size ``a.shape[1]``.
        If ``b`` is a sparse matrix, then ``x`` is a matrix of size
        ``(a.shape[1], b.shape[1])``.

    """
    a, used_cols = remove_null_cols(a)
    px = spsolve(a, b[used_cols], **kwargs)
    x = np.zeros(b.shape[0], dtype=b.dtype)
    x[used_cols] = px

    return x


def make_symmetric(m):
    """Returns a new coo_matrix which is symmetric

    Convenient function to populate a sparse matrix symmetricaly. Only the
    upper triagle of matrix ``m`` has to be defined.

    The rows, cols and values are evaluated such that where ``rows > cols``
    the values will be ignored and recreated from the region where ``cols >
    rows``, in order to obtain a symmetric matrix.

    Parameters