import numpy as np
# the Hilbert matrix is defined in scipy, so we import this function
from scipy.linalg import hilbert

# implementation of CG method
def CG(A,b,x,stopping_residual=1e-6,max_iter=2000):
    r = A@x-b
    p = -r
    normr = np.linalg.norm(r)
    n_iter = 0
    while((normr > stopping_residual) and (n_iter < max_iter)):
        # the term A@p is used several times, but
        # we only need to compute this matrix-vector product once
        Ap = A@p
        alpha = (normr**2)/np.inner(p,Ap)
        x += alpha*p
        r += alpha*Ap
        normr_old = normr
        normr = np.linalg.norm(r)
        beta = (normr**2)/(normr_old**2)
        p = -r + beta*p
        n_iter += 1
    return(x,n_iter)

# Test everything on the matrix in problem 3 to see whether
# the algorithm works
A = np.array([[2,-1,-1],[-1,3,-1],[-1,-1,2]])
b = np.array([1,0,1])
x0 = np.zeros(3)
result,num_iterations = CG(A,b,x0)
print('The method converged in {} iterations.'.format(num_iterations))

print('\n Starting the example with the Hilbert matrix.')
for n in [5,8,12,20]:
    A = hilbert(n)
    b = np.ones(n)
    x0 = np.ones(n)
    result,num_iterations = CG(A,b,x0)
    print('For n = {}, the method converged in {} iterations.'.format(n,num_iterations))
    print('The condition number of the matrix is {}.'.format(np.linalg.cond(A)))