import csv
import numpy as np
from numpy import sin, cos, tan, pi, exp
import sys

# data parameters
x0 = 0.0
x1 = 5.0
ndata = 101
xstep = (x1 - x0) / (ndata - 1)

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

dump  = 0.1
eps   = 1.0e-5
nmaxiter = 100
iprintinterval = 10

# d func (xi) / d ci
def diff1(i, x, c):
    if i == 0:
        return 1.0 / (1.0 + c[1] * x)
    if i == 1:
        return -c[0] * x / pow(1.0 + c[1] * x, 2)

def func(x, c):
    return c[0] / (1.0 + c[1] * x)

def y(x):
    return func(x, [1.0, 1.0]) + 0.0 * np.random.rand();

# minimize S2 = sum[ (func(xi) - yi) ]^2
def S2(x, y, c, func):
    sum = 0.0
    for k in range(len(x)):
        d = func(x[k], c) - y[k]
        sum += d * d
    return sum


def main():
    global c0, dump, eps, nmaxiter, iprintinterval

    print("Find minimum / maximum point by Marquart method")
    print("")

    xd = []
    yd = []
    print("input data:")
    for i in range(ndata):
        xd.append(x0 + i * xstep)
        yd.append(y(xd[i]))
        print("  {:12.4g}  {:12.4g}".format(xd[i], yd[i]))
    print("")

    n = len(c0)
    Afi  = np.empty([n, 1])
    AAij = np.empty([n, n])

# calculate initial parameters
    sum = S2(xd, yd, c0, func)
    print("c0 = ({}, {}): f = {}".format(c0[0], c0[1], sum))

# optimization start
    for iter in range(nmaxiter):
        for i in range(0, n):
            Afi[i, 0] = 0.0
            for k in range(ndata):
                Afi[i, 0] += diff1(i, xd[k], c0) * (func(xd[k], c0) - yd[k])
            for j in range(i, n):
                AAij[i, j] = 0.0
                for k in range(ndata):
                    AAij[i, j] += diff1(i, xd[k], c0) * diff1(j, xd[k], c0)
                AAij[j, i] = AAij[i, j]

#        print("At*f=", Afi)
#        print("At*A=", AAij)
#        print("")
        tr = np.trace(AAij)
        AAij += dump * tr
#        AAij = 1000 * np.eye(n)
#        dc = -np.linalg.inv(AAij) @ Afi
        dc = -np.linalg.solve(AAij, Afi)
        dc = dc.transpose()[0]
        c0 = c0 + dc
        dcmax = max(abs(dc))
        sum = S2(xd, yd, c0, func)
        if iter % iprintinterval ==0:
            print("   dc=", dc)
            print("{:4d}: c = ({:12.6g}, {:12.6g}) dcmax = {:12.6g}  S = {:12.6g}".format(
                    iter, c0[0], c0[1], dcmax, sum))

        if dcmax < eps:
            print("Converged at c = ", c0, " dcmax = {}".format(dcmax))
            break;
    else:
        print("Not converged")


if __name__ == "__main__":
    main()