import numpy


def cgSolver(A, b, lam, Tol=1e-16, CGMaxIter=20):
    '''
    Solve (A^T * A + lam * I) * w = b.
    '''
    d = A.shape[1]
    b = b.reshape(d, 1)
    tol = Tol * numpy.linalg.norm(b)
    w = numpy.zeros((d, 1))
    r = b - lam * w - numpy.dot(A.T, numpy.dot(A, w))
    p = r
    rsold = numpy.sum(r ** 2)

    # Compare with this https://en.wikipedia.org/wiki/Conjugate_gradient_method
    for i in range(CGMaxIter):
        # This is the hessian vector product
        Ap = lam * p + numpy.dot(A.T, numpy.dot(A, p))
        alpha = rsold / numpy.dot(p.T, Ap)
        w += alpha * p
        r -= alpha * Ap
        rsnew = numpy.sum(r ** 2)
        if numpy.sqrt(rsnew) < tol:
            print('Converged! res = ' + str(rsnew))
            break
        p = r + (rsnew / rsold) * p
        rsold = rsnew

    #print('res = ' + str(rsnew) + ',   iter = ' + str(i))
    # if i == MaxIter-1:
        #print('Warn: CG does not converge! Res = '  + str(rsnew))
    return w
