import math
import operator
import sets
from types import *

'''A module providing some basic linear algebra operations.

The functions cover the standard basic operations on vectors and matrices.
If you are looking for an efficient and complete linear algebra module, 
this isn't it. This module is intended just as a simple (no install, 
no dependencies) way of double-checking some of the more tedious 
linear algebra stuff as I worked my way through a linear algebra text.

The tests embedded in the docstring of each function shows basic usage of the
function, and all tests can be run by executing:

python linearalgebra.py

No output means success. If you want the output to display what is being
tested, run:

python linearalgebra.py -v

Some useful math functions and constants are re-exported from the standard
Python math module, as a convenience for interactive scripting. These include
e, pi, sqrt, cos, sin, tan, acos, asin, and atan.

Conventions for function parameters:
  * u and v always refer to a vector;
  * s and c always refers to a scalar;
  * A, B, and C always refer to a matrix;
  * m, n, r always refer to a scalar that is the dimension of a matrix
    (e.g., m = 3, n = 2 for a 3 x 2 or m x n matrix);
  * x may be either a vector or a scalar;
  * any of these may be postfixed by an integer, and is of the same type
    as the unpostfixed variable name.'''



# We re-export these from the math module for convenience.
degrees = math.degrees
radians = math.radians
e = math.e
pi = math.pi
sqrt = math.sqrt
cos = math.cos
sin = math.sin
tan = math.tan
acos = math.acos
asin = math.asin
atan = math.atan

# The value we use for determining whether a float is equal to an
# int: if the difference is less than epsilon, we deem them equal.
epsilon  = 0.000000000000001

class LAValueError(ValueError):
    pass
class InvalidVectorException(LAValueError):
    pass
class IncompatibleVectorException(LAValueError):
    pass
class InvalidMatrixException(LAValueError):
    pass
class IncompatibleMatrixException(LAValueError):
    pass

## First, some decorators that are used for verifying arguments, such as
## that a vector is a vector or that two matrices have dimensions that
## allow them to be multiplied.

# This decorator verifies args to the other decorators, all of which
# take integer arguments indicating the index of the arg to be checked
# before the function that is decorated is called.
def __intArgs(func):
    ''' Verifies that the indices that were passed to the given decorator for
    the given function are all ints. '''
    def wrapper(*args):
        for arg in args:
            if not isinstance(arg, int):
                msg = "Arg for '%s' decorator should be an int: %s"
                raise LAValueError(msg % (func.__name__, repr(arg)))
        return func(*args)
    wrapper.__doc__  = func.__doc__
    wrapper.__name__ = func.__name__
    return wrapper

# Verify that two vectors have the same size.
@__intArgs
def __samesizeV(index1, index2):
    ''' Decorates a function such that an exception is raised if
    the vector arg at position index1 does not have the same size as
    the vector arg at position index2, starting at 0 for the first arg.
    If the vectors have the same size, then the function is executed
    as it would have been.'''
    def decorator(func):
        def _decorator(*args):
            vec_u = args[index1]
            vec_v = args[index2]
            if not isV(vec_u):
                raise InvalidVectorException("First vector arg is not a vector.")
            elif not isV(vec_v):
                raise InvalidVectorException("Second vector arg is not a vector.")
            len_u = len(vec_u)
            len_v = len(vec_v)
            if len_u != len_v:
                raise IncompatibleVectorException(
                    "Vectors must be of equal size, not %d and %d."
                    % (len_u, len_v))
            return func(*args)
        _decorator.__doc__  = func.__doc__
        _decorator.__name__ = func.__name__
        return _decorator
    return decorator

# Verify that two matrices have the same number of rows and columns.
@__intArgs
def __samesizeM(index1, index2):
    ''' Decorates a function such that an exception is raised if
    the matrix arg at position index1 does not have the same dimensions
    as the matrix arg at position index2, starting at 0 for the first arg.
    If the matrices have the same number of rows and columns, then the
    function is executed as it would have been.'''
    def decorator(func):
        def _decorator(*args):
            (m1, n1) = dimensionsM(args[index1])
            (m2, n2) = dimensionsM(args[index2])
            if m1 != m2 or n1 != n2:
                msg = "Matrices must have same dimensions. Matrix 1 is %d x %d, but matrix 2 is %d x %d."
                raise IncompatibleMatrixException(msg % (m1, n1, m2, n2))
            return func(*args)
        _decorator.__doc__  = func.__doc__
        _decorator.__name__ = func.__name__
        return _decorator
    return decorator

# Verify that one or more args of a function is a scalar.
@__intArgs
def __verifyS(*argIndices):
    ''' Decorates a function such that an exception is raised if the scalar arg at
    any of the supplied indices is not actually a vector, and otherwise executes
    the function as usual.'''
    def decorator(func):
        def _decorator(*args):
            for argIndex in argIndices:
                arg = args[argIndex]
                if not isS(arg):
                    raise LAValueError("Arg at index %s is not a scalar: %s" % (argIndex, arg))
            return func(*args)
        _decorator.__doc__  = func.__doc__
        _decorator.__name__ = func.__name__
        return _decorator
    return decorator

# Verify that one or more args is a vector.
@__intArgs
def __verifyV(*argIndices):
    ''' Decorates a function such that an exception is raised if the vector
    arg at any of the supplied indices is not actually a vector, and otherwise
    executes the function as usual. '''
    def decorator(func):
        def _decorator(*args):
            for argIndex in argIndices:
                u = args[argIndex]
                if not isV(u):
                    raise InvalidVectorException("Arg should be a vector: %s" % repr(u))
            return func(*args)
        _decorator.__doc__  = func.__doc__
        _decorator.__name__ = func.__name__
        return _decorator
    return decorator

# Verify that one or more args is a matrix.
@__intArgs
def __verifyM(*argIndices):
    ''' Decorates a function such that an exception is raised if the matrix
    arg at any of the supplied indices is not a matrix, and otherwise executes
    the function as usual. '''
    def decorator(func):
        def _decorator(*args):
            for argIndex in argIndices:
                A = args[argIndex]
                if not isM(A):
                    raise InvalidMatrixException("Arg is not a matrix: %s" % repr(A))
            return func(*args)
        _decorator.__doc__  = func.__doc__
        _decorator.__name__ = func.__name__
        return _decorator
    return decorator

# Verify that the number of columns of the first matrix equals number of rows
# of the second matrix.
@__intArgs
def __columnRowCompatibleM(index1, index2):
    ''' Decorates a function such that an exception is raised if
    the matrix arg at position index1 is not column-row-compatible
    with the matrix arg at position index2, starting at 0 for the first arg.

    Matrices A (of size a x b) and B (of size c x d) are column-row-compatible
    if and only if b == c; i.e., if under variable name reassignment,
    A has size m x n and B has size n x r.'''
    def decorator(func):
        def _decorator(*args):
            (m1, n1) = dimensionsM(args[index1])
            (m2, n2) = dimensionsM(args[index2])
            if n1 != m2:
                msg = "Matrices are not size-compatible. Matrix 1 has %s column(s), but matrix 2 has %s row(s) instead of %s"
                raise IncompatibleMatrixException(msg % (n1, m2, n1))
            return func(*args)
        _decorator.__doc__  = func.__doc__
        _decorator.__name__ = func.__name__
        return _decorator
    return decorator

## Actual functions defined below here.

def isM(A):
    ''' Determine whether the given arg is a matrix, which must
    be a rectangular array of numbers, represented as a list/tuple of
    list/tuple elements, each of which is a vector of equal size.

    >>> isM([])
    True
    >>> isM([[]])
    True
    >>> isM([1,2])
    False
    >>> isM([[1,2], (3,4)])
    True
    >>> isM([[1,2], [3,4,5]])
    False
    >>> isM([[1,2], ['0', 1]])
    False
    >>> isM([[.5, -1, 9], [4, 3, 19], [3, 6, 8], [0, 0, 1]])
    True
    '''
    if not isinstance(A, list) and not isinstance(A, tuple):
        return False
    m = len(A)
    if m == 0:
        return True
    n = None
    for row in A:
        if not isinstance(row, list) and not isinstance(row, tuple):
            return False
        if not isV(row):
            return False
        if n is None:
            n = len(row)
        elif n != len(row):
            return False
    return True

@__verifyM(0)
def dimensionsM(A):
    ''' Calculates the dimensions of the given matrix, returning a pair consisting
    of the number of rows and the number of columns, raising an exception if any of
    the rows of the matrix have a different number of elements.

    >>> dimensionsM([])
    (0, 0)
    >>> dimensionsM([[1,2,3]])
    (1, 3)
    >>> dimensionsM([[1], [2], [3]])
    (3, 1)
    >>> dimensionsM([[1, 2], [1, 2, 3]])
    Traceback (most recent call last):
      ...
    InvalidMatrixException: Arg is not a matrix: [[1, 2], [1, 2, 3]]
    '''
    m = len(A)
    if m == 0:
        return (0,0)
    n = len(A[0])
    for (i, row) in enumerate(A):
        l = len(row)
        if n != l:
            msg = "Matrix has rows of different length: row 0 has %d element(s) and row %d has %d element(s)." % (n, i, l)
            raise InvalidMatrixException(msg)
    return (m, n)

@__verifyM(0)
def isSquareM(A):
    ''' Determine whether the given matrix is a square matrix, with the same number of
    rows and columns.

    >>> isSquareM([])
    True
    >>> isSquareM([[]])
    False
    >>> isSquareM([[1,1], [2,2], [3,3]])
    False
    >>> isSquareM([[1,2,3], [4,5,6], [7,8,9]])
    True
    >>> isSquareM([[1], [1,2]])
    Traceback (most recent call last):
      ...
    InvalidMatrixException: Arg is not a matrix: [[1], [1, 2]]
    '''
    (m, n) = dimensionsM(A)
    return m == n

@__verifyM(0)
def isDiagonalM(A):
    ''' Determine whether the given matrix is a diagonal matrix -- i.e., a square matrix with
    all nondiagonal entries equal to 0.

    >>> isDiagonalM([])
    True
    >>> isDiagonalM([[0]])
    True
    >>> isDiagonalM([[0, 0], [0, 0], [0, 0]])
    False
    >>> isDiagonalM([[1, 0], [0, 2]])
    True
    >>> isDiagonalM([[1, 0], [1, 0]])
    False
    '''
    if not isSquareM(A):
        return False
    for (i, row) in enumerate(A):
        for (j, elem) in enumerate(row):
            if i != j and elem != 0:
                return False
    return True

@__verifyM(0)
def isScalarM(A):
    ''' Determine whether the given matrix is a scalar matrix -- i.e., a diagonal matrix
    with all diagonal entries equal to each other.

    >>> isScalarM([])
    True
    >>> isScalarM([[2]])
    True
    >>> isScalarM([[3, 0], [0, 3]])
    True
    >>> isScalarM([[3, 0], [1, 3]])
    False
    '''
    if not isSquareM(A):
        return False
    elif len(A) == 0:
        return True
    diagVal = A[0][0]
    for (i, row) in enumerate(A):
        for (j, elem) in enumerate(row):
            if (i != j and elem != 0) or (i == j and elem != diagVal):
                return False
    return True

@__verifyM(0)
def isIdentityM(A):
    ''' Determine whether the given matrix is an identity matrix -- i.e., a scalar matrix
    with all diagonal entries equal to 1.

    >>> isIdentityM([])
    False
    >>> isIdentityM([[0]])
    False
    >>> isIdentityM([[1]])
    True
    >>> isIdentityM([[1, 1], [1, 1]])
    False
    >>> isIdentityM([[1, 0], [0, 1]])
    True
    >>> isIdentityM([[1, 0], [0, 1], [0, 0]])
    False
    >>> isIdentityM([[2, 0], [0, 2]])
    False
    '''
    if not isSquareM(A):
        return False
    elif len(A) == 0:
        return False
    for (i, row) in enumerate(A):
        for (j, elem) in enumerate(row):
            if (i != j and elem != 0) or (i == j and elem != 1):
                return False
    return True

@__verifyM(0)
def isSymmetricM(A):
    ''' Determine whether the given matrix is symmetric -- i.e., whether it is equal
    to its own transpose.

    >>> isSymmetricM([[1, 0], [0, 1]])
    True
    >>> isSymmetricM([[0, 1], [1, 0]])
    True
    >>> isSymmetricM([[1, 2], [2, 0]])
    True
    >>> isSymmetricM([[1, 2], [1, 2]])
    False
    >>> isSymmetricM([[1, 3, 2], [3, 5, 0], [2, 0, 4]])
    True
    '''
    return areEqualM(A, transposeM(A))

def __pairwiseM(A, B, func):
    ''' Helper function for plusM and minusM. '''
    (m, n) = dimensionsM(A)
    C = [None] * m
    for i in range(m):
        row = [None] * n
        for j in range(n):
            row[j] = func(A[i][j], B[i][j])
        C[i] = row
    return C

@__verifyM(0, 1)
@__samesizeM(0, 1)
def plusM(A, B):
    ''' Add the given matrices, which must have the same number of rows and columns.

    >>> plusM([[1, 2], [3, 4]], [[4, 3], [2, 1]])
    [[5, 5], [5, 5]]
    >>> plusM([[1,2]], [[3, 4]]), plusM([[1,2]], [[3, 4]])[0] == plusV([1,2], [3,4])
    ([[4, 6]], True)
    >>> plusM([[1]], [[1, 2]])
    Traceback (most recent call last):
      ...
    IncompatibleMatrixException: Matrices must have same dimensions. Matrix 1 is 1 x 1, but matrix 2 is 1 x 2.
    '''
    return __pairwiseM(A, B, operator.add)

@__verifyM(0, 1)
@__samesizeM(0, 1)
def minusM(A, B):
    ''' Subtract the second matrix from the first, both of which must have the
    same number of rows and columns.

    >>> minusM([[3, 4], [1, 2]], [[1, 1], [1, 1]])
    [[2, 3], [0, 1]]
    >>> minusM([[1,2]], [[3, 4]]), minusM([[1,2]], [[3, 4]])[0] == minusV([1,2], [3,4])
    ([[-2, -2]], True)
    >>> minusM([[1]], [[1, 2]])
    Traceback (most recent call last):
      ...
    IncompatibleMatrixException: Matrices must have same dimensions. Matrix 1 is 1 x 1, but matrix 2 is 1 x 2.
    '''
    return __pairwiseM(A, B, operator.sub)

@__verifyS(0)
@__verifyM(1)
def timesM(s, A):
    ''' Multiply the matrix A by the scalar s.

    >>> timesM(2, [[0,1], [1, 0]])
    [[0, 2], [2, 0]]
    >>> timesM(-1, [[1,2,3,4]])
    [[-1, -2, -3, -4]]
    >>> timesM(3, [[]])
    [[]]
    >>> timesM(1, [])
    []
    >>> timesM(3, [[1,2,3], [4,5,6], [7,8,9], [10, 11, 12]])
    [[3, 6, 9], [12, 15, 18], [21, 24, 27], [30, 33, 36]]
    '''
    return [[s*elem for elem in row] for row in A]

@__verifyM(0)
def negativeM(A):
    ''' Return the negative of the given matrix, or the matrix multiplied
    by scalar -1.

    >>> negativeM([])
    []
    >>> negativeM([[]])
    [[]]
    >>> negativeM(((-1, 2), (2, -3)))
    [[1, -2], [-2, 3]]
    '''
    return timesM(-1, A)

@__verifyM(0)
def transposeM(A):
    ''' Generate the transpose of the given m x n matrix, which is the n x m
    matrix with where each (i,j) element of the original matrix is the (j,i)
    element of the new matrix.

    >>> transposeM([[1, 2], [3, 4], [5, 6]])
    [[1, 3, 5], [2, 4, 6]]
    >>> transposeM([])
    []
    >>> transposeM([[]])
    []
    >>> transposeM([[1,2,3], [4,5,6], [7,8,9]])
    [[1, 4, 7], [2, 5, 8], [3, 6, 9]]
    >>> transposeM(transposeM([[1,2,3], [4,5,6], [7,8,9]]))
    [[1, 2, 3], [4, 5, 6], [7, 8, 9]]
    >>> transposeM([[1,2,3]])
    [[1], [2], [3]]
    '''
    (m, n) = dimensionsM(A)
    C = [[None] * m for i in range(n)]
    for i in range(m):
        for j in range(n):
            C[j][i] = A[i][j]
    return C

@__columnRowCompatibleM(0, 1)
@__verifyM(0, 1)
def productM(A, B):
    ''' Calculate the product of matrices A and B, where A must have the
    same number of columns and B has rows. The result for A of size m x n
    and B of size n x r is an m x r matrix.

    >>> productM([[1,1], [1,1]], [[1,1], [1,1]])
    [[2, 2], [2, 2]]
    >>> productM([], [])
    []
    >>> productM([[1, 3, -1], [-2, -1, 1]], [[-4, 0, 3, -1], [5, -2, -1, 1], [-1, 2, 0, 6]])
    [[12, -8, 0, -4], [2, 4, -5, 7]]
    '''
    (m, n) = dimensionsM(A)
    (n, r) = dimensionsM(B)
    C = [[None] * r for i in range(m)]
    for i in range(m):
        for j in range(r):
            C[i][j] = productV(A[i], [row[j] for row in B])
    return C

def areEqualS(s1, s2):
    ''' Determine whether the two scalars are equal to accuracy of epsilon.

    >>> areEqualS(0.9000000000000002, 0.9000000000000001)
    True
    >>> areEqualS(1, 1)
    True
    '''
    return abs(s1-s2) < epsilon

@__samesizeV(0, 1)
def areEqualV(u, v):
    ''' Determine whether the two vectors are equal by determining if the
    pairwise components are equal to accuracy of epsilon.

    >>> areEqualV([2], [2])
    True
    >>> areEqualV([1.0, 0.9000000000000002], [1, 0.9000000000000001])
    True
    '''
    return __allTrue(map(lambda (x,y): areEqualS(x,y), zip(u, v)))

@__verifyM(0, 1)
def areEqualM(A, B):
    ''' Determine whether the two matrices are equal by determining if
    the corresponding components are equal to accuracy of epsilon.

    >>> areEqualM([], [])
    True
    >>> areEqualM([[1, 2, 3], [3.9000000000000001, 2, 1]], [[1, 2, 3], [3.9, 2, 1]])
    True
    '''
    (m1, n1) = dimensionsM(A)
    (m2, n2) = dimensionsM(B)
    if m1 != m2 or n1 != n2:
        return False
    for i in range(m1):
        for j in range(n1):
            if not areEqualS(A[i][j], B[i][j]):
                return False
    return True

def zeroV(x):
    ''' Create a zero vector with same size as the given vector, or
    having as many elements as the given int.

    >>> zeroV(3)
    [0, 0, 0]
    >>> zeroV([1,2,3,4])
    [0, 0, 0, 0]
    >>> zeroV([])
    []
    '''
    if isinstance(x, int): return [0] * x
    elif not isinstance(x, list) and not isinstance(x, tuple):
        raise LAValueError("Arg should be a vector or int, not: %s" % repr(x))
    else: return [0 for n in x]

@__verifyV(0)
def unitV(v):
    ''' Create a unit vector in direction of v.
    '''
    return timesV(1/normV(v), v)

@__verifyV(1)
def timesV(s, v):
    ''' Multiplies the vector by the given scalar.

    >>> timesV(3, [1,2])
    [3, 6]
    >>> timesV(5, [4,-2,7])
    [20, -10, 35]
    >>> timesV(0, [1,2,3])
    [0, 0, 0]
     '''
    return map(lambda x: s * x, v)

@__verifyV(0, 1)
@__samesizeV(0, 1)
def plusV(u, v):
    ''' Calculate the vector resulting from the addition of
    the two given vectors.

    >>> plusV([1,2], [3,4])
    [4, 6]
    >>> plusV([0.5, 1], [2.5, -1])
    [3.0, 0]
    '''
    return map(lambda (x, y): x + y, zip(u, v))

@__verifyV(0, 1)
@__samesizeV(0, 1)
def minusV(u, v):
    ''' Subtract the second vector from the first vector.

    >>> minusV([2], [-3])
    [5]
    >>> minusV([1,2,3], [2,3,4])
    [-1, -1, -1]
    '''
    return map(lambda (x, y): x - y, zip(u, v))

@__verifyV(0, 1)
@__samesizeV(0, 1)
def productV(u, v):
    ''' Calculate the dot product of the given vectors, which is defined
    as the sum of the pairwise products of the vectors.

    >>> productV([0,0], [0,1])
    0
    >>> productV([1,2,-3], [-3,5,2])
    1
    '''
    return sum(map(lambda (x, y): x * y, zip(u, v)))

@__verifyV(0)
def normV(v):
    ''' Calculate the length, or norm, of the vector, which
    is defined to be the square root of the sum of the squares
    of the components of the vector.

    >>> normV([0,1])
    1.0
    >>> normV([1,1]), areEqualS(normV([1,1]), sqrt(2))
    (1.4142135623730951, True)
    >>> normV([2,-1,3]), areEqualS(normV([2,-1,3]), sqrt(14))
    (3.7416573867739413, True)
    >>> normV([sqrt(2), -1, 1])
    2.0
    '''
    return math.sqrt(sum(map(lambda x: x * x, v)))

@__verifyV(0, 1)
@__samesizeV(0, 1)
def distanceV(u, v):
    ''' Calculate the distance between the two vectors, which
    is defined to be the norm of the difference between the
    vectors.
    >>> distanceV([sqrt(2), 1, -1], [0, 2, -2])
    2.0
    >>> distanceV([0,1], [1,0])
    1.4142135623730951
    >>> distanceV([-1,2], [3,1])
    4.1231056256176606
    '''
    return normV(minusV(u, v))

@__verifyV(0, 1)
@__samesizeV(0, 1)
def areOrthogonalV(u, v):
    ''' Determine whether the two vectors are orthogonal, which
    is defined by whether the dot product is 0.

    >>> areOrthogonalV([0,1], [1,0])
    True
    >>> areOrthogonalV([2,-1,1], [1,-2,-1])
    False
    >>> areOrthogonalV([1,1,-2], [3,1,2])
    True
    '''
    return abs(productV(u, v)) < epsilon

@__verifyV(0, 1)
@__samesizeV(0, 1)
def angleV(u, v):
    ''' Determine the angle between the two vectors, in radians.

    >>> angleV([0,1,1], [1,0,1]), areEqualS(angleV([0,1,1], [1,0,1]), pi/3)
    (1.0471975511965979, True)
    >>> angleV([0.9, 2.1, 1.2], [-4.5, 2.6, -0.8])
    1.5376292299388497
    '''
    return math.acos(productV(u, v) / (normV(u) * normV(v)))

@__verifyV(0, 1)
@__samesizeV(0, 1)
def projectV(u, v):
    ''' Calculate the projection of v onto u.

    >>> projectV([2,1], [-1, 3]), areEqualV(projectV([2,1], [-1, 3]), [2.0/5, 1.0/5])
    ([0.40000000000000002, 0.20000000000000001], True)
    >>> projectV([0,0,1], [1,2,3]), areEqualV(projectV( [0,0,1], [1,2,3]), [0,0,3])
    ([0.0, 0.0, 3.0], True)
    '''
    return timesV(float(productV(u, v)) / productV(u, u), u)

# Don't include complex, because python can be built without support.
__NUMBER_TYPES = sets.Set([int, float, long])

def isV(u):
    ''' Determine whether the arg is a vector, or a list or tuple of numbers.

    >>> isV('')
    False
    >>> isV([])
    True
    >>> isV(())
    True
    >>> isV((1,))
    True
    >>> isV([1])
    True
    >>> isV([1,2,3,4])
    True
    >>> isV((1,2,3))
    True
    >>> isV([1, 0.5, -9L])
    True
    >>> isV([1, 2, "3"])
    False
    >>> isV([[0]])
    False
    >>> isV([[1,2]])
    False
    '''
    if not isinstance(u, list) and not isinstance(u, tuple):
        return False
    for s in u:
        if type(s) not in __NUMBER_TYPES:
            return False
    return True

def isS(s):
    ''' Determine whether the given arg is a scalar, which is either
    a float, int, or long. The integral values 0 and 1 may also be given
    as False and True, respectively.

    >>> isS(0)
    True
    >>> isS(True)
    True
    >>> isS(-9.2)
    True
    >>> isS(100L)
    True
    '''
    return isinstance(s, int) or isinstance(s, float) or isinstance(s, long)

# some utility functions

def __allTrue(lst):
    ''' Determine if all items in list evaluate to true. This is equivalent
    to all(iterable) in Python 2.5.

    >>> __allTrue([])
    True
    >>> __allTrue([1,0])
    False
    >>> __allTrue([True, 1])
    True
    >>> __allTrue([True, True, False])
    False
    '''
    for res in lst:
        if not res: return False
    return True

def __test():
    import doctest
    doctest.testmod()

if __name__ == "__main__":
    __test()