import sys
import numpy as np
import pandas as pd
import scipy.signal
import matplotlib.pyplot as plt
import matplotlib.widgets as wg


from tklib.tkutils import replace_path, getarg, getintarg, getfloatarg, pint, pfloat
from tklib.tkparams import tkParams
from tklib.tkvariousdata import tkVariousData
from tklib.tkapplication import tkApplication


"""
Smoothing by different methods and parameters
"""


#=============================
# parameters
#=============================
def usage(app):
    print("usage: python {} args".format(sys.argv[0]))

def initialize(app):
    app.cparams = tkParams()
    cparams = app.cparams
    cparams.infile = 'xrd.xlsx'
    cparams.mode = 'test'
    cparams.xmin = None
    cparams.xmax = None
    cparams.xlabel = 0
    cparams.ylabel = 1
    
    cparams.norder = 2
    cparams.nsmooth = 11
    cparams.nincrement = 10
    cparams.ndifferential = 0

    cparams.do_resampling  = 0
    cparams.sampling_step  = 0.02

    cparams.template = "StandardGraph.xlsm"

    cparams.figsize = [6, 6]
    cparams.fontsize   = 16
    cparams.fontsize_s = 10

def update_vars(app):
    cparams = app.cparams

    argv = sys.argv
    n = len(argv)
    if n >= 2:
        cparams.mode = argv[1]
    if n >= 3:
        cparams.infile = argv[2]
    if n >= 4:
        cparams.xmin = pfloat(argv[3], defval = argv[3])
    if n >= 5:
        cparams.xmax = pfloat(argv[4], defval = argv[4])
    if n >= 6:
        cparams.xlabel = pint(argv[5], defval = argv[5])
    if n >= 7:
        cparams.ylabel = pint(argv[6], defval = argv[6])
    if n >= 8:
        cparams.method = argv[7]
    if n >= 9:
        cparams.norder  = pint(argv[8], cparams.norder)
    if n >= 10:
        cparams.nsmooth = pint(argv[9], cparams.nsmooth)
    if n >= 11:
        cparams.nincrement = pint(argv[10], defval = cparams.nincrement)
    if n >= 12:
        cparams.ndifferential  = pint(argv[11], cparams.ndifferential)
    if n >= 13:
        cparams.do_resampling  = pint(argv[12], cparams.do_resampling)
    if n >= 14:
        cparams.sampling_step  = pfloat(argv[13], cparams.sampling_step)
    
def diff2forward(x, y, i):
    n = len(x)
    if(0 <= i <= n-2):
        return (y[i+1] - y[i]) / (x[i+1] - x[i]);
    return ''

def diff2backward(x, y, i):
    n = len(x)
    if(1 <= i <= n-1):
        return (y[i] - y[i-1]) / (x[i] - x[i-1]);
    return ''

def diff3(x, y, i):
    n = len(x)
    if(1 <= i <= n-2):
        return (y[i+1] - y[i-1]) / 2.0 / (x[i] - x[i-1]);
    return ''

def diff5(x, y, i):
    n = len(x)
    if(2 <= i <= n-3):
        return (-y[i+2] + 8.0*y[i+1] - 8.0*y[i-1] + y[i-2]) / 12.0 / (x[i] - x[i-1]);
    return ''

def diff7(x, y, i):
    n = len(x)
    if(3 <= i <= n-4):
        return (y[i+3] - 9.0*y[i+2] + 45.0*y[i+1] 
                - 45.0*y[i-1] + 9.0*y[i-2] - y[i-3]) / 60.0 / (x[i] - x[i-1]);
    return ''

def SmoothingBySimpleAverage(y, n):
    n2 = int(n / 2);
    ndata = len(y);
    ys = []
    for i in range(0, ndata):
        c = 0;
        ys.append(0.0);
        for k in range(i - n2, i + n2 + 1):
            if k < 0 or k >= ndata:
                continue
            ys[i] += y[k]
            c += 1
        if c > 0:
            ys[i] /= c;
        else:
            ys[i] = y[i]
    return ys;


def SmoothingByPolynomialFit(y, n):
    m = int(n / 2);
    W23 = (4.0 * m * m - 1.0) * (2.0 * m + 3.0) / 3.0;
    w23j = [0.0]*n
    for j in range(-m, m+1):
        w23j[j + m] = (3.0 * m * (m+1.0) - 1.0 - 5.0 * j * j) / W23

    ndata = len(y)
    ys = []
    for i in range(0, ndata):
        c = 0.0;
        ys.append(0.0);
        for j in range(-m, m+1):
            k = i + j
            if k < 0 or k >= ndata:
                continue
            ys[i] += w23j[j+m] * y[k]
            c += w23j[j+m]
        if c > 0:
            ys[i] /= c
        else:
            ys[i] = y[i]
    return ys;


def test(app):
    cparams = app.cparams

    cparams.logfile = app.replace_path(cparams.infile)
    cparams.outfile = app.replace_path(cparams.infile, template = ["{dirname}", "{filebody}-smoothed.xlss"])
    print(f"Open logfile [{cparams.logfile}]")
    app.redirect(targets = ["stdout", cparams.logfile], mode = 'w')

    print("")
    print(f"Smoothing data in [{cparams.infile}]")
    print("infile  : ", cparams.infile)
    print("mode    : ", cparams.mode)
    print("xlabel  : ", cparams.xlabel)
    print("ylabel  : ", cparams.ylabel)
    print("  x range : ", cparams.xmin, cparams.xmax)
    print("norder    : ", cparams.norder)
    print("nsmooth   : ", cparams.nsmooth)
    print("nincrement: ", cparams.nincrement)
    print("order of differentiation: ", cparams.ndifferential)
    print(f"do_resampling: {cparams.do_resampling}")
    print(f"sampling_step: {cparams.sampling_step}")
    
    if '*** choose ***' in cparams.xlabel:
        app.terminate(f"\nError: Choose x label\nTerminated.\n", usage = usage, pause = True)
    if '*** choose ***' in cparams.ylabel:
        app.terminate(f"\nError: Choose y label\nTerminated.\n", usage = usage, pause = True)

    print("Read data from [{}]".format(cparams.infile))
    datafile = tkVariousData(cparams.infile)
    header, datalist = datafile.Read_minimum_matrix(close_fp = True)#, usage = usage)
    xlabel, _xin = datafile.FindDataArray(cparams.xlabel, flag = 'i')
    ylabel, _yin = datafile.FindDataArray(cparams.ylabel, flag = 'i')
    if _xin is None:
        app.terminate(f"Error: Can not find the data with [{cparams.xlabel}]", usage = usage, pause = True)
    if _yin is None:
        app.terminate(f"Error: Can not find the data with [{cparams.xlabel}]", usage = usage, pause = True)

    x = []
    y = []
    for i in range(len(_xin)):
        if not (cparams.xmin is None or cparams.xmin == '' or cparams.xmin == '*') and _xin[i] < cparams.xmin:
            continue
        if not (cparams.xmax is None or cparams.xmax == '' or cparams.xmax == '*') and cparams.xmax < _xin[i]:
            continue
        x.append(_xin[i])
        y.append(_yin[i])
    ndata = len(x)
    print("  ndata = ", ndata)

#=============================
# prepare graph
#=============================
    print("")
    print("plot")
    if cparams.mode == 'test':
        fig, ax = plt.subplots(3, 3, figsize = cparams.figsize)
        ax = ax.flatten()
        markersize = 0.3
        linewidth = 0.5
        fontsize = cparams.fontsize_s
    else:
        fig, ax = plt.subplots(1, 1, figsize = cparams.figsize)
        ax = [ax]
        markersize = 1.0
        linewidth = 1.0
        fontsize = cparams.fontsize

    for axis in ax:
        axis.tick_params(labelsize = fontsize)

    def save_data(method, norder, nsmooth, ndifferential, x, y, ys, xret, yret):
        ext_dict = {"method": method, "norder": norder, "nsmooth": nsmooth, "ndifferential": ndifferential}
        if ndifferential == 0:
            if method == 'simple':
                outfile = app.replace_path(cparams.infile, template = ["{dirname}", 
                                "{filebody}-{method}-{nsmooth}p-order{norder}-smoothed.xlsm"], 
                                ext_dict = ext_dict)
            else:
                outfile = app.replace_path(cparams.infile, template = ["{dirname}", 
                                "{filebody}-{method}-{nsmooth}p-smoothed.xlsm"], 
                                ext_dict = ext_dict)
        else:
            if method == 'simple':
                outfile = app.replace_path(cparams.infile, template = ["{dirname}", 
                                "{filebody}-{method}-{nsmooth}p-order{norder}-smoothed-{ndifferential}-diff.xlsm"], 
                                ext_dict = ext_dict)
            else:
                outfile = app.replace_path(cparams.infile, template = ["{dirname}", 
                                "{filebody}-{method}-{nsmooth}p-smoothed-{ndifferential}-diff.xlsm"], 
                                ext_dict = ext_dict)

        print("Save calculated data to [{}]".format(outfile))
#        df = pd.DataFrame(np.array([x, y, ys]).T, columns = ['x', 'y(input)', 'y(smoothen)'])
#        df.to_excel(outfile, index = False)
        tkVariousData().to_excel(outfile, 
                                ['x', 'y(input)', 'y(smoothen)', 'x(resampled)', 'y(resampled)'], 
                                [x, y, ys, xret, yret], template = cparams.template)

    def execute_smooth(idx, method, norder, nsmooth, ndifferential, x, y, cparams):
        print("")
        print(f"Smoothing data #{idx+1} by {method} for order={norder} with {nsmooth} points...")
        ys = np.zeros(len(x))
        if method == 'simple':
            print("Smoothing by simple moving average")
            print(f"nsmooth={nsmooth}")
#            ys = SmoothingBySimpleAverage(y, nsmooth)
            if ndifferential == 0:
#                ys = SmoothingBySimpleAverage(y, nsmooth)
                print("y=", y)
                w = np.ones(nsmooth) / nsmooth
                ys = np.convolve(y, w, mode = 'same')
            else:
                print("")
                print("Differentiation is not implemented for simple moving average")
                print("Use polynomial smoothing instead")
                print("")
                return ys
                
            ax[idx].set_title(f"{nsmooth}p simple", fontsize = fontsize)
        elif method == 'polynomial':
            print("Smoothing by polynomial fit")
            print(f"nsmooth={nsmooth}  norder={norder}")
            if nsmooth <= norder:
                print(f"nsmooth={nsmooth} must be larger than norder={norder}. Skip.")
                return None
            else:
#                ys = SmoothingByPolynomialFit(y, nsmooth)
                ys = scipy.signal.savgol_filter(y, nsmooth, norder, deriv = ndifferential)
                if ndifferential >= 1:
                    ys /= pow(x[1] - x[0], ndifferential)
                ax[idx].set_title(f"{nsmooth}p {norder}-th order polynomial", fontsize = fontsize)

        if cparams.do_resampling:
            print(f"  Resampling...")
            #x, ysデータを補間し、xret[0] = x[0]からcparams.sampling_stepごとに、x[-1]を超えない範囲でデータを収集
            xret = np.arange(x[0], x[-1] + (x[1] - x[0]) * 0.01, cparams.sampling_step)    
            yret = np.interp(xret, x, ys)
        else:
            xret = x
            yret = ys

        print()
        print(f"{'x':10}\t{'y':10}\t{'ys':10}")
        for i in range(0, ndata):
            print(f"{x[i]:10.4g}\t{y[i]:10.4g}\t{ys[i]:10.4g}")

        print()
        print(f"{'xret':10}\t{'yret':10}")
        for i in range(0, len(xret)):
            print(f"{xret[i]:10.4g}\t{yret[i]:10.4g}")

#        save_data(method, norder, nsmooth, ndifferential, x, y, ys)
        ax[idx].plot(x, y,  linestyle = 'none', marker = 'o', markersize = markersize)
        if ndifferential == 0:
            ax[idx].plot(xret, yret, label = "{}p {}".format(nsmooth, method), linewidth = linewidth)
        else:
            ax2 = ax[idx].twinx()
            ax2.plot(xret, yret, label = "{}p {}".format(nsmooth, method), linewidth = linewidth, color = 'red')
            if ndifferential <= 3:
                s = ['First', 'Second', 'Third'][ndifferential - 1]
            else:
                s = f'{ndifferential}-th'
                
            ax2.set_ylabel(f"{s} differential", fontsize = fontsize)

        return x, ys, xret, yret

    tasks = []
    if cparams.mode == 'test':
        for i in range(3):
            tasks.append(['simple', cparams.norder, cparams.nsmooth + i * cparams.nincrement, cparams.ndifferential])
        for i in range(3):
            tasks.append(['polynomial', cparams.norder, cparams.nsmooth + i * cparams.nincrement, cparams.ndifferential])
        for i in range(3):
            tasks.append(['polynomial', cparams.norder + 2, cparams.nsmooth + i * cparams.nincrement, cparams.ndifferential])

        for i in range(len(tasks)):
            print(f"Smoothing with condition #{i+1}...")
            execute_smooth(i, tasks[i][0], tasks[i][1], tasks[i][2], tasks[i][3], x, y, cparams)
    else:
        tasks = [[cparams.method, cparams.norder, cparams.nsmooth, cparams.ndifferential]]
        print(f"Smoothing...")
        xs, ys, xret, yret = execute_smooth(0, tasks[0][0], tasks[0][1], tasks[0][2], tasks[0][3], x, y, cparams)
        ax[0].set_xlabel(xlabel, fontsize = fontsize)
        ax[0].set_ylabel(ylabel, fontsize = fontsize)
        save_data(cparams.method, cparams.norder, cparams.nsmooth, cparams.ndifferential, x, y, ys, xret, yret)

 
    plt.tight_layout()

    plt.pause(0.1)
    app.terminate("Press ENTER to exit>>", usage = usage, pause = True)

    exit()


def main(app):
    cparams = app.cparams

    if cparams.mode == 'test':
        test(app)
    elif cparams.mode == 'plot':
        test(app)
    else:
        app.terminate(f"\nError: Invalide mode=[{cparams.mode}]\n", usage = usage, pause = True)
#        print("")
#        print(f"Error: Invalide mode=[{cparams.mode}]")
#        print("")
#        exit()


if __name__ == '__main__':
    app = tkApplication()

    print(f"Initialize parameters")
    initialize(app)
    print(f"Update parameters by command-line arguments")
    update_vars(app)

    main(app)