"""ch08_matrix_inverse.py

Compute the inverse of a small matrix by Doolittle LU factorization.

Origin:
    Python translation of ch08_matrix_inverse.f90 by Alexander Godunov,
    based on Inverse.f90. The no-pivot LU factorization and repeated
    triangular solves match the Fortran version.

Example:
    A = [[3, 2, 4], [2, -3, 1], [1, 1, 2]].
    Solve A*x=e_k for each unit vector e_k, then put the solutions
    into the columns of A^{-1}.

Explicit inversion is seldom needed to solve A*x=b. If only x is
needed, solve the system directly. This example illustrates reuse
of the LU factors for multiple right-hand sides.
"""


def main():
    """Set up the example and check A times its computed inverse."""
    a = [[3.0,  2.0, 4.0],
         [2.0, -3.0, 1.0],
         [1.0,  1.0, 2.0]]
    n = len(a)

    print("Matrix inverse from Doolittle LU factorization")
    print("\nMatrix A:")
    for row in a:
        print("".join(f"{value:12.6f}" for value in row))

    # Factorization changes its matrix. Preserve a for verification.
    a_work = [row.copy() for row in a]
    inverse, info = inverse_lu(a_work)
    if info != 0:
        if info < 0:
            raise ValueError("Invalid matrix size or shape")
        raise ArithmeticError(f"Zero pivot at step {info}; pivoting is needed")

    print("\nInverse matrix A^{-1}:")
    for row in inverse:
        print("".join(f"{value:12.6f}" for value in row))

    # Multiply A and A^{-1} explicitly and compare with the identity.
    max_error = 0.0
    for i in range(n):
        for j in range(n):
            element = 0.0
            for k in range(n):
                element += a[i][k] * inverse[k][j]
            if i == j:
                element -= 1.0
            max_error = max(max_error, abs(element))
    print(f"\nMaximum |A*A^{{-1}}-I| = {max_error:.4e}")


def inverse_lu(a):
    """Form A^{-1} using one no-pivot LU factorization.

    Input/output:
        a - square coefficient matrix, changed during elimination

    Output:
        c    - computed inverse if info = 0
        info - 0 success; -1 invalid shape;
               positive one-based step of an exact zero pivot

    Factor A=L*U with unit diagonal in L. For each column e_k of the
    identity, solve L*d=e_k and U*x=d, then store x in column k of c.
    A nonsingular matrix may require row pivoting; this routine does
    not perform it. A small identity-check error alone does not show
    that an ill-conditioned inverse is accurate.
    """
    n = len(a)
    if n < 1 or any(len(row) != n for row in a):
        return [], -1
    c = [[0.0] * n for _ in range(n)]
    lower = [[0.0] * n for _ in range(n)]
    upper = [[0.0] * n for _ in range(n)]

    # Record each elimination multiplier and change the trailing part.
    for k in range(n-1):
        if a[k][k] == 0.0:
            return c, k + 1
        for i in range(k+1, n):
            coeff = a[i][k] / a[k][k]
            lower[i][k] = coeff
            for j in range(k+1, n):
                a[i][j] -= coeff * a[k][j]
    if a[n-1][n-1] == 0.0:
        return c, n

    # Complete L with ones; U is the upper triangle after elimination.
    for i in range(n):
        lower[i][i] = 1.0
    for j in range(n):
        for i in range(j+1):
            upper[i][j] = a[i][j]

    # Solve once for each identity column, reusing the LU factors.
    for k in range(n):
        rhs = [0.0] * n
        rhs[k] = 1.0
        d = [0.0] * n
        x = [0.0] * n

        # Forward substitution: L*d=e_k.
        for i in range(n):
            d[i] = rhs[i]
            for j in range(i):
                d[i] -= lower[i][j] * d[j]

        # Back substitution: U*x=d.
        for i in range(n-1, -1, -1):
            x[i] = d[i]
            for j in range(i+1, n):
                x[i] -= upper[i][j] * x[j]
            x[i] /= upper[i][i]

        # Solution of A*x=e_k is column k of the inverse.
        for i in range(n):
            c[i][k] = x[i]
    return c, 0


if __name__ == "__main__":
    main()