"""
arpack - Scipy module to find a few eigenvectors and eigenvalues of a matrix

Uses ARPACK: http://www.caam.rice.edu/software/ARPACK/

"""
__all___=['eigen','eigen_symmetric']

import _arpack 
import numpy as sb
import warnings

# inspired by iterative.py
# so inspired, in fact, that some of it was copied directly
try:
    False, True
except NameError:
    False, True = 0, 1

_type_conv = {'f':'s', 'd':'d', 'F':'c', 'D':'z'}

class get_matvec:
    methname = 'matvec'
    def __init__(self, obj, *args):
        self.obj = obj
        self.args = args
        if isinstance(obj, sb.matrix):
            self.callfunc = self.type1m
            return
        if isinstance(obj, sb.ndarray):
            self.callfunc = self.type1
            return
        meth = getattr(obj, self.methname, None)
        if not callable(meth):
            raise ValueError, "Object must be an array "\
                  "or have a callable %s attribute." % (self.methname,)

        self.obj = meth
        self.callfunc = self.type2

    def __call__(self, x):
        return self.callfunc(x)

    def type1(self, x):
        return sb.dot(self.obj, x)

    def type1m(self, x):
        return sb.dot(self.obj.A, x)

    def type2(self, x):
        return self.obj(x, *self.args)


def eigen(A, k=6, M=None, ncv=None, which='LM',
          maxiter=None, tol=0, return_eigenvectors=True, v0=None):
    """Return k eigenvalues and eigenvectors of the matrix A.

    Solves A * x[i] = w[i] * x[i], the standard eigenvalue problem for
    w[i] eigenvalues with corresponding eigenvectors x[i].

    Inputs:

    A --  A matrix, array or an object with matvec(x) method to perform
          the matrix vector product A * x.  The sparse matrix formats
          in scipy.sparse are appropriate for A.

    k -- The number of eigenvalue/eigenvectors desired

    M -- (Not implemented)
          A symmetric positive-definite matrix for the generalized
          eigenvalue problem A * x = w * M * x
    v0 -- Initial starting solution (n x 1)

    Outputs:

    w -- An array of k eigenvalues

    v -- An array of k eigenvectors, k[i] is the eigenvector corresponding
         to the eigenvector w[i]

    Optional Inputs:

    ncv -- Number of Lanczos vectors generated, ncv must be greater than k
           and is recommended to be ncv > 2*k

    which -- String specifying which eigenvectors to compute.
             Compute the k eigenvalues of:
             'LM' - largest magnitude.
             'SM' - smallest magnitude.
             'LR' - largest real part.
             'SR' - smallest real part.
             'LI' - largest imaginary part.
             'SI' - smallest imaginary part.

    maxiter --  Maximum number of Arnoldi update iterations allowed

    tol -- Relative accuracy for eigenvalues (stopping criterion)

    return_eigenvectors -- True|False, return eigenvectors 

    """
    try:
        n, ny = A.shape
        n == ny
    except:
        raise AttributeError("matrix is not square")
    if M is not None:
        raise NotImplementedError("generalized eigenproblem not supported yet")
        
    # some defaults
    if ncv is None:
        ncv = 2*k + 1
    ncv = min(ncv, n)
    if maxiter == None:
        maxiter = n*10

    # guess type      
    resid = sb.zeros(n,'f')
    try:
        typ = A.dtype.char
    except AttributeError:
        typ = A.matvec(resid).dtype.char
    if typ not in 'fdFD':
        raise ValueError("matrix type must be 'f', 'd', 'F', or 'D'")

    # some sanity checks
    if k <= 0:
        raise ValueError("k must be positive, k=%d"%k)
    if k == n:
        raise ValueError("k must be less than rank(A), k=%d"%k)
    if maxiter <= 0:
        raise ValueError("maxiter must be positive, maxiter=%d"%maxiter)
    whiches = ['LM','SM','LR','SR','LI','SI']
    if which not in whiches:
        raise ValueError("which must be one of %s"%' '.join(whiches))
    if ncv > n or ncv < k:
        raise ValueError("ncv must be k<=ncv<=n, ncv=%s"%ncv)

    # assign solver and postprocessor
    ltr = _type_conv[typ]
    eigsolver = _arpack.__dict__[ltr+'naupd']
    eigextract = _arpack.__dict__[ltr+'neupd']
    matvec = get_matvec(A)

    v = sb.zeros((n, ncv), typ) # holds Ritz vectors
    if v0 == None:
        resid = sb.zeros(n, typ) # residual
        info = 0
    else: # starting vector is given
        nn, kk = v0.shape
        if nn != n:
            raise ValueError("starting vector must be: (%d, 1), got: (%d, %d)" %(n, nn, kk))
        resid = v0[:,0].astype(typ)
        info = 1
        
    workd = sb.zeros(3*n, typ) # workspace
    workl = sb.zeros(3*ncv*ncv+6*ncv, typ) # workspace
    iparam = sb.zeros(11, 'int') # problem parameters
    ipntr = sb.zeros(14, 'int') # pointers into workspaces
    
    ido = 0

    if typ in 'FD':
        rwork = sb.zeros(ncv, typ.lower())

    # only supported mode is 1: Ax=lx
    ishfts = 1
    mode1 = 1 
    bmat = 'I'
    iparam[0] = ishfts
    iparam[2] = maxiter
    iparam[6] = mode1
    
    while True:
        if typ in 'fd':
            ido,resid,v,iparam,ipntr,info =\
                eigsolver(ido,bmat,which,k,tol,resid,v,iparam,ipntr,
                          workd,workl,info)
        else:
            ido,resid,v,iparam,ipntr,info =\
                eigsolver(ido,bmat,which,k,tol,resid,v,iparam,ipntr,
                          workd,workl,rwork,info)

        if (ido == -1 or ido == 1):
            # compute y = A * x
            xslice = slice(ipntr[0] - 1, ipntr[0] - 1 + n)
            yslice = slice(ipntr[1] - 1, ipntr[1] - 1 + n)
            workd[yslice] = matvec(workd[xslice])
        else: # done
            break

    if  info < -1 :
        raise RuntimeError("Error info=%d in arpack"%info)
        return None
    if info == -1:
        warnings.warn("Maximum number of iterations taken: %s"%iparam[2])
#    if iparam[3] != k:
#        warnings.warn("Only %s eigenvalues converged"%iparam[3])


    # now extract eigenvalues and (optionally) eigenvectors        
    rvec = return_eigenvectors
    ierr = 0
    howmny = 'A' # return all eigenvectors
    sselect = sb.zeros(ncv,'int') # unused 
    sigmai = 0.0 # no shifts, not implemented
    sigmar = 0.0 # no shifts, not implemented
    workev = sb.zeros(3*ncv,typ) 

    if typ in 'fd':
        dr=sb.zeros(k+1,typ)
        di=sb.zeros(k+1,typ)
        zr=sb.zeros((n,k+1),typ)
        dr,di,zr,info=\
            eigextract(rvec,howmny,sselect,sigmar,sigmai,workev,
                   bmat,which,k,tol,resid,v,iparam,ipntr,
                   workd,workl,info)

        # make eigenvalues complex
        d=(dr+1.0j*di)[:k]
        # futz with the eigenvectors:
        # complex are stored as real,imaginary in consecutive columns
        z=zr.astype(typ.upper())[:,:k]
        for i in range(k): # fix c.c. pairs
            if di[i] > 0 :
                z[:,i]=zr[:,i]+1.0j*zr[:,i+1]
                z[:,i+1]=z[:,i].conjugate()
                     
    else:
        d,z,info =\
              eigextract(rvec,howmny,sselect,sigmar,workev,
                         bmat,which,k,tol,resid,v,iparam,ipntr,
                         workd,workl,rwork,ierr)

        

    if ierr != 0:
        raise RuntimeError("Error info=%d in arpack"%info)
        return None
    if return_eigenvectors:
        return d,z
    return d


def eigen_symmetric(A,k=6,M=None,ncv=None,which='LM',
                    maxiter=None,tol=0, return_eigenvectors=True, v0=None):
    """ Return k eigenvalues and eigenvectors of the real symmetric matrix A.

    Solves A * x[i] = w[i] * x[i], the standard eigenvalue problem for
    w[i] eigenvalues with corresponding eigenvectors x[i].
    A must be real and symmetric.
    See eigen() for nonsymmetric or complex symmetric (Hermetian) matrices.

    Inputs:

    A --  A symmetric matrix, array or an object with matvec(x) method
          to perform the matrix vector product A * x.
          The sparse matrix formats in scipy.sparse are appropriate for A.

    k -- The number of eigenvalue/eigenvectors desired

    M -- (Not implemented)
          A symmetric positive-definite matrix for the generalized
          eigenvalue problem A * x = w * M * x

    v0 -- Starting vector (n, 1)
    
    Outputs:

    w -- An real array of k eigenvalues 

    v -- An array of k real eigenvectors, k[i] is the eigenvector corresponding
         to the eigenvector w[i]

    Optional Inputs:

    ncv -- Number of Lanczos vectors generated, ncv must be greater than k
           and is recommended to be ncv > 2*k

    which -- String specifying which eigenvectors to compute.
             Compute the k 
             'LA' - largest (algebraic) eigenvalues.
             'SA' - smallest (algebraic) eigenvalues.
             'LM' - largest (in magnitude) eigenvalues.
             'SM' - smallest (in magnitude) eigenvalues. 
             'BE' - eigenvalues, half from each end of the
                    spectrum.  When NEV is odd, compute one more from the
                    high end than from the low end.

    maxiter --  Maximum number of Arnoldi update iterations allowed

    tol -- Relative accuracy for eigenvalues (stopping criterion)

    return_eigenvectors -- True|False, return eigenvectors 

    """
    try:
        n,ny=A.shape
        n==ny
    except:
        raise AttributeError("matrix is not square")
    if M is not None:
        raise NotImplementedError("generalized eigenproblem not supported yet")
    if ncv is None:
        ncv=2*k+1
    ncv=min(ncv,n)
    if maxiter==None:
        maxiter=n*10


    # guess type        
    resid = sb.zeros(n,'f')
    try:
        typ = A.dtype.char
    except AttributeError:
        typ = A.matvec(resid).dtype.char
    if typ not in 'fd':
        raise ValueError("matrix type must be 'f' or 'd'")

    # some sanity checks
    if k <= 0:
        raise ValueError("k must be positive, k=%d"%k)
    if k == n:
        raise ValueError("k must be less than rank(A), k=%d"%k)
    if maxiter <= 0:
        raise ValueError("maxiter must be positive, maxiter=%d"%maxiter)
    whiches=['LM','SM','LA','SA','BE']
    if which not in whiches:
        raise ValueError("which must be one of %s"%' '.join(whiches))
    if ncv > n or ncv < k:
        raise ValueError("ncv must be k<=ncv<=n, ncv=%s"%ncv)

    # assign solver and postprocessor
    ltr = _type_conv[typ]
    eigsolver = _arpack.__dict__[ltr+'saupd']
    eigextract = _arpack.__dict__[ltr+'seupd']
    matvec = get_matvec(A)

    v = sb.zeros((n,ncv),typ)
    if v0 == None:
        resid = sb.zeros(n, typ) # residual
        info = 0
    else: # starting solution is given
        nn, kk = v0.shape
        if nn != n:
            raise ValueError("starting vectors must be: (%d, %d), got: (%d, %d)" %(n, k, nn, kk))
        resid = v0[:,0].astype(typ)
        info = 1
    
    #resid = sb.zeros(n,typ)
    workd = sb.zeros(3*n,typ)
    workl = sb.zeros(ncv*(ncv+8),typ)
    iparam = sb.zeros(11,'int')
    ipntr = sb.zeros(11,'int')
    #info = 0
    ido = 0

    # only supported mode is 1: Ax=lx
    ishfts = 1
    mode1 = 1
    bmat='I'
    iparam[0] = ishfts
    iparam[2] = maxiter
    iparam[6] = mode1
    
    while True:
        ido,resid,v,iparam,ipntr,info =\
            eigsolver(ido,bmat,which,k,tol,resid,v,iparam,ipntr,
               workd,workl,info)
        if (ido == -1 or ido == 1):
            xslice = slice(ipntr[0]-1, ipntr[0]-1+n)
            yslice = slice(ipntr[1]-1, ipntr[1]-1+n)
            workd[yslice]=matvec(workd[xslice])
        else:
            break

    if  info < -1 :
        raise RuntimeError("Error info=%d in arpack"%info)
        return None
    if info == -1:
        warnings.warn("Maximum number of iterations taken: %s"%iparam[2])

    # now extract eigenvalues and (optionally) eigenvectors        
    rvec = return_eigenvectors
    ierr = 0
    howmny = 'A' # return all eigenvectors
    sselect = sb.zeros(ncv,'int') # unused 
    sigma = 0.0 # no shifts, not implemented

    d,z,info =\
             eigextract(rvec,howmny,sselect,sigma,
                        bmat,which, k,tol,resid,v,iparam[0:7],ipntr,
                        workd[0:2*n],workl,ierr)

    if ierr != 0:
        raise RuntimeError("Error info=%d in arpack"%info)
        return None
    if return_eigenvectors:
        return d,z
    return d