# Author : Fabian Jakobs
"""
Functions for solving linear systems.

The library provides simple linear algebra operations which operate
directly on the gsl_vector and gsl_matrix objects. These are intended
for use with "small" systems where simple algorithms are acceptable.

Anyone interested in large systems will want to use the sophisticated
routines found in LAPACK. The Fortran version of LAPACK is recommended
as the standard package for linear algebra. It supports blocked algorithms,
specialized data representations and other optimizations.
"""
import pygsl
from . import _gslwrap
from .permutation import Permutation
import pygsl._numobj as numx
zeros = numx.zeros
array = numx.array
get_typecode = pygsl.get_typecode
array_typed_copy = pygsl.array_typed_copy
Float = pygsl.Float
Complex = pygsl.Complex
#
# LU Decomposition
#

def LU_decomp(A):
    """
    returns (LU, P, signum)
    
    This function factorizes the square matrix A into the LU decomposition
    PA = LU. On output the diagonal and upper triangular part of the return
    matrix contain the matrix U. The lower triangular part of the input matrix
    (excluding the diagonal) contains L. The diagonal elements of L are unity,
    and are not stored.
    
    The permutation matrix P is encoded in the permutation p. The j-th column
su
of the matrix P is given by the k-th column of the identity matrix, where
    k = p_j the j-th element of the permutation vector. The sign of the
    permutation is given by signum. It has the value (-1)^n, where n is the
    number of interchanges in the permutation.
    
    The algorithm used in the decomposition is Gaussian Elimination with
    partial pivoting (Golub & Van Loan, Matrix Computations, Algorithm 3.4.1).
    """
    
    p = Permutation(A.shape[1])
    code = get_typecode(A)
    An = array_typed_copy(A)
    if code == Complex:
        # Now all error flags are turned into python exceptions. So no
        # unpack necessary any longer.
        signum = _gslwrap.gsl_linalg_complex_LU_decomp(An, p)
    elif code == Float:
        signum = _gslwrap.gsl_linalg_LU_decomp(An, p)
    else:
        #print code, Float, Complex
        raise TypeError("LU must be of type Float or Complex")
    return (An, p, signum)


def LU_unpack(LU):
    """
    returns (L,U)
    
    This function splits the matrix LU into the the upper matrix U and
    the lower matrix L. The diagonal of L is the identity.
    """
    code = get_typecode(LU)
    u = zeros(LU.shape, code)
    l = numx.identity(LU.shape[0], code)
    for i in  range(LU.shape[0]):
        u[i, i: ] = LU[i, i:]
        l[i, 0:i] = LU[i, :i]
    return (l, u)

        
def LU_solve(LU, p, b):
    """
    This function solves the system A x = b using the LU decomposition of A
    into (LU, p) given by LU_decomp.
    """
    code = get_typecode(LU)
    x = zeros(LU.shape[1], code)
    if code == Complex:
        _gslwrap.gsl_linalg_complex_LU_solve(LU, p, b, x)
    elif code == Float:
        _gslwrap.gsl_linalg_LU_solve(LU, p, b, x)
    else:
        raise TypeError( "LU must be of type Float or Complex")
    return x


def LU_refine(A, LU, p, b, x):
    """
    returns (x, residual)

    This functions applies an iterative improvement to x, the solution
    of A x = b, using the LU decomposition of A into (LU,p). The initial
    residual r = A x - b is also computed and stored in residual. 
    """
    code = get_typecode(LU)
    raise NotImplementedError("This function is not (yet implemented)")
#    if code == Numeric.Complex:
#        _gslwrap.gsl_linalg_complex_LU_refine(A, LU, p, b, x, residual)
#    elif code == Numeric.Float:
#        _gslwrap.gsl_linalg_LU_refine
#    else:
#        raise TypeError, "LU must be of type Float or Complex"
    

def LU_invert(LU, p):
    """
    returns inverse

    This function computes the inverse of a matrix A from its LU
    decomposition (LU,p), storing the result in the matrix inverse.
    The inverse is computed by solving the system A x = b for each column of
    the identity matrix. It is preferable to avoid direct computation of the
    inverse whenever possible.
    """
    code = get_typecode(LU)
    inverse = zeros(LU.shape, code)
    if code == Complex:
        _gslwrap.gls_linalg_complex_LU_invert(LU, p, inverse)
    elif code == Float:
        _gslwrap.gsl_linalg_LU_invert(LU, p, inverse)
    else:
        raise TypeError("LU must be of type Float or Complex")
    return inverse


def LU_det(LU, signum):
    """
    returns determinant

    This function computes the determinant of a matrix A from its LU
    decomposition, LU. The determinant is computed as the product of
    the diagonal elements of U and the sign of the row permutation signum.
    """
    code = get_typecode(LU)
    if code == Complex:
        return _gslwrap.gls_linalg_complex_LU_det(LU, signum)
    elif code == Float:
        return _gslwrap.gsl_linalg_LU_det(LU, signum)
    else:
        raise TypeError("LU must be of type Float or Complex")


def LU_lndet(LU):
    """
    This function computes the logarithm of the absolute value of the
    determinant of a matrix A, \ln|det(A)|, from its LU decomposition, LU.
    This function may be useful if the direct computation of the determinant
    would overflow or underflow.
    """
    code = get_typecode(LU)
    if code == Complex:
        return _gslwrap.gls_linalg_complex_LU_lndet(LU)
    elif code == Float:
        return _gslwrap.gsl_linalg_LU_lndet(LU)
    else:
        raise TypeError("LU must be of type Float or Complex")


def LU_sgndet(LU, signum):
    """
    This function computes the sign or phase factor of the determinant of a
    matrix A, det(A)/|det(A)|, from its LU decomposition, LU.
    """
    code = get_typecode(LU)
    if code == Complex:
        return _gslwrap.gls_linalg_complex_LU_sgndet(LU, signum)
    elif code == Float:
        return _gslwrap.gsl_linalg_LU_sgndet(LU, signum)
    else:
        raise TypeError("LU must be of type Float or Complex")
    

#
# QR Decomposition
#

def QR_decomp(A):
    """
    returns (QR, tau)
    
    Function: int gsl_linalg_QR_decomp (gsl_matrix * A, gsl_vector * tau)
    This function factorizes the M-by-N matrix A into the QR
    decomposition A = Q R. On output the diagonal and upper triangular
    part of the input matrix contain the matrix R. The vector tau and the
    columns of the lower triangular part of the matrix QR contain the
    Householder coefficients and Householder vectors which encode the
    orthogonal matrix Q. The vector tau must be of length k=\min(M,N).
    The matrix Q is related to these components by, Q = Q_k ... Q_2 Q_1
    where Q_i = I - \tau_i v_i v_i^T and v_i is the Householder vector
    v_i = (0,...,1,A(i+1,i),A(i+2,i),...,A(m,i)).
    This is the same storage scheme as used by LAPACK.

    The algorithm used to perform the decomposition is Householder QR
    (Golub & Van Loan, Matrix Computations, Algorithm 5.2.1).
    """
    code = get_typecode(A)
    qr = array_typed_copy(A, code)
    tau = zeros((min(A.shape),), code)
    _gslwrap.gsl_linalg_QR_decomp(qr, tau)
    return (qr, tau)


def QR_solve(QR, tau, b):
    """
    returns x
    
    This function solves the system A x = b using the QR decomposition of
    A into (QR, tau) given by gsl_linalg_QR_decomp.
    """
    code = get_typecode(b)
    x = zeros(QR.shape[1], code)
    _gslwrap.gsl_linalg_QR_solve(QR, tau, b, x)
    return x


def QR_lssolve(QR, tau, b):
    """
    returns (x, residual)
    
    This function finds the least squares solution to the overdetermined
    system A x = b where the matrix A has more rows than columns. The least
    squares solution minimizes the Euclidean norm of the residual, ||Ax - b||.
    The routine uses the QR decomposition of A into (QR, tau) given by
    gsl_linalg_QR_decomp. The solution is returned in x. The residual is
    computed as a by-product and stored in residual.
    """
    code = get_typecode(QR)
    x = zeros(QR.shape[1], code)
    residual = zeros(QR.shape[0], code)
    _gslwrap.gsl_linalg_QR_lssolve(QR, tau, b, x, residual)
    return (x, residual)


def QR_QTvec(QR, tau, v):
    """
    returns v'

    This function applies the matrix Q^T encoded in the decomposition
    (QR,tau) to the vector v, storing the result Q^T v in v'. The matrix
    multiplication is carried out directly using the encoding of the
    Householder vectors without needing to form the full matrix Q^T.
    """
    vn = array_typed_copy(v)
    _gslwrap.gsl_linalg_QR_QTvec(QR,tau,vn)
    return vn


def QR_Qvec(QR, tau, v):
    """
    returns v'
    
    This function applies the matrix Q encoded in the decomposition
    (QR,tau) to the vector v, storing the result Q v in v'. The matrix
    multiplication is carried out directly using the encoding of the
    Householder vectors without needing to form the full matrix Q.
    """
    vn = array_typed_copy(v)
    _gslwrap.gsl_linalg_QR_Qvec(QR,tau,vn)
    return vn
    

def QR_Rsolve(QR, b):
    """
    returns x
    
    This function solves the triangular system R x = b for x.
    It may be useful if the product b' = Q^T b has already been computed
    using gsl.linalg.QR_QTvec.
    """
    x = zeros(QR.shape[1], get_typecode(b))
    _gslwrap.gsl_linalg_QR_Rsolve(QR, b, x)
    return x
    
    
def QR_unpack(QR, tau):
    """
    returns (Q, R)
    
    This function unpacks the encoded QR decomposition (QR,tau) into the
    matrices Q and R, where Q is M-by-M and R is M-by-N.
    """
    (m, n) = QR.shape
    code = get_typecode(QR)
    q = zeros([m,m], code)
    r = zeros([m,n], code)
    _gslwrap.gsl_linalg_QR_unpack(QR, tau, q, r)
    return (q,r)
    

def QR_QRsolve(Q, R, b):
    """
    returns x
    
    This function solves the system R x = Q^T b for x. It can be used when
    the QR decomposition of a matrix is available in unpacked form as (Q,R).
    """
    x = zeros(R.shape[1], get_typecode(R))