import csv
import numpy as np
from numpy import sin, cos, tan, pi, exp
import matplotlib.pyplot as plt
import time
import sys

# optimization parameters
x0    = -2.0
if len(sys.argv) >= 2:
    x0 = float(sys.argv[1])

dump  = 0.0
eps   = 1.0e-10
nmaxiter = 200
iprintiterval = 1

# differentiation x mesh
h = 0.01

# graph plot parameters
ngdata = 51
xgmin = -4.0
xgmax =  4.0
tsleep = 1.0

# second derivative of func(x)
def diff2(x):
#    return exp(x)
#    return -60.0 + 9.0 * x + 36.0 * x*x
    return (func(x+h) - 2.0 * func(x) + func(x-h)) / h / h

# first derivative of func(x)
def diff1(x):
#    return exp(x) - 3.0
#    return 1.0 - 60.0 * x + 4.5 * x*x + 12.0 * x*x*x
    return (func(x+h) - func(x-h)) / 2.0 / h

def func(x):
#    return exp(x) - 3.0 * x
    return -3.0 + x - 30.0 * x*x + 1.5 * x*x*x + 3.0 * x*x*x*x

def main():
    global plt, ngdata, xgmin, xgmax
    global x0, dump, eps, nmaxiter, iprintinterval

    print("Find minimum / maximum point by Newton-Raphson method")
    print("")

    xg = []
    yg = []
    yz = [0.0] * ngdata
    xgstep = (xgmax - xgmin) / (ngdata - 1)
    for i in range(ngdata):
        xg.append(xgmin + i * xgstep)
        yg.append(func(xg[i]))

    fig, ax = plt.subplots(1, 1)
    plt.title("Find minimum / maximum")
    plt.xlabel("x")
    plt.ylabel("f(x)")
    ax.set_xlim([xgmin, xgmax])
    ax.set_ylim([min(yg), max(yg)])
    data, = ax.plot(xg, yg, color = 'black', linewidth = 0.3)
    yzero, = ax.plot(xg, yz, color = 'red', linewidth = 0.3)
    solve, = ax.plot([], color = 'blue', marker = 'o', linestyle = '')
    plt.pause(0.1)
#    time.sleep(3.0)

    x = x0
    xt = []
    yt = []
    for i in range(nmaxiter):
        f  = func(x)
        f1 = diff1(x)
        f2 = diff2(x)
        if f2 < 0.0:
            f2 -= dump
        else:
            f2 += dump

        xnext = x - f1 / f2
        dx = xnext - x
        if i % iprintiterval == 0:
            print("Iter {:5d}: x: {:>16.12f} => {:>16.12f}, dx = {:>10.4g}   f={:>16.12f}".
                    format(i, x, xnext, dx, f))
        if abs(dx) < eps:
            print("  Success: Convergence reached: dx = {} < eps = {}".format(dx, eps))
            break

        xt.append(x)
        yt.append(f)
        solve.set_data(xt, yt)
        plt.pause(0.1)
        time.sleep(tsleep)

        x = xnext
    else:
        print("  Failed: Convergence did not reach: dx = {} > eps = {}".format(dx, eps))

    print("Press enter to terminate:", end = '')
    ret = input()


if __name__ == "__main__":
    main()
