import numpy as np
from scipy.linalg import lu, solve_triangular


def table4_15(max_iterations=1000000):
    # orthogonal iteration with inverse power method

    # error tolerance
    tol = 1e-7

    # number of iterations
    K = max_iterations

    # matrix
    a = 3
    m = 5
    n = 200
    print(f"n = {n}")
    A = np.zeros((n, n))
    for i in range(n):
        A[i, i] = a
        if i < n - 1:
            A[i, i + 1] = -1
        if i > 0:
            A[i, i - 1] = -1

    # step zero
    B = np.random.rand(n, m)
    Q, R = np.linalg.qr(B, mode="reduced")
    eigens0 = np.ones(m)
    eigens = np.ones(m)
    P, L, U = lu(A)

    for k in range(1, K + 1):
        Z = solve_triangular(
            L, P.T @ Q, lower=True, unit_diagonal=True, check_finite=False
        )
        B = solve_triangular(U, Z, lower=False, check_finite=False)
        # B=np.linalg.inv(A)@Q
        Q, R = np.linalg.qr(B, mode="reduced")
        errs = np.zeros(m)
        for ic in range(m):
            eigens[ic] = np.dot(Q[:, ic], A @ Q[:, ic])
            errs[ic] = (eigens[ic] - eigens0[ic]) / eigens[ic]
        err = np.linalg.norm(errs, np.inf)
        if err < tol:
            break
        eigens0 = eigens.copy()
    print(f"k = {k}")

    # exact values
    evx = np.ones(m)
    evx[0] = a + 2 * np.cos(n * np.pi / (n + 1))
    evx[1] = a + 2 * np.cos((n - 1) * np.pi / (n + 1))
    evx[2] = a + 2 * np.cos((n - 2) * np.pi / (n + 1))
    evx[3] = a + 2 * np.cos((n - 3) * np.pi / (n + 1))
    evx[4] = a + 2 * np.cos((n - 4) * np.pi / (n + 1))

    # errors
    print("\n i  \t lambda_i \t   RIE \t       RE ")
    errs2 = np.zeros(m)
    for im in range(m):
        errs2[im] = abs((eigens[im] - evx[im]) / evx[im])
        print(
            f"{n - im:d} & {eigens[im]:15.12e} & {errs[im]:5.2e} & "
            f"{errs2[im]:5.2e}  "
        )
    print()

    # np.set_printoptions(precision=15)
    # ee=np.linalg.eigvals(A)
    # ee[0:5]


if __name__ == "__main__":
    table4_15()
