"""Literal Python conversion of table3_5.m."""

import numpy as np


def table3_5(n_values=None, Nv=5):
    # tridaigonal
    if n_values is None:
        n = np.array([4, 8, 12, 16, 160, 1600, 16000], dtype=int)
    else:
        n = np.asarray(n_values, dtype=int)
    # n=np.array([4, 8, 12, 16, 160])
    N = len(n)
    NN = len(n)
    eM = np.zeros(N)
    cA = np.zeros(N)
    ek = np.zeros(N)
    for in_ in range(N):
        A = np.zeros((n[in_], n[in_]))
        for i in range(n[in_]):
            A[i, i] = 2
            if i > 0:
                A[i, i - 1] = 1
            if i < n[in_] - 1:
                A[i, i + 1] = 1
        xe = np.ones(n[in_])
        a = 2 * np.ones(n[in_])
        b = np.ones(n[in_])
        c = np.ones(n[in_])
        f = 4 * np.ones(n[in_])
        f[0] = 3
        f[n[in_] - 1] = 3
        xc = tridiag(a, b, c, f)
        err = xe - xc
        eM[in_] = np.linalg.norm(err, np.inf)
        cA[in_] = np.linalg.cond(A, np.inf)
        ek[in_] = np.finfo(float).eps * np.linalg.cond(A, np.inf)

    # vandermonde
    eMv = np.zeros(N)
    cAv = np.zeros(N)
    ekv = np.zeros(N)
    # n=np.array([4, 8, 12, 16, 160])
    Nv = min(Nv, N)
    for in_ in range(Nv):
        A = np.zeros((n[in_], n[in_]))
        for i in range(n[in_]):
            a = 1 / (i + 1)
            A[i, 0] = 1
            for j in range(1, n[in_]):
                A[i, j] = a ** j
        xe = np.ones(n[in_])
        b = A @ xe
        xc = np.linalg.solve(A, b)
        r = b - A @ xc
        err = xe - xc
        eMv[in_] = np.linalg.norm(err, np.inf)
        cAv[in_] = np.linalg.cond(A, np.inf)
        ekv[in_] = np.finfo(float).eps * np.linalg.cond(A, np.inf)
        # print("\n %d & %17.2e & %17.1e & %17.1e" %
        #       (n[in_], eM[in_], cA[in_], ek[in_]), end="")

    for in_ in range(NN):
        print(r"   %i  &  %5.1e  &  %5.1e  &  %5.1e      &  %5.1e  &  %5.1e  &  %5.1e  \\\hline" %
              (n[in_], eM[in_], cA[in_], ek[in_], eMv[in_], cAv[in_], ekv[in_]))

    print()


def tridiag(a, b, c, f):
    #  Solve the  N x N  tridiagonal system for y:
    #
    #  [ a(1)  c(1)                                  ] [  y(1)  ]   [  f(1)  ]
    #  [ b(2)  a(2)  c(2)                            ] [  y(2)  ]   [  f(2)  ]
    #  [       b(3)  a(3)  c(3)                      ] [        ]   [        ]
    #  [            ...   ...   ...                  ] [  ...   ] = [  ...   ]
    #  [                    ...    ...    ...        ] [        ]   [        ]
    #  [                        b(N-1) a(N-1) c(N-1) ] [ y(N-1) ]   [ f(N-1) ]
    #  [                                 b(N)  a(N)  ] [  y(N)  ]   [  f(N)  ]
    #
    #  f must be a vector (row or column) of length N
    #  a, b, c must be vectors of length N (note that b(1) and c(N) are not used)

    N = len(f)
    v = np.zeros(N)
    y = v.copy()

    w = a[0]
    y[0] = f[0] / w
    for i in range(1, N):
        v[i - 1] = c[i - 1] / w
        w = a[i] - b[i] * v[i - 1]
        y[i] = (f[i] - b[i] * y[i - 1]) / w

    for j in range(N - 2, -1, -1):
        y[j] = y[j] - v[j] * y[j + 1]

    return y


if __name__ == "__main__":
    table3_5()
