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


mode = 'sim' # sim|fit

infile = "peak.xlsx"

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


#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

iprint_interval = 5

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


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


argv = sys.argv
nargs = len(argv)
if nargs >= 2: mode = argv[1]
if nargs >= 3: infile = argv[2]
if nargs >= 4: method = argv[3]
if nargs >= 5: x0[0] = float(argv[4])
if nargs >= 6: x0[1] = float(argv[5])
if nargs >= 7: x0[2] = float(argv[6])
if nargs >= 8: iprint_interval = int(argv[7])


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

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

#    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 True

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

    return ret

def ycal_list(xin, xk):
    print("parameters=", xk)
    y = []
    for i in range(len(xin)):
        y.append(ycal(xin[i], xk))

    return y

def minimize_func(xk, xin, yin):
    nx = len(xin)

    S2 = 0.0
    for i in range(nx):
        yc = ycal(xin[i], xk)
        d  = yin[i] - yc
        S2 += d * d

    return S2

def diff1i(i, xk, xin, yin):
    nx = len(xin)

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

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

        yc = ycal(xin[id], [I0, x0, w])
        dy = yin[id] - yc

        a  = 0.832554611 / w
        dx = a * (xin[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(xk, xin, yin):
    nparams = len(xk)

    df = np.empty(nparams)
    for i in range(nparams):
        df[i] = diff1i(i, xk, xin, yin)

    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, xin, yin):
    global iter
    global figure, ax_fit, ax_conv
    global fitdata, convdata
    
    fmin = minimize_func(xk, xin, yin)
    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(xin, xk)
    fitdata[0].set_data(xin, yc)

    plt.pause(0.2)

def read_file(path):
    df = pd.read_excel(path)
    x = df['x'].values
    y = df['y'].values
    return x, y

#==========================================
# Main routines
#==========================================
       

def sim():
    print()
    print("Simulation:")

    print(f"Read data from [{infile}]")
    xin, yin = read_file(infile)

    print()
    print("Simulate using initial parameters:", x0)
    ysim = ycal_list(xin, x0)
    print(f"{'xin':10}\t{'yin':10}\t{'ycal':10}")
    for _xin, _yin, _ysim in zip(xin, yin, ysim):
        print(f"{_xin:10.4g}\t{_yin:10.4g}\t{_ysim:10.4g}")

    print()
    print("Plot")
    plt.plot(xin, yin, label='y')
    plt.plot(xin, ysim, label='ysim')
    plt.legend()
    plt.pause(1.0e-4)
    input("\nPress ENTER to terminate:")


def fit():
    global xiter, yfmin
    global figure, ax_fit, ax_conv
    global fitdata, convdata

    print()
    print("Fitting:")

    print(f"Read data from [{infile}]")
    xin, yin = read_file(infile)

    print()
    print("Simulate initial curve:", x0)
    yini = ycal_list(xin, x0)
    print(f"{'xin':10}\t{'yin':10}\t{'yini':10}")
    for _xin, _yin, _yini in zip(xin, yin, yini):
        print(f"{_xin:10.4g}\t{_yin:10.4g}\t{_yini:10.4g}")


    print("Initial data are saved to [{}]".format(initial_csv))
    ret = save_csv(initial_csv, ['x(in)', 'y(data)', 'y(ini)'], [xin, yin, 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(xin, yin, color = 'blue', linestyle = '', linewidth = 0.5,
                            fillstyle = 'full', marker = 'x', markersize = 5)
        inidata = ax_fit.plot(xin, 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(lambda xk: minimize_func(xk, xin, yin), x0, 
                   jac = lambda xk: diff1(xk, xin, yin), 
                   method = method, tol = tol, 
                   callback = lambda xk: callback(xk, xin, yin),
                   options = {'maxiter':maxiter, "disp":True})
    print(res)


    yopt = ycal_list(xin, res.x)
    print("Final data are saved to [{}]".format(final_csv))
    ret = save_csv(final_csv, ['x', 'y(data)', 'y(ini)', 'y(opt)'], [xin, yin, 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()


def main():
    print()
    print(f"{mode=}")
    print(f"{infile=}")
    print(f"{method=}")
    print(f"x0=", x0)
    print(f"{iprint_interval=}")

    if mode == 'sim': sim()
    elif mode == 'fit': fit()
    else:
        print(f"\nError: Invalide mode [{mode}]\n")
        exit()


if __name__ == "__main__":
    main()
