# 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.
from math import isnan, ceil, log
from numpy import spacing

DEFAULTDELTA = 0.1
DEFAULTSIGMA = 0.9 
eps = spacing

# Generate initial guess for step size (HZ, stage I0)
def alphainit(alpha, x, gr, f_x, psi0=0.01):
    if isnan(alpha):
        alpha = 1
        gr_max = maximum(abs(gr))
        if gr_max != 0.0:
            x_max = maximum(abs(x))
            if x_max != 0.0:
                alpha = psi0 * x_max / gr_max
            elif f_x != 0.0:
                alpha = psi0 * abs(f_x) / norm(gr)
    return alpha

# Code used to use phi, a 1-parameter function induced by f and s
# Need to pass s as an explicit parameter
def alphatry(alpha, d, x, s, xtmp, gtmp, lsr, psi1=0.2, psi2=2, psi3=0.1, iterfinitemax=ceil(-log(eps(1), 2)), alphamax=float('Inf'), display=False):

    phi0 = lsr.value[0]
    dphi0 = lsr.slope[0]

    alphatest = psi1 * alpha
    alphatest = min(alphatest, a,
            alphamax=self.stpmaxlphamax)

    # Use xtmp here
    phitest = d.f(x + alphatest * s)

    a = (phitest - alphatest * dphi0 - phi0) / alphatest**2 # quadratic fit
    if display:
        print(("quadfit: alphatest = ", alphatest,
                ", phi0 = ", phi0,
                ", phitest = ", phitest,
                ", quadcoef = ", a))

    mayterminate = False
    if a > 0 and phitest <= phi0:
        alpha = -dphi0 / 2. / a # if convex, choose minimum of quadratic
        if alpha == 0:
            error("alpha is zero. dphi0 = ", dphi0, ", a = ", a)
        if alpha <= alphamax:
            mayterminate = True
        else:
            alpha = alphamax
            mayterminate = False
        if display:
            print(("alpha guess (quadratic): ", alpha,
                    ",(mayterminate = ", mayterminate, ")"))
    else:
        if phitest > phi0:
            alpha = alphatest
        else:
            alpha *= psi2 # if not convex, expand the interval

    alpha = min(alphamax, alpha)
    if display:
        print(("alpha guess (expand): ", alpha))
    return alpha, mayterminate


class LSR:
    def __init__(self, value, slope, alpha):
        self.value = value
        self.slope = slope
        self.alpha = alpha

def hz_linesearch(df, c, mayterminate, delta=DEFAULTDELTA, sigma=DEFAULTSIGMA, alphamax=float('Inf'), rho=5., epsilon=1e-6, gamma=0.66, linesearchmax=50, psi3=0.1, iterfinitemax=ceil(-log(eps(1), 2)), display=False):

    lsr = []  # List consisting of LSR objects

    phi0, dphi0 = linefunc(df, 0, calc_grad=True)
    lsr.append(LSR(phi0, dphi0, 0))

    philim = phi0 + epsilon * abs(phi0)
    assert c > 0
    assert not isnan(c) and c <= alphamax

    phic, dphic = linefunc(df, c, calc_grad=True)
    iterfinite = 1

    lsr.append(LSR(phic, dphic, c))

    # If c was generated by quadratic interpolation, check whether it
    # satisfies the Wolfe conditions
    if mayterminate and satisfies_wolfe(c, phic, dphic, phi0, dphi0, philim, delta, sigma):
        if display:
            print(("Wolfe condition satisfied on point alpha = ", c))
        return c

    # Initial bracketing step (HZ, stages B0-B3)
    isbracketed = False
    ia = 1
    ib = 2
    assert len(lsr) == 2
    iter = 1
    cold = -1
    while not isbracketed and iter < linesearchmax:
        if display:
            print(("bracketing: ia = ", ia,
                    ", ib = ", ib,
                    ", c = ", c,
                    ", phic = ", phic,
                    ", dphic = ", dphic))
        if dphic >= 0:
            # We've reached the upward slope, so we have b; examine
            # previous values to find a
            ib = len(lsr)-1
            for i in range(ib, -1, -1):
                if lsr[i].value <= philim:
                    ia = i
                    break

            isbracketed = True
        elif lsr[-1].value > philim:
            # The value is higher, but the slope is downward, so we must
            # have crested over the peak. Use bisection.
            ib = len(lsr)-1
            ia = ib - 1
            if c != lsr[ib].alpha or lsr[ib].slope >= 0:
                raise ValueError("c = ", c, ", lsr = ", lsr)

            ia, ib = bisect(df, lsr, ia, ib, philim, display)
            isbracketed = True
        else:
            # We'll still going downhill, expand the interval and try again
            cold = c
            c *= rho
            if c > alphamax:
                if display:
                    print("bracket: exceeding alphamax, truncating")
                c = alphamax

            phic, dphic = linefunc(df, c, calc_grad=True)
            iterfinite = 1

            if dphic < 0 and c == alphamax:
                # We're on the edge of the allowed region, and the
                # value is still decreasing. This can be due to
                # roundoff error in barrier penalties, a barrier
                # coefficient being so small that being eps() away
                # from it still doesn't turn the slope upward, or
                # mistakes in the user's function.
                if iterfinite >= iterfinitemax:
                    print("Warning: failed to expand interval to bracket with finite values. If this happens frequently, check your function and gradient.")
                    print(("c = ", c,
                            ", alphamax = ", alphamax,
                            ", phic = ", phic,
                            ", dphic = ", dphic))
                return c
            lsr.append(LSR(phic, dphic, c))
        iter += 1

    while iter < linesearchmax:
        a = lsr[ia].alpha
        b = lsr[ib].alpha
        assert b > a
        if display:
            print(("linesearch: ia = ", ia,
                    ", ib = ", ib,
                    ", a = ", a,
                    ", b = ", b,
                    ", phi(a) = ", lsr[ia].value,
                    ", phi(b) = ", lsr[ib].value))

        if b - a <= eps(b):
            return a

        iswolfe, iA, iB = secant2(df, lsr, ia, ib, philim, delta, sigma, display)
        if iswolfe:
            return lsr[iA].alpha

        A = lsr[iA].alpha
        B = lsr[iB].alpha
        assert B > A
        if B - A < gamma * (b - a):
            if display:
                print("Linesearch: secant succeeded")
            ia = iA
            ib = iB
        else:
            # Secant is converging too slowly, use bisection
            if display:
                println("Linesearch: secant failed, using bisection")

            c = (A + B) / 2.
            # phic = phi(gphi, c) # TODO: Replace
            phic, dphic = linefunc(df, c, calc_grad=True)
            assert not isnan(phic)
            lsr.append(LSR(phic, dphic, c))
            ia, ib = update(df, lsr, iA, iB, len(lsr)-1, philim, display)

        iter += 1
    error("Linesearch failed to converge")

# Check Wolfe & approximate Wolfe
def satisfies_wolfe(c, phic, dphic, phi0, dphi0, philim, delta=DEFAULTDELTA, sigma=DEFAULTSIGMA):
    wolfe1 = delta * dphi0 >= (phic - phi0) / c and dphic >= sigma * dphi0
    wolfe2 = (2.0 * delta - 1.0) * dphi0 >= dphic >= sigma * dphi0 and phic <= philim
    return wolfe1 or wolfe2

# HZ, stages S1-S4
def secant(a, b, dphia, dphib):
    return (a * dphib - b * dphia) / (dphib - dphia)

def secant_(lsr, ia, ib):
    return secant(lsr[ia].alpha, lsr[ib].alpha, lsr[ia].slope, lsr[ib].slope)

# phi
def secant2(df, lsr, ia, ib, philim, delta=DEFAULTDELTA, sigma=DEFAULTSIGMA, display=False):

    phi0 = lsr[0].value
    dphi0 = lsr[0].slope
    a = lsr[ia].alpha
    b = lsr[ib].alpha
    dphia = lsr[ia].slope
    dphib = lsr[ib].slope
    assert dphia < 0
    assert dphib >= 0
    c = secant(a, b, dphia, dphib)
    if display:
        print(("secant2: a = ", a, ", b = ", b, ", c = ", c))

    assert not isnan(c)
    # phic = phi(tmpc, c) # Replace

    phic, dphic = linefunc(df, c, calc_grad=True)
    assert not isnan(phic)
    lsr.append(LSR(phic, dphic, c))
    ic = len(lsr)-1
    if satisfies_wolfe(c, phic, dphic, phi0, dphi0, philim, delta, sigma):
        if display:
            print("secant2: first c satisfied Wolfe conditions")
        return True, ic, ic

    iA, iB = update(df, lsr, ia, ib, ic, philim, display)
    if display:
        print(("secant2: iA = ", iA, ", iB = ", iB, ", ic = ", ic))

    a = lsr[iA].alpha
    b = lsr[iB].alpha
    doupdate = false
    if iB == ic:
        # we updated b, make sure we also update a
        c = secant_(lsr, ib, iB)
    elif iA == ic:
        # we updated a, do it for b too
        c = secant_(lsr, ia, iA)
    if a <= c <= b:
        if display:
            print(("secant2: second c = ", c))

        phic, dphic = linefunc(df, c, calc_grad=True)
        assert not isnan(phic)
        lsr.append(LSR(phic, dphic, c))
        ic = len(lsr)-1
        # Check arguments here
        if satisfies_wolfe(c, phic, dphic, phi0, dphi0, philim, delta, sigma):
            if display:
                print("secant2: second c satisfied Wolfe conditions")
            return True, ic, ic

        iA, iB = update(df, lsr, iA, iB, ic, philim, display)

    if display: 
        print(("secant2 output: a = ", lsr[iA].alpha, ", b = ", lsr[iB].alpha))

    return False, iA, iB

# HZ, stages U0-U3
# Given a third point, pick the best two that retain the bracket
# around the minimum (as defined by HZ, eq. 29)
# b will be the upper bound, and a the lower bound
def update(df, lsr, ia, ib, ic, philim, display=False):

    a = lsr[ia].alpha
    b = lsr[ib].alpha
    # Debugging (HZ, eq. 4.4):
    assert lsr[ia].slope < 0
    assert lsr[ia].value <= philim
    assert lsr[ib].slope >= 0
    assert b > a
    c = lsr[ic].alpha
    phic = lsr[ic].value
    dphic = lsr[ic].slope
    if display:
        print(("update: ia = ", ia,
                ", a = ", a,
                ", ib = ", ib,
                ", b = ", b,
                ", c = ", c,
                ", phic = ", phic,
                ", dphic = ", dphic))

    if c < a or c > b:
        return ia, ib, 0, 0  # it's out of the bracketing interval

    if dphic >= 0:
        return ia, ic, 0, 0  # replace b with a closer point

    # We know dphic < 0. However, phi may not be monotonic between a
    # and c, so check that the value is also smaller than phi0.  (It's
    # more dangerous to replace a than b, since we're leaving the
    # secure environment of alpha=0; that's why we didn't check this
    # above.)
    if phic <= philim:
        return ic, ib, 0, 0  # replace a

    # phic is bigger than phi0, which implies that the minimum
    # lies between a and c. Find it via bisection.
    return bisect(df, lsr, ia, ic, philim, display)

# HZ, stage U3 (with theta=0.5)
def bisect(df, lsr, ia, ib, philim, display=False):

    gphi = nan(T)
    a = lsr[ia].alpha
    b = lsr[ib].alpha
    # Debugging (HZ, conditions shown following U3)
    assert lsr[ia].slope < 0
    assert lsr[ia].value <= philim
    assert lsr[ib].slope < 0       # otherwise we wouldn't be here
    assert lsr[ib].value > philim
    assert b > a

    while b - a > eps(b):
        if display:
            print(("bisect: a = ", a, ", b = ", b, ", b - a = ", b - a))

        d = (a + b) / 2.
        phid, gphi = linefunc(df, d, calc_grad=True)
        assert isfinite(phid)
        lsr.append(LSR(phid, gphi, d))
        id = len(lsr)-1
        if gphi >= 0:
            return ia, id # replace b, return

        if phid <= philim:
            a = d # replace a, but keep bisecting until dphib > 0
            ia = id
        else:
            b = d
            ib = id

    return ia, ib

# Define one-parameter function for line searches
def linefunc(df, alpha, calc_grad):

    if calc_grad:
        return df.phi_dphi(alpha)
    else:
        return df.phi(alpha)
