import torch


def cg(f_Ax, b, cg_iters=10, residual_tol=1e-10):
    """
    Conjugate gradient descent.
    Algorithm is from page 312 of book
    "Applied Numerical Linear Algebra" by James W. Demmel

    Approximately solve x = A^{-1}b, or Ax = b,
    Here A is F and b is g.

    A good thing about this algorithm is that, to compute x,
    we only need to have access of function f: d -> Ad, where
    function f takes an arbitary vector d and output the
    Fisher-vector product Ad, where A is the fisher information
    matrix, i.e. the hessian of KL divergence.

    This method uses krylov space solver.
    See
    "A Brief Introduction to Krylov Space Methods for Solving Linear Systems"
    http://www.sam.math.ethz.ch/~mhg/pub/biksm.pdf

    Specifically, it create a sequence of Krylov subspaces K_n for Ax = b,
    and find approximate solutions x_n in K_n

    Parameters
    ----------
    f_Ax (function): A funtion takes a vector x as input and output
        Ax where A is the Fisher Information matrix.
    b (torch.Tensor):
    cg_iter (int): A integer, indicate the iteration for conjugate
        gradient.
    residual_tol (float): A float, indicate the terminate condition
        for the algorithm.

    Returns
    -------
    x (torch.Tensor): result of conjugate gradient
    """
    p = b.clone()
    r = b.clone()
    x = torch.zeros(b.size())
    rdotr = r.dot(r)

    for i in range(cg_iters):
        z = f_Ax(p)
        v = rdotr / p.dot(z)
        x += v * p
        r -= v * z
        newrdotr = r.dot(r)
        mu = newrdotr / rdotr
        p = r + mu * p
        rdotr = newrdotr
        if rdotr < residual_tol:
            break
    return x