import math

N = 1000

EPS = 1e-4

ALPHA = 0.01
BETA = 10


def mul(bottom, middle, top, x):
    ret = [0] * N

    for i in range(N):
        ret[i] = middle[i] * x[i]
        if i > 0:
            ret[i] += bottom[i] * x[i - 1]
        if i < N - 1:
            ret[i] += top[i] * x[i + 1]

    return ret


def diff(u, v):
    return [u[i] - v[i] for i in range(N)]


def norm(u):
    return math.sqrt(sum(x * x for x in u))


def universal_solver(bottom, middle, top, f, x, tau=None):
    ret_e = []
    inv = [1.0 / w for w in middle]
    Ax = mul(bottom, middle, top, x)
    r0 = norm(diff(Ax, f))
    while True:
        if tau is not None:
            Ax = mul(bottom, middle, top, x)
            Ax = [tau * w for w in Ax]
        else:
            Ax = mul([0] * N, inv, [0] * N, mul(bottom, middle, top, x))

        r = norm(diff(Ax, f))
        e = r / r0

        ret_e.append(e)
        if e < EPS:
            break

        x = [x[i] + (f[i] - Ax[i]) for i in range(N)]

    return ret_e, x


def main():
    tau = 0.1651512666112257
    bottom = [-1] * N
    top = [-1] * N
    middle = [2 + ALPHA] * N
    middle[0] = 2 + BETA
    x = [0] * N

    f = [1 if 494 <= i < 505 else 0 for i in range(N)]
    f_tau = [tau * w for w in f]
    inv_f = [1.0 / middle[i] if 494 <= i < 505 else 0 for i in range(N)]

    mti_e, mti_x = universal_solver(bottom, middle, top, f_tau, [0] * N, tau)
    jacobi_e, jacobi_x = universal_solver(bottom, middle, top, inv_f, x)

    index_nev = [i + 1 for i in range(len(mti_e))]
    plt.plot(index_nev, mti_e)
    plt.show()

    index_solve = [i + 1 for i in range(len(mti_x))]
    plt.plot(index_solve, mti_x)
    plt.show()


if __name__ == '__main__':
    main()