# gaussian elimination with scaled pivoting
#
# M. Zingale (2013-02-25)

from __future__ import print_function
import numpy as np

def gauss_elim(A, b, return_det=0, quiet=0):
    """ perform gaussian elimination with pivoting, solving A x = b.

        A is an NxN matrix, x and b are an N-element vectors.  Note: A
        and b are changed upon exit to be in upper triangular (row
        echelon) form """

    # b is a vector
    if not b.ndim == 1:
        print("ERROR: b should be a vector")
        return None

    N = len(b)

    # A is square, with each dimension of length N
    if not A.shape == (N, N):
        print("ERROR: A should be square with each dim of same length as b")
        return None

    # allocation the solution array
    x = np.zeros((N), dtype=A.dtype)

    # find the scale factors for each row -- this is used when pivoting
    scales = np.max(np.abs(A), 1)

    # keep track of the number of times we swapped rows
    num_row_swap = 0

    if not quiet: print_Ab(A, b)

    # main loop over rows
    for k in range(N):

        # find the pivot row based on the size of column k -- only consider
        # the rows beyond the current row
        row_max = np.argmax(A[k:, k]/scales[k:])
        if k > 0: row_max += k  # we sliced A from k:, correct for total rows

        # swap the row with the largest scaled element in the current column
        # with the current row (pivot) -- do this with b too!
        if not row_max == k:
            A[[k, row_max],:] = A[[row_max, k],:]
            b[[k, row_max]] = b[[row_max, k]]
            if not quiet: print("pivoted")
            num_row_swap += 1

        # do the forward-elimination for all rows below the current
        for i in range(k+1, N):
            coeff = A[i,k]/A[k,k]

            for j in range(k+1, N):
                A[i,j] += -A[k,j]*coeff

            A[i,k] = 0.0
            b[i] += -coeff*b[k]

        if not quiet: print_Ab(A, b)

    # back-substitution

    # last solution is easy
    x[N-1] = b[N-1]/A[N-1,N-1]

    for i in reversed(range(N-1)):
        sum = b[i]
        for j in range(i+1,N):
            sum += -A[i,j]*x[j]
        x[i] = sum/A[i,i]


    # determinant
    det = np.prod(np.diagonal(A))*(-1.0)**num_row_swap

    if not return_det:
        return x
    else:
        return x, det



# for debugging:
# output: np.savetxt("test.out", a, fmt="%5.2f", delimiter="  ")
# convert -font Courier-New-Regular -pointsize 20 text:test.out test.png

def print_Ab(A, b):
    """ printout the matrix A and vector b in a pretty fashion.  We
        don't use the numpy print here, because we want to make them
        side by side"""

    N = len(b)

    open_t = "/"
    close_t = "\\"

    open_b = "\\"
    close_b = "/"

    # numbers take 6 positions + 2 spaces
    a_fmt = " {:6.3f} "
    space = 8*" "

    line = "|" + N*a_fmt + "|" + space + "|" + a_fmt + "|"
    top = open_t + N*space + close_t  + space + open_t + space + close_t
    bottom = open_b + N*space + close_b + space + open_b + space + close_b + "\n"

    print(top)
    for i in range(N):
        out = tuple(A[i,:]) + (b[i],)
        print(line.format(*out))
    print(bottom)