import sys
import csv
import numpy as np
from math import exp
from scipy.optimize import minimize
import matplotlib.pyplot as plt


#nelder-mead    Downhill simplex
#powell         Modified Powell
#cg             conjugate gradient (Polak-Ribiere method)
#bfgs           BFGS法
#newton-cg      Newton-CG
#trust-ncg      信頼領域 Newton-CG 法
#dogleg         信頼領域 dog-leg 法
#L-BFGS-B’ (see here)
#TNC’ (see here)
#COBYLA’ (see here)
#SLSQP’ (see here)
#trust-constr’(see here)
#dogleg’ (see here)
#trust-exact’ (see here)
#trust-krylov’ (see here)
method = "cg"

maxiter = 100
tol = 1.0e-5

#==========================================
# Source parameters to be fitted
#==========================================
# Data parameters: I0, x0, w
peaks = [1.1, 0.4, 0.4]
nparams = len(peaks)

xrange = [-1.0, 2.0]
nx     = 101
xstep  = (xrange[1] - xrange[0]) / (nx - 1)
xd = []
for i in range(nx):
   xd.append(xrange[0] + i * xstep)
yd = []

# optimization parameters
x0    = [1.3,  0.6, 0.1]


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

#==========================================
# File configurations
#==========================================
initial_csv = 'initial.csv'
final_csv   = 'final.csv'
conv_csv    = 'convergence.csv'


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


#==========================================
# functions
#==========================================
def save_csv(path, headerlist, datalist, is_print = 0):
    f = open(path, 'w')
    if not f:
        return 0

#    print("len=", len(datalist), len(datalist[0]))
    writer = csv.writer(f, delimiter=',', lineterminator='\n')
    writer.writerow(headerlist)
    for i in range(len(datalist[0])):
        dlist = [datalist[id][i] for id in range(len(datalist))]
        dliststr = joinf(dlist, "%12.8g", ", ")
        if is_print:
            print("  {:3d}: {}".format(i, dliststr))
        writer.writerow(dlist)
    f.close()
    return 1

def joinf(list, format, sep):
    s = format % (list[0])
    for i in range(1, len(list)):
        s += sep + format % (list[i])
    return s


def ycal(x, params):
    ret = 0.0
    for ip in range(0, len(params), 3):
        I0 = params[ip]
        x0 = params[ip+1]
        w  = params[ip+2]
        if w < 0.001:
            w = 0.001
        a = 0.832554611 / w
        X = a * (x - x0)
        ret += I0 * exp(-X * X)
    return ret

def ycal_list(xd, params):
    print("pa=", params)
    y = []
    for i in range(len(xd)):
        y.append(ycal(xd[i], params))
    return y

def CalS2(params):
    global xd, yd, nx
    S2 = 0.0
    for i in range(nx):
        yc = ycal(xd[i], params)
        d  = yd[i] - yc
        S2 += d * d
    return S2

def diff1i(i, params):
    global xd, yd, nx

    ip  = i // 3
    idx = i % 3
    I0 = params[ip*3]
    x0 = params[ip*3+1]
    w  = params[ip*3+2]

    ret = 0.0
    for id in range(nx):
        if w < 0.001:
            w = 0.001

        yc = ycal(xd[id], [I0, x0, w])
        dy = yd[id] - yc

        a  = 0.832554611 / w
        dx = a * (xd[id] - x0)
        e  = exp(-dx*dx)
        if idx == 0: # I0
            dif = e
        elif idx == 1: # x0
            dif = 2.0 * a * dx * I0 * e
        elif idx == 2: # w
            dif = 2.0 * dx * dx / w * I0 * e
        else:
            print("Error in diff1i: Invalid index (ivar={}, ipeak={}, diff_idx={})".format(i, ip, idx))
            exit()
        d = -2.0 * dif * dy
        ret += d

    return ret

def diff1(params):
    global nparams
    df = np.empty(len(params))
    for i in range(len(params)):
        df[i] = diff1i(i, params)
    return df

iter = 0
xiter  = []
yfmin  = []
figure   = None
ax_fit   = None
rawdata  = None
inidata  = None
fitdata  = None
ax_conv  = None
convdata = None
def callback(xk):
    global xd
    global iter
    global figure, ax_fit, ax_conv
    global fitdata, convdata
    
    fmin = CalS2(xk)
    print("callback {}: xk={}".format(iter, xk))
    print("   fmin={}".format(fmin))
    iter += 1
    xiter.append(iter)
    yfmin.append(fmin)

    convdata[0].set_data(xiter, yfmin)
    ax_conv.set_xlim((0.0, max(xiter) + 1.0))
    ax_conv.set_ylim((min(yfmin) * 0.8, max(yfmin) * 1.2))

    yc = ycal_list(xd, xk)
    fitdata[0].set_data(xd, yc)

    plt.pause(0.2)


#==========================================
# Main routine
#==========================================
def main():
    global xd, yd, xiter, yfmin
    global figure, ax_fit, ax_conv
    global fitdata, convdata

    print("")
    print("Peak fitting by python scipy.optimize")

    yd   = ycal_list(xd, peaks)
    yini = ycal_list(xd, x0)

    print("Initial data are saved to [{}]".format(initial_csv))
    ret = save_csv(initial_csv, ['x', 'y(data)', 'y(ini)'], [xd, yd, yini])
    if ret == 0:
        print("Error: Can not write to [{}]".format(initial_csv))


    if fplot == 1:
        figure = plt.figure(figsize = (10, 5))
        ax_fit  = figure.add_subplot(1, 2, 1)
        rawdata = ax_fit.plot(xd, yd, color = 'blue', linestyle = '', linewidth = 0.5,
                            fillstyle = 'full', marker = 'x', markersize = 5)
        inidata = ax_fit.plot(xd, yini, color = 'black', linestyle = 'dashed', linewidth = 0.5)
        fitdata = ax_fit.plot([], [], color = 'red', linestyle = '-', linewidth = 0.5)

        ax_conv = figure.add_subplot(1, 2, 2)
        ax_conv.set_yscale('log')
        convdata  = ax_conv.plot([], [], color = 'black', linestyle = '-', linewidth = 0.5,
                            fillstyle = 'full', marker = 'o', markersize = 5)
        plt.pause(0.001)


    res = minimize(CalS2, x0, jac=diff1, method=method, tol = tol, callback = callback,
                options = {'maxiter':maxiter, "disp":True})
    print(res)


    yopt = ycal_list(xd, res.x)
    print("Final data are saved to [{}]".format(final_csv))
    ret = save_csv(final_csv, ['x', 'y(data)', 'y(ini)', 'y(opt)'], [xd, yd, yini, yopt])
    if ret == 0:
        print("Error: Can not write to [{}]".format(final_csv))

    print("Convergence data are saved to [{}]".format(conv_csv))
    ret = save_csv(conv_csv, ['iter', 'error'], [xiter, yfmin])
    if ret == 0:
        print("Error: Can not write to [{}]".format(conv_csv))


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


if __name__ == "__main__":
    main()
