import matplotlib as plt

def run(A, f, w, eps=1e-12):
    iterations = 0
    f = np.matmul(A.T, f)
    A = np.matmul(A.T, A)
    x = np.zeros(A.shape[1], dtype=np.float64)
    B = np.zeros(A.shape, dtype=np.float64)
    for i in range(A.shape[0]):
        B[i][i] = A[i][i] / w
        for j in range(A.shape[0]):
            if j < i:
                B[i][j] = A[i][j]
    loss = np.sum((f - np.matmul(A, x)) ** 2) ** (1 / 2)
    while loss > eps:
        iterations += 1
        x = np.matmul(np.linalg.inv(B), f - np.matmul(A - B, x))
        loss = np.sum((f - np.matmul(A, x)) ** 2) ** (1 / 2)
    return x, iterations

def execute(par, A, f, w, eps=1e-12):
    x, iters = run(A, f, w, eps)
    if par == "solution":
        return x
    if par == "iterations":
        return iters

def learn_iter(A, f, eps=1e-12):
    w_step = 0.1
    W = np.arange(w_step, 2, w_step)
    iters = []
    for w in W:
        iters += [execute("iterations", A, f, w, eps)]
    plt.plot(W, iters)
    plt.show()

def check(A, f):
## Для проверки функций я использую функции библиотеки для анализа данных numpy
    return np.linalg.solve(A, f)

def main():
    print("Insert the dimension of matrix")
    n = int(input())

    print("\nType of input\n")
    print("\t1 - File\n\t2 - Use formula")
    typ = int(input())

    A = np.zeros(shape=(n, n), dtype=np.float64)
    b = np.zeros(shape=(n), dtype=np.float64)
    if typ == 1:
        print("\nEnter filename\n")
        filename = input()
        with open(filename) as f:
            for i in range(n):
                line = list(f.readline().split())
                for j, elem in enumerate(line[:n]):
                    A[i][j] = np.float64(elem)
                b[i] = np.float64(line[n])
    elif typ == 2:
        print("\nEnter parametr m\n")
        m = np.float64(input())
        for i in range(n):
            for j in range(n):
                if i == j:
                    A[i][j] = n + m ** 2 + j / m + i / n
                else:
                    A[i][j] = (i + j) / (m + n)
            b[i] = 200 + 50 * i

    print("\nChoose operation\n")
    print("\t1 - Find solution\n" +
          "\t2 - Learn iterations\n")
    operation = int(input())
    w = 1
    if operation == 1:
        print("\nEnter w\n")
        w = int(input())
    print("solution =\n {0}".format(execute("solution", A, b, w)))
    print("check:\n {0}\n".format(check(A, b)))
    if operation == 2:
        learn_iter(A, b)

main()