from numpy import hstack, eye, ndarray, where, allclose
from numpy.linalg import inv, LinAlgError, matrix_rank
from numpy.random import standard_normal, randint


def inv_by_gauss_elimination(mat):
    """Calculate inverse matrix by gauss elimination

    Parameters
    ----------
    mat : ndarray
        matrix

    Returns
    -------
    mat_inv : ndarray
        inverse matrix
    null_dim : int
        dimension of null space

    """
    n_row, n_col = mat.shape

    if n_row != n_col:
        raise IndexError("The rows and columns are different sizes.")

    data = hstack([mat, eye(n_row)])

    null_dim = 0
    # calculate row echelon form
    for idx_row in range(n_row):
        idx_pivot_candidate = where(data[idx_row:, idx_row] != 0.0)[0]
        if len(idx_pivot_candidate) > 0:
            idx_pivot = idx_pivot_candidate[0] + idx_row
        else:
            null_dim += 1

        if idx_pivot != idx_row:
            tmp = data[idx_row, :].copy()
            data[idx_row, :] = data[idx_pivot, :]
            data[idx_pivot, :] = tmp

        if data[idx_row, idx_row] != 0.0:
            data[idx_row, :] /= data[idx_row, idx_row]

        for idx_row_dst in range(idx_row + 1, n_row):
            data[idx_row_dst, :] -= (data[idx_row, :]
                                     * data[idx_row_dst, idx_row])

    # calculate inverse matrix
    for idx_row in range(n_row - 1, 0, -1):
        for idx_row_dst in range(idx_row - 1, -1, -1):
            data[idx_row_dst, :] -= (data[idx_row, :]
                                     * data[idx_row_dst, idx_row])

    mat_inv = data[-n_row:, -n_col:]

    return mat_inv, null_dim

n_row = 128

for idx in range(1):
    mat = standard_normal((n_row, n_row))

    # Test Inverse
    try:
        mat_inv_exact = inv(mat)
        mat_inv, null_dim = inv_by_gauss_elimination(mat)
        is_same = allclose(mat_inv_exact, mat_inv)
    
        if is_same is True:
            print("same")
        else:
            print("=====Error=====")
            break
    except LinAlgError:
        null_dim_exact = n_row - matrix_rank(mat)
        mat_inv, null_dim = inv_by_gauss_elimination(mat)
        if null_dim_exact == null_dim:
            print("null dim same")
        else:
            print("=====null dim Error=====")
            break