
#
# Conjugate gradient
#
# This is an independent implementation of:
#   W. W. Hager and H. Zhang (2006) Algorithm 851: CG_DESCENT, a
#     conjugate gradient method with guaranteed descent. ACM
#     Transactions on Mathematical Software 32: 113–137.
#
# Code comments such as "HZ, stage X" or "HZ, eqs Y" are with
# reference to a particular point in this paper.
#
# Several aspects of the following have also been incorporated:
#   W. W. Hager and H. Zhang (2012) The limited memory conjugate
#     gradient method.
#
# This paper will be denoted HZ2012 below.
#
# There are some modifications and/or extensions from what's in the
# paper (these may or may not be extensions of the cg_descent code
# that can be downloaded from Hager's site; his code has undergone
# numerous revisions since publication of the paper):
#
# cgdescent: the termination condition employs a "unit-correct"
#   expression rather than a condition on gradient
#   components---whether this is a good or bad idea will require
#   additional experience, but preliminary evidence seems to suggest
#   that it makes "reasonable" choices over a wider range of problem
#   types.
#
# linesearch: the Wolfe conditions are checked only after alpha is
#   generated either by quadratic interpolation or secant
#   interpolation, not when alpha is generated by bisection or
#   expansion. This increases the likelihood that alpha will be a
#   good approximation of the minimum.
#
# linesearch: In step I2, we multiply by psi2 only if the convexity
#   test failed, not if the function-value test failed. This
#   prevents one from going uphill further when you already know
#   you're already higher than the point at alpha=0.
#
# both: checks for Inf/NaN function values
#
# both: support maximum value of alpha (equivalently, c). This
#   facilitates using these routines for constrained minimization
#   when you can calculate the distance along the path to the
#   disallowed region. (When you can't easily calculate that
#   distance, it can still be handled by returning Inf/NaN for
#   exterior points. It's just more efficient if you know the
#   maximum, because you don't have to test values that won't
#   work.) The maximum should be specified as the largest value for
#   which a finite value will be returned.  See, e.g., limits_box
#   below.  The default value for alphamax is Inf. See alphamaxfunc
#   for cgdescent and alphamax for linesearch_hz.

macro cgtrace()
    quote
        if tracing
            dt = Dict()
            if o.extended_trace
                dt["x"] = copy(x)
                dt["g(x)"] = copy(g)
                dt["Current step size"] = alpha
            end
            g_norm = vecnorm(g, Inf)
            update!(tr,
                    iteration,
                    f_x,
                    g_norm,
                    dt,
                    o.store_trace,
                    o.show_trace,
                    o.show_every,
                    o.callback)
        end
    end
end

immutable ConjugateGradient{T} <: Optimizer
    eta::Float64
    P::T
    precondprep!::Function
    linesearch!::Function
end

function ConjugateGradient(;
                           linesearch!::Function = hz_linesearch!,
                           eta::Real = 0.4,
                           P::Any = nothing,
                           precondprep! = (P, x) -> nothing)
    ConjugateGradient{typeof(P)}(Float64(eta),
                                 P, precondprep!,
                                 linesearch!)
end

function optimize{T}(df::DifferentiableFunction,
                     initial_x::Array{T},
                     mo::ConjugateGradient,
                     o::OptimizationOptions)
    # Print header if show_trace is set
    print_header(o)

    # Maintain current state in x and previous state in x_previous
    x, x_previous = copy(initial_x), copy(initial_x)

    # Count the total number of iterations
    iteration = 0

    # Track calls to function and gradient
    f_calls, g_calls = 0, 0

    # Count number of parameters
    n = length(x)

    # Maintain current gradient in g and previous gradient in g_previous
    g, g_previous = similar(x), similar(x)

    # Maintain the preconditioned gradient in pg
    pg = similar(x)

    # The current search direction
    s = similar(x)

    # Buffers for use in line search
    x_ls, g_ls = similar(x), similar(x)

    # Intermediate value in CG calculation
    y = similar(x)
    py = similar(x)

    # Store f(x) in f_x
    f_x = df.fg!(x, g)
    @assert typeof(f_x) == T
    f_x_previous = convert(T, NaN)
    f_calls, g_calls = f_calls + 1, g_calls + 1
    copy!(g_previous, g)

    # Keep track of step-sizes
    alpha = alphainit(one(T), x, g, f_x)

    # TODO: How should this flag be set?
    mayterminate = false

    # Maintain a cache for line search results
    lsr = LineSearchResults(T)

    # Trace the history of states visited
    tr = OptimizationTrace{typeof(mo)}()
    tracing = o.store_trace || o.show_trace || o.extended_trace || o.callback != nothing
    @cgtrace

    # Output messages
    if !isfinite(f_x)
        error("Must have finite starting value")
    end
    if !all(isfinite(g))
        @show g
        @show find(!isfinite(g))
        error("Gradient must have all finite values at starting point")
    end

    # Determine the intial search direction
    #    if we don't precondition, then this is an extra superfluous copy
    #    TODO: consider allowing a reference for pg instead of a copy
    mo.precondprep!(mo.P, x)
    A_ldiv_B!(pg, mo.P, g)
    scale!(copy!(s, pg), -1)

    # Assess multiple types of convergence
    x_converged, f_converged = false, false
    g_converged = vecnorm(g, Inf) < o.g_tol

    # Iterate until convergence
    converged = g_converged
    while !converged && iteration < o.iterations
        # Increment the number of steps we've had to perform
        iteration += 1

        # Reset the search direction if it becomes corrupted
        dphi0 = vecdot(g, s)
        if dphi0 >= 0
            @simd for i in 1:n
                @inbounds s[i] = -pg[i]
            end
            dphi0 = vecdot(g, s)
            if dphi0 >= 0
                break
            end
        end

        # Refresh the line search cache
        clear!(lsr)
        @assert typeof(f_x) == T
        @assert typeof(dphi0) == T
        push!(lsr, zero(T), f_x, dphi0)

        # Pick the initial step size (HZ #I1-I2)
        alpha, mayterminate, f_update, g_update =
          alphatry(alpha, df, x, s, x_ls, g_ls, lsr)
        f_calls, g_calls = f_calls + f_update, g_calls + g_update

        # Determine the distance of movement along the search line
        alpha, f_update, g_update =
          mo.linesearch!(df, x, s, x_ls, g_ls, lsr, alpha, mayterminate)
        f_calls, g_calls = f_calls + f_update, g_calls + g_update

        # Maintain a record of previous position
        copy!(x_previous, x)

        # Update current position # x = x + alpha * s
        LinAlg.axpy!(alpha, s, x)

        # Maintain a record of the previous gradient
        copy!(g_previous, g)

        # Update the function value and gradient
        f_x_previous, f_x = f_x, df.fg!(x, g)
        f_calls, g_calls = f_calls + 1, g_calls + 1

        x_converged,
        f_converged,
        g_converged,
        converged = assess_convergence(x,
                                       x_previous,
                                       f_x,
                                       f_x_previous,
                                       g,
                                       o.x_tol,
                                       o.f_tol,
                                       o.g_tol)

        # Check sanity of function and gradient
        if !isfinite(f_x)
            error("Function must finite function values")
        end

        # Determine the next search direction using HZ's CG rule
        #  Calculate the beta factor (HZ2012)
        # -----------------
        # Comment on py: one could replace the computation of py with
        #    ydotpgprev = vecdot(y, pg)
        #    vecdot(y, py)  >>>  vecdot(y, pg) - ydotpgprev
        # but I am worried about round-off here, so instead we make an
        # extra copy, which is probably minimal overhead.
        # -----------------
        mo.precondprep!(mo.P, x)
        dPd = dot(s, mo.P, s)
        etak::T = mo.eta * vecdot(s, g_previous) / dPd
        @simd for i in 1:n
            @inbounds y[i] = g[i] - g_previous[i]
        end
        ydots = vecdot(y, s)
        copy!(py, pg)        # below, store pg - pg_previous in py
        A_ldiv_B!(pg, mo.P, g)
        @simd for i in 1:n     # py = pg - py
           @inbounds py[i] = pg[i] - py[i]
        end
        betak = (vecdot(y, pg) - vecdot(y, py) * vecdot(g, s) / ydots) / ydots
        beta = max(betak, etak)
        @simd for i in 1:n
            @inbounds s[i] = beta * s[i] - pg[i]
        end

        @cgtrace
    end

    return MultivariateOptimizationResults("Conjugate Gradient",
                                           initial_x,
                                           x,
                                           Float64(f_x),
                                           iteration,
                                           iteration == o.iterations,
                                           x_converged,
                                           o.x_tol,
                                           f_converged,
                                           o.f_tol,
                                           g_converged,
                                           o.g_tol,
                                           tr,
                                           f_calls,
                                           g_calls)
end
