import random

## format: m is a matrix of floats. pp_matrix(m) prints out m in
## matrix form
def pp_matrix(m):
  print('\n'.join([''.join(['{:10,.2f}'.format(item) for item in row]) 
                   for row in m]))
  return None

## format: v is a vector of floats. pp_vector(v) prints out v
## vertically
def pp_vector(v):
  pp_matrix ([[elt] for elt in v])
  return None

## format: A is a list, i and j two indices. list_exch(A,i,j) returns
## a copy of A with rows i and j exchanged.
def list_exch(A,i,j):
  if i == j: return A
  return A[:i] + [A[j]] + A[i+1:j] + [A[i]] + A[j+1:]

## format: A is a square matrix (of floats), b a vector (of
## floats). gauss(A,b) returns a vector x such that Ax=b.
def gauss(A,b):
  # make the pivot a float for correct computation
  a,n = float(A[0][0]), len(A)
  
  # base case: when n = 1, ax = b (a != 0 by hypothesis) is solved by
  # x = b/a
  if n == 1: return [b[0]/a]

  # if pivot is 0, exchange first line with one with non-zero pivot
  if a == 0:
    ipiv = 1
    while A[ipiv][0] == 0 : ipiv += 1
    return gauss (list_exch(A,0,ipiv), list_exch(b,0,ipiv))

  # else, do the operation on lines and recursively call gauss on the
  # smaller system
  B = [ row[1:] for row in A[1:] ] # A without first row and column
  c = b[1:] # b without the first row
  for i in range(n-1):
    c[i] -= (b[0]*A[i+1][0]/a)
    for j in range(n-1):
      B[i][j] -= (A[0][j+1]*A[i+1][0]/a)
  # recursive call
  y = gauss(B,c) # By = c
  # reconstruction from y of x such that Ax = b
  x = [b[0]] + y
  for i in range(n-1): x[0] -= (A[0][i+1]*y[i])
  x[0] /= a
  return x
## If we denote C(n) the complexity of gauss(A,b) with matrix A and
## vector b of size n, then C(n) = (n-1)^2 + C(n-1) and C(1) = K
## constant. Hence C(n) = O(sum for i=1 to n-1 of i^2) = O(n^3). In
## practice, it computes the solution to the system in less than a
## second up to n = 200.

## format: A and B are matices. matrix_mult(A,B) returns the product
## AB; A should be of size nxm and B mxp.
def matrix_mult(A,B):
  n,m,p = len(A),len(B),len(B[0])
  prod = [[0 for j in range(p)] for i in range(n)]
  for i in range(n):
    for j in range(p):
      prod[i][j] = sum([A[i][k]*B[k][j] for k in range(m)])
  return prod

## format: p is a list. perm_matrix(p) returns the permutation matrix
## corresponding to the permutation p. In particular, p should be a
## rearrangement of range(len(p)).
def perm_matrix(p):
  n = len(p)
  P = [[0 for j in range(n)] for i in range(n)]
  for i in range(n): P[i][p[i]] = 1
  return P

## format: A,B,C,D are matrices. block_matrix(A,B,C,D) returns the
## matrix whose block-writting is
##
##   -      - 
##  |  A | B |
##  |  ----- |
##  |  C | D |
##   -      -
##
## In particular, A should be of size nxm, B of size nxp, C of size
## qxm and D of size qxp such that the resultiong matrix is of size
## (n+q)x(m+p).
def block_matrix(A,B,C,D):
  n1,n2 = len(A), len(C) # = len(B),len(D)
  M = []
  for i in range(n1): M += [A[i]+B[i]]
  for i in range(n2): M += [C[i]+D[i]]
  return M
  
## format: A is a square matrix. lup(A) returns a tuple (L,U,P) such
## that PA = LU with P permutation matrix, L lower triangular matrix
## with 1 over the diagonal and U upper triangular matrix
def lup(A):
  n = len(A)
  # base case: n=1; in that case L=(1),U=A,P=(1) is a match
  if n == 1:
    return ([[1]],A,[[1]])

  # otherwise start by searching for the first non zero A[k][0]
  ipiv = 0
  while A[ipiv][0] == 0: ipiv += 1
  piv = float(A[ipiv][0])
  # Q perm matrix associated to transposition (0,ipiv)
  Q = perm_matrix(list_exch(range(n),0,ipiv))
  # compute QA and B
  QA = list_exch(A,0,ipiv)
  B = [row[1:] for row in QA[1:]]
  for i in range(n-1):
    for j in range(n-1):
      B[i][j] -= (QA[0][j+1]*QA[i+1][0]/piv)
  # recursive call
  (LL,UU,PP) = lup(B)
  # create L,U,P from LL,UU,PP ; a useful trick to turn a line vector
  # v into a column one is to use [[elt] for elt in v] ; it is used
  # several time in the following
  P = matrix_mult(block_matrix([[1]],[[0 for i in range(n-1)]],[[0] for i in range(n-1)],PP),Q)
  L = block_matrix([[1]],[[0 for i in range(n-1)]],[[sum([PP[i][k]*QA[k+1][0]/piv for k in range(n-1)])] for i in range(n-1)],LL)
  U = block_matrix([[piv]],[QA[0][1:]],[[0] for i in range(n-1)],UU)
  # return the newly created L, U and P
  return (L,U,P)

## format: A is a matrix. matrix_trans(A) returns the transpose of A.
def matrix_trans(A):
  return [[A[j][i] for j in range(len(A))] for i in range(len(A))]  

## format: A is a non singular upper triangular matrix, b a
## vector. solve_upper_triangular(A,b) returns x such that Ax=b.
def solve_upper_triangular(A,b):
  n = len(A)
  x = [0. for i in range(n)]
  for i in range(n)[::-1]:
    x[i] = (b[i] - sum([A[i][j]*x[j] for j in range(i+1,n)]))/float(A[i][i])
  return x

## format: A is a non singular lower triangular matrix, b a
## vector. solve_lower_triangular(A,b) returns x such that Ax=b; it
## uses the trick to turn the lower triangular system into an upper
## one then apply the previous function.
def solve_lower_triangular(A,b):
  return solve_upper_triangular([row[::-1] for row in A[::-1]],b[::-1])[::-1]

## format: L,U,P as before, b is a vector. solve_lup(L,U,P,b) returns
## x such that Ax=b where A is the matrix whose (L,U,P) is the LUP
## decomposition.
def solve_lup(L,U,P,b):
  Pb = [row[0] for row in matrix_mult(P,b)] # Pb = P.b
  Ux = solve_lower_triangular(L,Pb) # L.Ux = P.b
  x = solve_upper_triangular(U,Ux) # L.U.x = P.b, ie P.A.x = P.b
  return x # A.x=b because P invertible

## format: A is an invertible matrix. matrix_inv(A) returns the
## inverse of A by solving the sytem Ax = e for each vector e of the
## canonical basis.
def matrix_inv(A):
  n,IA = len(A),[]
  (L,U,P) = lup(A)
  e = [[1]] + [[0] for i in range(n-1)] # first vector of can. basis
  for i in range(n):
    IA += [solve_lup(L,U,P,e)]
    if i < n-1: e[i],e[i+1] = e[i+1],e[i] # next vector in can. basis
  return matrix_trans(IA)

## format: A is an invertible matrix. matrix_det(A) returns the
## determinant of A, by taking its LUP decomposition: taking
## determinant in PA = LU gives back det(P)det(A) = det(L)det(U)
## because of multiplicativity of det; but det(P) = 1 and det(L) = 1,
## hence det(A) = det(U).
def matrix_det(A):
  n, (_,U,_) = len(A), lup(A)
  r = 1
  for i in range(n): r *= U[i][i]
  return r


## TEST ##

print "### Random matrix and vector ###"
n = 10
A = [[random.uniform(-10,10) for j in range(n)] for i in range(n)]
b = [random.uniform(-100,100) for j in range(n)]
print "A:"
pp_matrix(A)
print "b:"
pp_matrix([b])

# test pivot
print
print "### TEST PIVOT ###"
x = gauss(A,b)
print "x:"
pp_matrix([x])
print "Verif, computing Ax, should equal b:"
pp_matrix(matrix_mult(A,[[elt] for elt in x]))


#test LUP
print
print "### TEST LUP ###"
(L,U,P) = lup(A)
PA,LU = matrix_mult(P,A), matrix_mult(L,U)
print "L:"
pp_matrix(L)
print "U:"
pp_matrix(U)
print "P:"
pp_matrix(P)
print "Verif, computing PA and LU, should be equal:"
print "PA:"
pp_matrix(PA)
print "LU:"
pp_matrix(LU)

#test inverse
print
print "### TEST INVERSE AND DET ###"
invA, detA = matrix_inv(A), matrix_det(A)
print "Inverse of A:"
pp_matrix(invA)
print "Verif, computing product with A, should equal In:"
pp_matrix(matrix_mult(A,invA))
print "det A =", detA, "(no easy verification)"
