"""ch08_jacobi_eigen.py

Find all eigenvalues and eigenvectors of a real symmetric matrix.

Origin:
    Python translation of ch08_jacobi_eigen.f90 by Alexander Godunov,
    based on Jacobi.f90 (December 2009). The cyclic, thresholded Jacobi
    rotations and the original formulas for c and s are retained.

Example:
    A = [ 1  2  3; 2  2 -2; 3 -2  4 ]. Its eigenvalues are approximately
    -2.54138127, 3.54138127, and 6.00000000.

Columns of X correspond to the eigenvalues on the diagonal of the
transformed matrix, in the order produced by Jacobi. Eigenvectors
may differ by an overall sign. Use library routines for large matrices.
"""

from math import sqrt


def main():
    """Run the three-dimensional example and check every eigenpair."""
    a = [[1.0,  2.0,  3.0],
         [2.0,  2.0, -2.0],
         [3.0, -2.0,  4.0]]
    # As in the Fortran code, the tolerance bounds the SUM OF SQUARED
    # off-diagonal entries, not the off-diagonal norm itself.
    abserr = 1.0e-18
    max_sweeps = 100
    n = len(a)

    print("Eigenvalues and eigenvectors (Jacobi method)")
    print("\nOriginal symmetric matrix A:")
    for row in a:
        print("".join(f"{value:12.6f}" for value in row))
    print(f"Squared off-diagonal tolerance = {abserr:.4e}")

    # The routine changes its matrix. Keep A for an independent check.
    a_work = [row.copy() for row in a]
    x, sweeps, rotations, offdiag_norm, info = jacobi(
        a_work, abserr, max_sweeps)
    if info == 1:
        raise ValueError("Invalid input: A must be symmetric and controls positive")

    print("\nEigenvalues and corresponding unit eigenvectors:")
    for j in range(n):
        print(f"lambda({j+1}) = {a_work[j][j]:12.6f}")
        for i in range(n):
            print(f"  v({i+1}) = {x[i][j]:12.6f}")

    # Check each component of A*X-X*D using the original coefficients.
    max_residual = 0.0
    max_orthogonality = 0.0
    for j in range(n):
        for i in range(n):
            value = 0.0
            for k in range(n):
                value += a[i][k]*x[k][j]
            max_residual = max(max_residual,
                               abs(value-a_work[j][j]*x[i][j]))

            value = 0.0
            for k in range(n):
                value += x[k][i]*x[k][j]
            defect = 1.0 if i == j else 0.0
            max_orthogonality = max(max_orthogonality,
                                    abs(value-defect))

    print(f"\nSweeps = {sweeps}")
    print(f"Rotations = {rotations}")
    print(f"Off-diagonal norm of final matrix = {offdiag_norm:.4e}")
    print(f"Maximum |A*X-X*D| component = {max_residual:.4e}")
    print(f"Maximum |X^T*X-I| component = {max_orthogonality:.4e}")
    if info == 2:
        raise ArithmeticError("Sweep limit reached before the tolerance was met")


def jacobi(a, abserr, max_sweeps):
    """Diagonalize a small real symmetric matrix by Jacobi rotations.

    Input/output:
        a            - symmetric matrix, overwritten by a nearly
                       diagonal similar matrix; its diagonal holds
                       the eigenvalues
    Output:
        x            - orthonormal eigenvectors in columns
        sweeps       - completed passes over all off-diagonal pairs
        rotations    - number of rotations performed
        offdiag_norm - sqrt of the sum of off-diagonal squares
        info         - 0 converged; 1 invalid input; 2 sweep limit
    Controls:
        abserr       - positive bound on SUM of squared off-diagonal entries
        max_sweeps   - positive maximum number of complete passes

    Start X as the identity. For each pair above the threshold, rotate
    rows and columns of A and columns of X. Recompute the off-diagonal
    sum each sweep so bookkeeping roundoff cannot falsely report success.
    The angle formulas follow the original Jacobi.f90.
    """
    n = len(a)
    sweeps = 0
    rotations = 0
    offdiag_norm = float("inf")
    if (n == 0 or any(len(row) != n for row in a)
            or not abserr > 0.0 or max_sweeps < 1):
        return [], sweeps, rotations, offdiag_norm, 1
    for i in range(n):
        for j in range(i+1, n):
            if a[i][j] != a[j][i]:
                return [], sweeps, rotations, offdiag_norm, 1

    x = [[0.0]*n for _ in range(n)]
    for i in range(n):
        x[i][i] = 1.0

    while True:
        # Measure the actual off-diagonal entries at each sweep.
        b2 = 0.0
        for i in range(n):
            for j in range(n):
                if i != j:
                    b2 += a[i][j]*a[i][j]
        offdiag_norm = sqrt(b2)
        if b2 <= abserr:
            return x, sweeps, rotations, offdiag_norm, 0
        if sweeps >= max_sweeps:
            return x, sweeps, rotations, offdiag_norm, 2

        # Half the mean off-diagonal square sets the skip threshold.
        bar = 0.5*b2/(n*n)
        for i in range(n-1):
            for j in range(i+1, n):
                if a[j][i]*a[j][i] <= bar:
                    continue
                b2 = max(0.0, b2-2.0*a[j][i]*a[j][i])
                bar = 0.5*b2/(n*n)

                # Choose c and s as in Jacobi.f90.
                beta = (a[j][j]-a[i][i])/(2.0*a[j][i])
                coeff = 0.5*beta/sqrt(1.0+beta*beta)
                s = sqrt(max(0.5+coeff, 0.0))
                c = sqrt(max(0.5-coeff, 0.0))

                # First rotate rows i and j of A.
                for k in range(n):
                    cs =  c*a[i][k]+s*a[j][k]
                    sc = -s*a[i][k]+c*a[j][k]
                    a[i][k] = cs
                    a[j][k] = sc

                # Then rotate columns i and j of A and X.
                for k in range(n):
                    cs =  c*a[k][i]+s*a[k][j]
                    sc = -s*a[k][i]+c*a[k][j]
                    a[k][i] = cs
                    a[k][j] = sc

                    cs =  c*x[k][i]+s*x[k][j]
                    sc = -s*x[k][i]+c*x[k][j]
                    x[k][i] = cs
                    x[k][j] = sc
                rotations += 1
        sweeps += 1


if __name__ == "__main__":
    main()