# optimize-steepdescent2d-linesearch-fixed.py
#
# Usage:
#   python optimize-steepdescent2d-linesearch-fixed.py x0 y0 lsmode functype
#
# Examples:
#   python optimize-sd2d-linesearch_modified.py  0.0  0.0 armijo general
#   python optimize-sd2d-linesearch_modified.py  3.0  3.0 armijo ellipsoid
#   python optimize-sd2d-linesearch_modified.py -1.0 -1.0 simple general
#
# lsmode:
#   one, simple, exact, golden, armijo
#
# functype:
#   general, ellipsoid

import sys
import time
import traceback

import numpy as np
from numpy import sqrt

import matplotlib.pyplot as plt


# ============================================================
# optimization parameters
# ============================================================

x0 = [0.0, 0.0]

# for line search: '', 'one', 'simple', 'exact', 'golden', 'armijo'
# 推奨デフォルトは armijo
lsmode = 'armijo'

nmaxlsiter = 100

# for golden search
alphaeps = 1.0e-2

# for exact line search
dump = 0.3

# for Armijo backtracking
armijo_c1 = 1.0e-4     # sufficient decrease parameter
armijo_tau = 0.5       # reduction factor
armijo_alpha_min = 1.0e-12

# for global optimization
eps = 1.0e-5
nmaxiter = 200
iprintiterval = 1

# target function
# functype = 'ellipsoid'
functype = 'general'

# graph plot parameters
fplot = 1
ngdata = 51
xgmin = -4.0
xgmax =  4.0
ygmin = -4.0
ygmax =  4.0
tsleep = 0.3


# ============================================================
# command line
# ============================================================

if len(sys.argv) >= 2:
    x0[0] = float(sys.argv[1])
if len(sys.argv) >= 3:
    x0[1] = float(sys.argv[2])
if len(sys.argv) >= 4:
    lsmode = sys.argv[3]
if len(sys.argv) >= 5:
    functype = sys.argv[4]

x0 = np.asarray(x0, dtype=float)

# 基本ステップ幅
if functype == 'ellipsoid':
    alpha = 0.03
else:
    alpha = 0.01

# 線探索を行う場合は、初期 alpha を 1.0 から縮める方が説明しやすい
if lsmode == 'golden':
    alpha = 1.0
elif lsmode == 'armijo':
    alpha = 1.0


# ============================================================
# line search routines
# ============================================================

def linesearch(mode, func, x0, diff, alpha, dump, nmaxiter, eps):
    if mode == '' or mode == 'one':
        return linesearch_one(func, x0, diff, alpha)

    if mode == 'simple':
        return linesearch_simple(func, x0, diff, alpha, nmaxiter)

    if mode == 'exact':
        return linesearch_exact(
            func, x0, diff,
            h=1.0e-3 * alpha,
            dump=dump,
            nmaxiter=1,
            eps=eps,
        )

    if mode == 'golden':
        return linesearch_golden(func, x0, diff, alpha, nmaxiter, eps)

    if mode == 'armijo':
        return linesearch_armijo(
            func, x0, diff,
            alpha0=alpha,
            c1=armijo_c1,
            tau=armijo_tau,
            nmaxiter=nmaxiter,
            alpha_min=armijo_alpha_min,
        )

    raise ValueError("unknown line search mode: {}".format(mode))


def linesearch_one(func, x0, diff, alpha):
    """
    Fixed step:
        x_new = x0 - alpha * grad
    """
    x0 = np.asarray(x0, dtype=float)
    diff = np.asarray(diff, dtype=float)

    dx = -alpha * diff
    return 1, x0 + dx, dx, alpha


def linesearch_simple(func, x0, diff, alpha, nmaxiter):
    """
    Very simple trial search.

    alpha, 2 alpha, 3 alpha, ... と進み、
    f が悪化したら直前の点を採用する。

    教育用の素朴な方法。
    Armijo や golden search より推奨度は低い。
    """
    x0 = np.asarray(x0, dtype=float)
    diff = np.asarray(diff, dtype=float)

    fbest = func(x0)
    xbest = x0.copy()
    dxbest = np.zeros_like(x0)
    alphabest = 0.0

    for i in range(1, nmaxiter + 1):
        alpha2 = i * alpha
        dx = -alpha2 * diff
        x = x0 + dx
        f1 = func(x)

        if f1 > fbest:
            if alphabest == 0.0:
                # 最初の試行から悪化したので、ステップ失敗
                return -1, x0.copy(), np.zeros_like(x0), 0.0
            return i, xbest, dxbest, alphabest

        xbest = x.copy()
        dxbest = dx.copy()
        alphabest = alpha2
        fbest = f1

    return -1, xbest, dxbest, alphabest


def linesearch_armijo(func, x0, diff,
                      alpha0=1.0,
                      c1=1.0e-4,
                      tau=0.5,
                      nmaxiter=50,
                      alpha_min=1.0e-12):
    """
    Armijo backtracking line search for steepest descent.

    Steepest descent direction:
        p = -grad

    Armijo condition:
        f(x + alpha p) <= f(x) + c1 alpha grad^T p

    Since p = -grad,
        grad^T p = -||grad||^2

    Therefore:
        f(x - alpha grad) <= f(x) - c1 alpha ||grad||^2

    注意:
      alpha が小さくなりすぎた場合は、悪いステップを無理に採用せず、
      line search failure としてゼロステップを返す。
    """
    x0 = np.asarray(x0, dtype=float)
    diff = np.asarray(diff, dtype=float)

    grad2 = float(np.inner(diff, diff))
    if grad2 <= 0.0:
        return 0, x0.copy(), np.zeros_like(x0), 0.0

    f0 = func(x0)
    alpha = float(alpha0)

    for i in range(nmaxiter):
        dx = -alpha * diff
        xnew = x0 + dx
        fnew = func(xnew)

        armijo_rhs = f0 - c1 * alpha * grad2

        if fnew <= armijo_rhs:
            return i, xnew, dx, alpha

        alpha *= tau

        if alpha < alpha_min:
            return -1, x0.copy(), np.zeros_like(x0), 0.0

    return -1, x0.copy(), np.zeros_like(x0), 0.0


def linesearch_golden(func, x0, diff, xrange, nmaxiter, eps):
    """
    Golden-section search on alpha in [0, xrange].

    Minimize:
        phi(alpha) = f(x0 - alpha * grad)

    注意:
      厳密には、最小値が [0, xrange] の内部にあることが望ましい。
      講義用の比較として残している。
    """
    x0 = np.asarray(x0, dtype=float)
    diff = np.asarray(diff, dtype=float)

    eta = 0.381966011   # (sqrt(5) - 1) / (sqrt(5) + 1)

    if np.inner(diff, diff) <= 0.0:
        return 0, x0.copy(), np.zeros_like(x0), 0.0

    a = 0.0
    b = float(xrange)

    x1 = a + eta * (b - a)
    x2 = b - eta * (b - a)

    f1 = func(x0 - x1 * diff)
    f2 = func(x0 - x2 * diff)

    alpha2 = 0.5 * (a + b)

    for i in range(nmaxiter):
        if abs(b - a) < eps:
            alpha2 = 0.5 * (a + b)
            dx = -alpha2 * diff
            return i, x0 + dx, dx, alpha2

        if f1 < f2:
            b = x2
            x2 = x1
            f2 = f1
            x1 = a + eta * (b - a)
            f1 = func(x0 - x1 * diff)
        else:
            a = x1
            x1 = x2
            f1 = f2
            x2 = b - eta * (b - a)
            f2 = func(x0 - x2 * diff)

        alpha2 = 0.5 * (a + b)

    dx = -alpha2 * diff
    return -1, x0 + dx, dx, alpha2


def linesearch_exact(func, x0, diff, h, dump, nmaxiter, eps):
    """
    Approximate one-dimensional Newton step along steepest descent direction.

    This is not a true exact line search.
    It estimates the first and second derivative of
        phi(s) = f(x0 - s * grad / ||grad||)
    near s = 0.
    """
    x0 = np.asarray(x0, dtype=float)
    diff = np.asarray(diff, dtype=float)

    dr2 = float(np.inner(diff, diff))
    if dr2 <= 0.0:
        return 0, x0.copy(), np.zeros_like(x0), 0.0

    dr = sqrt(dr2)
    direction = -diff / dr

    for i in range(nmaxiter):
        f0 = func(x0)
        fp = func(x0 + h * direction)
        fm = func(x0 - h * direction)

        # phi'(0), phi''(0)
        f1 = (fp - fm) / (2.0 * h)
        f2 = (fp - 2.0 * f0 + fm) / (h * h)

        step = -f1 / (abs(f2) + dump)
        dx = step * direction
        x0 = x0 + dx

        if np.max(np.abs(dx)) < eps:
            return i, x0, dx, abs(step)

    return -1, x0, dx, abs(step)


# ============================================================
# target functions and gradients
# ============================================================

def ediff1(i, x):
    if i == 0:
        return 2.0 * x[0]
    if i == 1:
        return 18.0 * x[1]
    raise ValueError("bad derivative index")


def gdiff1(i, x):
    if i == 0:
        return (
            -10.0
            - 60.0 * x[0]
            + 4.5 * x[0] * x[0]
            + 12.0 * x[0] * x[0] * x[0]
            + 3.0 * x[1] * x[1]
        )
    if i == 1:
        return (
            30.0
            - 60.0 * x[1]
            + 12.0 * x[1] * x[1] * x[1]
            + 6.0 * x[0] * x[1]
        )
    raise ValueError("bad derivative index")


def efunc(x):
    return x[0] * x[0] + 9.0 * x[1] * x[1]


def gfunc(x):
    return (
        -3.0
        - 10.0 * x[0]
        - 30.0 * x[0] * x[0]
        + 1.5 * x[0] * x[0] * x[0]
        + 3.0 * x[0] * x[0] * x[0] * x[0]
        + 30.0 * x[1]
        - 30.0 * x[1] * x[1]
        + 3.0 * x[1] * x[1] * x[1] * x[1]
        + 3.0 * x[0] * x[1] * x[1]
    )


def gradient_vector(diff1func, x):
    x = np.asarray(x, dtype=float)
    g = np.empty(len(x), dtype=float)

    for i in range(len(x)):
        g[i] = diff1func(i, x)

    return g


# global function variables
if functype == 'ellipsoid':
    diff1 = ediff1
    func = efunc
else:
    diff1 = gdiff1
    func = gfunc


# ============================================================
# main
# ============================================================

def main():
    global x0
    global func, diff1

    print("Find minimum point by steepest descent method")
    print("line search mode =", lsmode)
    print("function type    =", functype)
    print("")

    # ------------------------------------------------------------
    # plot surface graph
    # ------------------------------------------------------------
    if fplot == 1:
        xs = np.linspace(xgmin, xgmax, ngdata)
        ys = np.linspace(ygmin, ygmax, ngdata)

        xg, yg = np.meshgrid(xs, ys)
        zg = func([xg, yg])

        fig = plt.figure(figsize=(10, 5))

        ax = fig.add_subplot(1, 2, 1)
        ax2 = fig.add_subplot(1, 2, 2, projection='3d')

        ax2.plot_wireframe(xg, yg, zg, rstride=2, cstride=2)
        ax2.set_title('surface')
        ax2.set_xlabel('x')
        ax2.set_ylabel('y')
        ax2.set_zlabel('f')

        ax.set_title('contour')
        ax.contourf(xg, yg, zg, levels=51, cmap='Spectral')
        ax.set_xlabel('x')
        ax.set_ylabel('y')
        ax.set_aspect('equal', adjustable='box')

        solve, = ax.plot([], [], color='blue', marker='o', linestyle='-')

        plt.pause(0.1)

    # ------------------------------------------------------------
    # calculate initial parameters
    # ------------------------------------------------------------
    f = func(x0)

    print("x0 = ({}, {}): f = {}".format(x0[0], x0[1], f))

    xt = [x0[0]]
    yt = [x0[1]]

    # ------------------------------------------------------------
    # optimization start
    # ------------------------------------------------------------
    for iter in range(nmaxiter):
        grad = gradient_vector(diff1, x0)
        gradnorm = sqrt(float(np.inner(grad, grad)))

        if gradnorm < eps:
            print("Converged by gradient norm at x = ", x0, " |grad| = {}".format(gradnorm))
            break

        insiter, xnew, dx, alpha2 = linesearch(
            lsmode,
            func,
            x0,
            grad,
            alpha,
            dump,
            nmaxlsiter,
            alphaeps,
        )

        dxmax = float(np.max(np.abs(dx)))

        if insiter < 0 and dxmax == 0.0:
            print("Line search failed at iteration {}".format(iter))
            print("x = ({:12.6g}, {:12.6g}) f = {:12.6g} |grad| = {:12.6g}".format(
                x0[0], x0[1], func(x0), gradnorm
            ))
            break

        x0 = xnew
        f = func(x0)

        print(
            "{:4d}({:3d}): x = ({:12.6g}, {:12.6g}) "
            "alpha = {:12.6g} dx = {:12.6g} |grad| = {:12.6g}: f = {:12.6g}".format(
                iter,
                insiter,
                x0[0],
                x0[1],
                alpha2,
                dxmax,
                gradnorm,
                f,
            )
        )

        if fplot == 1:
            xt.append(x0[0])
            yt.append(x0[1])
            solve.set_data(xt, yt)
            plt.pause(0.1)
            time.sleep(tsleep)

        if dxmax < eps:
            print("Converged by step size at x = ", x0, " dx = {}   f = {}".format(dxmax, f))
            break

    else:
        print("Not converged")

    if fplot == 1:
        print("Press enter to terminate:", end='')
        input()


if __name__ == "__main__":
    try:
        main()
    except Exception:
        traceback.print_exc()
        sys.exit(1)