"""ch08_shifted_inverse_iteration.py

Find an eigenpair near a chosen shift by inverse iteration.

Origin:
    Python translation of ch08_shifted_inverse_iteration.f90 by
    Alexander Godunov. Each step solves (A-shift*I)*z=y and then
    normalizes z. Pivoted LU factors are computed once and reused.

Example:
    A = [ 1  2  3; 2  2 -2; 3 -2  4 ], shift = 3.4.
    Its eigenvalues are approximately -2.54138127, 3.54138127,
    and 6.00000000. The nearest eigenvalue to 3.4 is 3.54138127.

A shift at an eigenvalue makes A-shift*I singular. A nearby shift can
be ill-conditioned. The initial vector must have a component along
the desired eigenvector.
"""

from math import isfinite, sqrt
from sys import float_info


def main():
    """Run the three-dimensional example and check its eigenpair."""
    a = [[1.0,  2.0,  3.0],
         [2.0,  2.0, -2.0],
         [3.0, -2.0,  4.0]]
    y = [1.0, 1.0, 1.0]
    shift = 3.4
    tol = 1.0e-10
    max_iter = 100

    print("Shifted inverse iteration for one eigenpair")
    print("\nMatrix A:")
    for row in a:
        print("".join(f"{value:12.6f}" for value in row))
    print("\nInitial vector:" + "".join(f"{value:12.6f}" for value in y))
    print(f"Shift = {shift:12.6f}")
    print(f"Eigenpair residual tolerance = {tol:.4e}")

    y, eigenvalue, iterations, residual, pivot_step, info = (
        shifted_inverse(a, y, shift, tol, max_iter))
    if info == 1:
        raise ValueError("Invalid input: check symmetry, initial vector, and controls")
    if info == 2:
        raise ArithmeticError("A-shift*I is singular or nearly singular "
                              f"at pivot step {pivot_step}; choose a different shift")
    if info == 3:
        raise ArithmeticError("The inverse iteration vector is zero; cannot normalize")

    print(f"\nEigenvalue near shift = {eigenvalue:12.6f}")
    print("Unit eigenvector:")
    for i, value in enumerate(y, start=1):
        print(f"y({i}) = {value:12.6f}")
    print(f"Iterations = {iterations}")
    print(f"Eigenpair residual ||A*y-lambda*y|| = {residual:.4e}")
    if info == 4:
        raise ArithmeticError("Iteration limit reached before the tolerance was met")


def shifted_inverse(a, y, shift, tol, max_iter):
    """Approximate the eigenvalue of a real symmetric A nearest shift.

    Input:
        a        - square symmetric matrix, left unchanged
        y        - nonzero starting vector, left unchanged
        shift    - fixed estimate of the desired eigenvalue
        tol      - positive tolerance for ||A*y-eigenvalue*y||_2
        max_iter - maximum number of inverse iterations

    Output:
        y          - unit vector approximating an eigenvector
        eigenvalue - signed Rayleigh quotient y^T*A*y / (y^T*y)
        iterations - completed inverse iterations
        residual   - Euclidean norm of the eigenpair residual
        pivot_step - one-based small-pivot step if info = 2
        info       - 0 converged; 1 invalid input; 2 singular shifted
                     matrix; 3 zero inverse vector; 4 iteration limit

    Factor A-shift*I once with row pivoting. Each step solves with
    the stored factors and normalizes the solution. Compute the
    Rayleigh quotient and residual from the ORIGINAL matrix A.
    """
    n = len(y)
    eigenvalue = 0.0
    iterations = pivot_step = 0
    residual = float("inf")
    if (n == 0 or len(a) != n or any(len(row) != n for row in a)
            or not isfinite(shift) or not tol > 0.0 or max_iter < 1):
        return y.copy(), eigenvalue, iterations, residual, pivot_step, 1
    for i in range(n):
        for j in range(i+1, n):
            if a[i][j] != a[j][i]:
                return y.copy(), eigenvalue, iterations, residual, pivot_step, 1

    norm_y = sqrt(sum(value*value for value in y))
    if norm_y == 0.0:
        return y.copy(), eigenvalue, iterations, residual, pivot_step, 1
    y = [value/norm_y for value in y]

    # Subtract the fixed shift from the diagonal, leaving A unchanged.
    lu = [row.copy() for row in a]
    for i in range(n):
        lu[i][i] -= shift
    ipiv, pivot_step = lu_factor(lu)
    if pivot_step != 0:
        return y, eigenvalue, iterations, residual, pivot_step, 2

    for k in range(1, max_iter+1):
        # Solve with the same LU factors at every iteration.
        z = lu_solve(lu, ipiv, y)
        norm_z = sqrt(sum(value*value for value in z))
        if norm_z == 0.0:
            return y, eigenvalue, iterations, residual, pivot_step, 3
        y = [value/norm_z for value in z]

        ay = matrix_vector(a, y)
        eigenvalue = sum(y[i]*ay[i] for i in range(n))
        residual = sqrt(sum((ay[i]-eigenvalue*y[i])**2 for i in range(n)))
        iterations = k
        if residual <= tol:
            return y, eigenvalue, iterations, residual, pivot_step, 0
    return y, eigenvalue, iterations, residual, pivot_step, 4


def lu_factor(lu):
    """Factor B as P*B=L*U in place, with partial row pivoting.

    The lower triangle stores L multipliers; diagonal and upper
    entries store U. ipiv[k] records the row swapped with k. Exchange
    entire rows, including multipliers from previous columns. Return
    one-based pivot_step for a zero or nearly zero pivot, else zero.
    """
    n = len(lu)
    ipiv = [0]*n
    scale = max(abs(value) for row in lu for value in row)
    pivot_tol = 10.0*n*float_info.epsilon*scale

    for k in range(n):
        pivot_row = k
        best = abs(lu[k][k])
        for i in range(k+1, n):
            if abs(lu[i][k]) > best:
                best = abs(lu[i][k])
                pivot_row = i
        ipiv[k] = pivot_row
        if best <= pivot_tol:
            return ipiv, k+1
        if pivot_row != k:
            lu[k], lu[pivot_row] = lu[pivot_row], lu[k]

        for i in range(k+1, n):
            lu[i][k] /= lu[k][k]
            for j in range(k+1, n):
                lu[i][j] -= lu[i][k]*lu[k][j]
    return ipiv, 0


def lu_solve(lu, ipiv, b):
    """Apply saved row swaps, then solve L*w=P*b and U*z=w."""
    n = len(b)
    z = b.copy()
    for i in range(n):
        if ipiv[i] != i:
            z[i], z[ipiv[i]] = z[ipiv[i]], z[i]
    for i in range(n):
        for j in range(i):
            z[i] -= lu[i][j]*z[j]
    for i in range(n-1, -1, -1):
        for j in range(i+1, n):
            z[i] -= lu[i][j]*z[j]
        z[i] /= lu[i][i]
    return z


def matrix_vector(a, v):
    """Form A*v explicitly from the original matrix."""
    n = len(v)
    w = [0.0]*n
    for i in range(n):
        for j in range(n):
            w[i] += a[i][j]*v[j]
    return w


if __name__ == "__main__":
    main()