import csv
import numpy as np
from numpy import sin, cos, tan, pi, exp
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d.axes3d import Axes3D
from matplotlib import cm
import time
import sys

# optimization parameters
x0 = [0.0, 0.0]
if len(sys.argv) >= 2:
    x0[0] = float(sys.argv[1])
if len(sys.argv) >= 3:
    x0[1] = float(sys.argv[2])

# target function
functype = 'ellipsoid'

if functype == 'ellipsoid':
   alpha = 0.03
else:
   alpha = 0.01

eps   = 1.0e-5
nmaxiter = 200
iprintiterval = 1

# differentiation x mesh
h = 0.01

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

# first derivative of func(x)
def ediff1(i, x):
    if i == 0:
        return 2.0 * x[0]
    if i == 1:
        return 18.0 * x[1]

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]

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]

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


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

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

# plot surface graph
    if fplot == 1:
        xgstep = (xgmax - xgmin) / (ngdata - 1)
        ygstep = (ygmax - ygmin) / (ngdata - 1)
        xg = np.empty([ngdata, ngdata])
        yg = np.empty([ngdata, ngdata])
        zg = np.empty([ngdata, ngdata])
        for ix in range(ngdata):
            for iy in range(ngdata):
                xg[ix, iy] = xgmin + ix * xgstep
                yg[ix, iy] = ygmin + iy * ygstep
                zg[ix, iy] = func([xg[ix][iy], yg[ix][iy]])

        fig = plt.figure(figsize = (10, 5))
        ax  = fig.add_subplot(1, 2, 1)
        ax2 = fig.add_subplot(1, 2, 2, projection='3d')
#        ax = fig.add_subplot(111, projection='3d')
#        surf = ax.plot_surface(xg, yg, zg, cmap=cm.coolwarm)
        surf = ax2.plot_wireframe(xg, yg, zg, rstride=2, cstride=2)
#        ax.set_zlim3d(-3.01, 3.01)
#        plt.colorbar(surf, shrink=0.5, aspect=10)
        ax.set_title('contour')
        contour = ax.contourf(xg, yg, zg, levels = 51, cmap = 'Spectral')
#        contour = ax.contour(xg, yg, zg, levels = 51, cmap = 'Spectral')
        solve, = ax.plot([], color = 'blue', marker = 'o', linestyle = '-')
#        solve, = ax.plot([], color = 'blue', marker = 'o', linestyle = '')
        plt.pause(0.1)
#        plt.show()

    n = len(x0)
    Si = np.empty(n)

# calculate initial parameters
    f = func(x0)
    print("x0 = ({}, {}): f = {}".format(x0[0], x0[1], f))
    xt = []
    yt = []
    xt.append(x0[0])
    yt.append(x0[1])

# optimization start
    for iter in range(nmaxiter):
        for i in range(0, n):
            Si[i] = diff1(i, x0)

        dx = -alpha * Si
        print("   dx=", dx)
        x0 = x0 + dx
        f = func(x0)
        dxmax = max(abs(dx))
        print("{:4d}: x = ({:12.6g}, {:12.6g}) dx = {:12.6g}: f = {:12.6g}".format(iter, x0[0], x0[1], dxmax, 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 at x = ", x0, " dx = {}   f = {}".format(dxmax, f))
            break;
    else:
        print("Not converged")

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


if __name__ == "__main__":
    main()

