import copy
import numpy as np

# # Question 1 optimisation problem
A = np.array([[1., 2., 1., -1.],
             [-1., 1., 0., 2.],
             [0., -1., -2., 1.]])
b = np.array([[3.], [2.], [-2.]])
x = np.array([[1.], [1.], [1.], [1.]])
k = 0
alpha = 0.1
gamma = 0.2

def gradient_step_1(A, b, x, gamma):
    updated_x = np.array(A.T.dot((A.dot(x) - b)) + gamma*x)
    return np.round(updated_x,4)

print_form = []

def edit_form(print_form,x):
    print_form.pop(0)
    x_k = copy.copy(x)
    print_form.append([k, x_k])

while True:
    if k > 0 and k <= 5:
        x_k= copy.copy(x)
        print_form.append((k,x_k))
    if k == 6:
        for item in print_form:
            print(f"k = {item[0]}, x({item[0]}) = {item[1]}")
    if k > 6:
        edit_form(print_form,x)
    k += 1
    x -= alpha*gradient_step_1(A, b, x, gamma)

    if np.linalg.norm(gradient_step_1(A, b, x, gamma), ord=2)<0.001:
        edit_form(print_form, x)
        k += 1
        x -= alpha * gradient_step_1(A, b, x, gamma)
        edit_form(print_form, x)
        for item in print_form:
            print(f"k = {item[0]}, x({item[0]}) = {item[1]}")
        break



# x = np.array([[1.123], [1.123], [1.123], [1.123]])
# x =np.array([[-4.59627504],
#  [8.95686683],
#  [-9.76586351],
#  [ 8.87335576]])
# print(np.around(x,2))
# for i in range(10):
#     x = A.T.dot((A.dot(x) - b)) + gamma*x
#     print(np.round(x,4))
