from __future__ import absolute_import
import hashlib

import numpy as nm
import warnings

import scipy.sparse as sps
import six
from six.moves import range

warnings.simplefilter('ignore', sps.SparseEfficiencyWarning)

from sfepy.base.base import output, get_default, assert_, try_imports
from sfepy.base.timing import Timer
from sfepy.solvers.solvers import LinearSolver

def solve(mtx, rhs, solver_class=None, solver_conf=None):
    """
    Solve the linear system with the matrix `mtx` and the right-hand side
    `rhs`.

    Convenience wrapper around the linear solver classes below.
    """
    solver_class = get_default(solver_class, ScipyDirect)
    solver_conf = get_default(solver_conf, {})

    solver = solver_class(solver_conf, mtx=mtx)
    solution = solver(rhs)

    return solution

def _get_cs_matrix_hash(mtx, chunk_size=100000):
    def _gen_array_chunks(arr):
        ii = 0
        while len(arr[ii:]):
            yield arr[ii:ii+chunk_size].tobytes()
            ii += chunk_size

    sha1 = hashlib.sha1()
    for chunk in _gen_array_chunks(mtx.indptr):
        sha1.update(chunk)
    for chunk in _gen_array_chunks(mtx.indices):
        sha1.update(chunk)
    for chunk in _gen_array_chunks(mtx.data):
        sha1.update(chunk)

    digest = sha1.hexdigest()
    return digest

def _is_new_matrix(mtx, mtx_digest, force_reuse=False):
    if not isinstance(mtx, sps.csr_matrix):
        return True, mtx_digest

    if force_reuse:
        return False, mtx_digest

    id0, digest0 = mtx_digest
    id1 = id(mtx)
    digest1 = _get_cs_matrix_hash(mtx)
    if (id1 == id0) and (digest1 == digest0):
        return False, (id1, digest1)

    return True, (id1, digest1)

def standard_call(call):
    """
    Decorator handling argument preparation and timing for linear solvers.
    """
    def _standard_call(self, rhs, x0=None, conf=None, eps_a=None, eps_r=None,
                       i_max=None, mtx=None, status=None, context=None,
                       **kwargs):
        timer = Timer(start=True)

        conf = get_default(conf, self.conf)
        mtx = get_default(mtx, self.mtx)
        status = get_default(status, self.status)
        context = get_default(context, self.context)

        assert_(mtx.shape[0] == mtx.shape[1] == rhs.shape[0])
        if x0 is not None:
            assert_(x0.shape[0] == rhs.shape[0])

        result = call(self, rhs, x0, conf, eps_a, eps_r, i_max, mtx, status,
                      context=context, **kwargs)
        if isinstance(result, tuple):
            result, n_iter = result

        else:
            n_iter = -1 # Number of iterations is undefined/unavailable.

        elapsed = timer.stop()
        if status is not None:
            status['time'] = elapsed
            status['n_iter'] = n_iter

        return result

    return _standard_call

def petsc_call(call):
    """
    Decorator handling argument preparation and timing for PETSc-based linear
    solvers.
    """
    def _petsc_call(self, rhs, x0=None, conf=None, eps_a=None, eps_r=None,
                    i_max=None, mtx=None, status=None, comm=None,
                    context=None, **kwargs):
        timer = Timer(start=True)

        conf = get_default(conf, self.conf)
        mtx = get_default(mtx, self.mtx)
        status = get_default(status, self.status)
        context = get_default(context, self.context)
        comm = get_default(comm, self.comm)

        mshape = mtx.size if isinstance(mtx, self.petsc.Mat) else mtx.shape
        rshape = [rhs.size] if isinstance(rhs, self.petsc.Vec) else rhs.shape

        assert_(mshape[0] == mshape[1] == rshape[0])
        if x0 is not None:
            xshape = [x0.size] if isinstance(x0, self.petsc.Vec) else x0.shape
            assert_(xshape[0] == rshape[0])

        result = call(self, rhs, x0, conf, eps_a, eps_r, i_max, mtx, status,
                      comm, context=context, **kwargs)

        elapsed = timer.stop()
        if status is not None:
            status['time'] = elapsed
            status['n_iter'] = self.ksp.getIterationNumber()

        return result

    return _petsc_call


class ScipyDirect(LinearSolver):
    """
    Direct sparse solver from SciPy.
    """
    name = 'ls.scipy_direct'

    _parameters = [
        ('method', "{'auto', 'umfpack', 'superlu'}", 'auto', False,
         'The actual solver to use.'),
        ('use_presolve', 'bool', False, False,
         'If True, pre-factorize the matrix.'),
    ]

    def __init__(self, conf, method=None, **kwargs):
        LinearSolver.__init__(self, conf, solve=None, **kwargs)
        um = self.sls = None
        if method is None:
            method = self.conf.method

        aux = try_imports(['import scipy.linsolve as sls',
                           'import scipy.splinalg.dsolve as sls',
                           'import scipy.sparse.linalg.dsolve as sls'],
                          'cannot import scipy sparse direct solvers!')
        if 'sls' in aux:
            self.sls = aux['sls']
        else:
            raise ValueError('SuperLU not available!')

        if method in ['auto', 'umfpack']:
            aux = try_imports([
                'import scipy.linsolve.umfpack as um',
                'import scipy.splinalg.dsolve.umfpack as um',
                'import scipy.sparse.linalg.dsolve.umfpack as um',
                'import scikits.umfpack as um'])

            is_umfpack = True if 'um' in aux\
                and hasattr(aux['um'], 'UMFPACK_OK') else False
            if method == 'umfpack' and not is_umfpack:
                raise ValueError('UMFPACK not available!')
        elif method == 'superlu':
            is_umfpack = False
        else:
            raise ValueError('uknown solution method! (%s)' % method)

        if is_umfpack:
            self.sls.use_solver(useUmfpack=True,
                                assumeSortedIndices=True)
        else:
            self.sls.use_solver(useUmfpack=False)

    @standard_call
    def __call__(self, rhs, x0=None, conf=None, eps_a=None, eps_r=None,
                 i_max=None, mtx=None, status=None, **kwargs):

        if conf.use_presolve:
            self.presolve(mtx)

        if self.solve is not None:
            # Matrix is already prefactorized.
            return self.solve(rhs)
        else:
            return self.sls.spsolve(mtx, rhs)

    def presolve(self, mtx):
        is_new, mtx_digest = _is_new_matrix(mtx, self.mtx_digest)
        if is_new:
            self.solve = self.sls.factorized(mtx)
            self.mtx_digest = mtx_digest


class ScipySuperLU(ScipyDirect):
    """
    SuperLU - direct sparse solver from SciPy.
    """
    name = 'ls.scipy_superlu'

    _parameters = [
        ('use_presolve', 'bool', False, False,
         'If True, pre-factorize the matrix.'),
    ]

    def __init__(self, conf, **kwargs):
        ScipyDirect.__init__(self, conf, method='superlu', **kwargs)


class ScipyUmfpack(ScipyDirect):
    """
    UMFPACK - direct sparse solver from SciPy.
    """
    name = 'ls.scipy_umfpack'

    _parameters = [
        ('use_presolve', 'bool', False, False,
         'If True, pre-factorize the matrix.'),
    ]

    def __init__(self, conf, **kwargs):
        ScipyDirect.__init__(self, conf, method='umfpack', **kwargs)


class ScipyIterative(LinearSolver):
    """
    Interface to SciPy iterative solvers.

    The `eps_r` tolerance is both absolute and relative - the solvers
    stop when either the relative or the absolute residual is below it.
    """
    name = 'ls.scipy_iterative'

    _parameters = [
        ('method', 'str', 'cg', False,
         'The actual solver to use.'),
        ('setup_precond', 'callable', lambda mtx, context: None, False,
         """User-supplied function for the preconditioner initialization/setup.
            It is called as setup_precond(mtx, context), where mtx is the
            matrix, context is a user-supplied context, and should return one
            of {sparse matrix, dense matrix, LinearOperator}.
         """),
        ('callback', 'callable', None, False,
         """User-supplied function to call after each iteration. It is called
            as callback(xk), where xk is the current solution vector, except
            the gmres method, where the argument is the residual.
         """),
        ('i_max', 'int', 100, False,
         'The maximum number of iterations.'),
        ('eps_a', 'float', 1e-8, False,
         'The absolute tolerance for the residual.'),
        ('eps_r', 'float', 1e-8, False,
         'The relative tolerance for the residual.'),
        ('*', '*', None, False,
         'Additional parameters supported by the method.'),
    ]

    # All iterative solvers in scipy.sparse.linalg pass a solution vector into
    # a callback except those below, that take a residual vector.
    _callbacks_res = ['gmres']

    def __init__(self, conf, context=None, **kwargs):
        import scipy.sparse.linalg.isolve as la

        LinearSolver.__init__(self, conf, context=context, **kwargs)

        try:
            solver = getattr(la, self.conf.method)
        except AttributeError:
            output('scipy solver %s does not exist!' % self.conf.method)
            output('using cg instead')
            solver = la.cg
        self.solver = solver
        self.converged_reasons = {
            0 : 'successful exit',
            1 : 'number of iterations',
            -1 : 'illegal input or breakdown',
        }

    @standard_call
    def __call__(self, rhs, x0=None, conf=None, eps_a=None, eps_r=None,
                 i_max=None, mtx=None, status=None, context=None, **kwargs):
        solver_kwargs = self.build_solver_kwargs(conf)

        eps_a = get_default(eps_a, self.conf.eps_a)
        eps_r = get_default(eps_r, self.conf.eps_r)
        i_max = get_default(i_max, self.conf.i_max)

        setup_precond = get_default(kwargs.get('setup_precond', None),
                                    self.conf.setup_precond)
        callback = get_default(kwargs.get('callback', lambda sol: None),
                               self.conf.callback)

        self.iter = 0
        def iter_callback(sol):
            self.iter += 1
            msg = '%s: iteration %d' % (self.conf.name, self.iter)
            if conf.verbose > 2:
                if conf.method not in self._callbacks_res:
                    res = mtx * sol - rhs

                else:
                    res = sol

                rnorm = nm.linalg.norm(res)
                msg += ': |Ax-b| = %e' % rnorm
            output(msg, verbose=conf.verbose > 1)

            # Call an optional user-defined callback.
            callback(sol)

        precond = setup_precond(mtx, context)

        if conf.method == 'qmr':
            prec_args = {'M1' : precond, 'M2' : precond}

        else:
            prec_args = {'M' : precond}

        solver_kwargs.update(prec_args)

        try:
            sol, info = self.solver(mtx, rhs, x0=x0, atol=eps_a, tol=eps_r,
                                    maxiter=i_max, callback=iter_callback,
                                    **solver_kwargs)
        except TypeError:
            sol, info = self.solver(mtx, rhs, x0=x0, tol=eps_r,
                                    maxiter=i_max, callback=iter_callback,
                                    **solver_kwargs)

        output('%s: %s convergence: %s (%s, %d iterations)'
               % (self.conf.name, self.conf.method,
                  info, self.converged_reasons[nm.sign(info)], self.iter),
               verbose=conf.verbose)

        return sol, self.iter

class PyAMGSolver(LinearSolver):
    """
    Interface to PyAMG solvers.

    The `method` parameter can be one of: 'smoothed_aggregation_solver',
    'ruge_stuben_solver'. The `accel` parameter specifies the Krylov
    solver name, that is used as an accelerator for the multigrid solver.
    """
    name = 'ls.pyamg'

    _parameters = [
        ('method', 'str', 'smoothed_aggregation_solver', False,
         'The actual solver to use.'),
        ('accel', 'str', None, False,
         'The accelerator.'),
        ('callback', 'callable', None, False,
         """User-supplied function to call after each iteration. It is called
            as callback(xk), where xk is the current solution vector, except
            the gmres accelerator, where the argument is the residual norm.
         """),
        ('i_max', 'int', 100, False,
         'The maximum number of iterations.'),
        ('eps_r', 'float', 1e-8, False,
         'The relative tolerance for the residual.'),
        ('force_reuse', 'bool', False, False,
         """If True, skip the check whether the MG solver object corresponds
            to the `mtx` argument: it is always reused."""),
        ('*', '*', None, False,
         """Additional parameters supported by the method. Use the 'method:'
            prefix for arguments of the method construction function
            (e.g. 'method:max_levels' : 5), and the 'solve:' prefix for
            the subsequent solver call."""),
    ]

    # All iterative solvers in pyamg.krylov pass a solution vector into
    # a callback except those below, that take a residual vector norm.
    _callbacks_res = ['gmres']

    def __init__(self, conf, **kwargs):
        try:
            import pyamg
        except ImportError:
            msg =  'cannot import pyamg!'
            raise ImportError(msg)

        LinearSolver.__init__(self, conf, mg=None, **kwargs)

        try:
            solver = getattr(pyamg, self.conf.method)
        except AttributeError:
            output('pyamg.%s does not exist!' % self.conf.method)
            output('using pyamg.smoothed_aggregation_solver instead')
            solver = pyamg.smoothed_aggregation_solver
        self.solver = solver

    @standard_call
    def __call__(self, rhs, x0=None, conf=None, eps_a=None, eps_r=None,
                 i_max=None, mtx=None, status=None, **kwargs):
        solver_kwargs = self.build_solver_kwargs(conf)

        eps_r = 
get_default(eps_r, self.conf.eps_r)