import os
import platform
import sys
import re

sys.path.append('.')

try:
    import tklib.tkimport as imp
except Exception as e:
    print()
    print("######################################################################")
    print("###########  ERROR ERROR ERROR ERROR ERROR ERROR #####################")
    print("######################################################################")
    print(f"# Failed to import [tklib.tkimport] module ({e}).")
    print(f"#  Add [tkProg]{os.sep}tklib{os.sep}python to PYTHONPATH variable")
    print(f"#  Current PYTHONPATH:", sys.path)
    print("######################################################################")
    input("Press ENTER to terminate>>")
    exit()

np  = imp.import_lib("numpy", stop_by_error = False)
pd  = imp.import_lib("pandas", stop_by_error = False)
mpl = imp.import_lib("matplotlib", stop_by_error = False)
imp.messages(stop_by_error = True)


from numpy import sqrt, exp, log, log10
from matplotlib import pyplot as plt

from sklearn.linear_model import LinearRegression, Lasso
from sklearn.preprocessing import StandardScaler
from sklearn.metrics import mean_absolute_error


from tklib.tkutils import print_data, pint, pfloat, format_strlist
from tklib.tkvariousdata import tkVariousData
from tklib.tkapplication import tkApplication
from tklib.tkparams import tkParams

from tklib.tksci.tksci import e, kB
from tklib.tksci.tkFit_lib import conv_input, save_data, add_history
import tklib.tksci.tkoptimize_flex as opt
import optimize_mup_mf as mf
import mobility_pi_model as omodel
import mobility_pi as cmu


"""
Fitting by Hall - T data to polynomial mobility model
"""


ProgramName = "optimize_mu_pi"


def init(app, cfg, mf):
    print("")
    print("Initialize fitting condition:")
    fit = opt.clean(app, cfg, mf)

    config_path   = cfg.parameterfile
    fitparam_path = fit.fitfile_path

    fit = mf.init_fit(app, cfg)

    pk_all = opt.linear_fit(app, cfg, mf, omodel, fit, print_level = True)
    fit.pk = pk_all

    print()
    mf.save_parameter_files(config_path, fitparam_path, fit, cfg)

    app.terminate(pause = True)

def plot_fit(app, cfg, mf):
    if cfg.file1 == 'fit':
        if cfg.file2 == '':
            app.terminate("\nError: Fit file must be given by optional arg #2.\n", pause = True)
        else:
            infile = cfg.file2
    elif cfg.file1 == '':
        app.terminate("\nError: Fit file must be given by optional arg #2.\n", pause = True)
    else:
        infile = cfg.file1

    print()
    print(f"Read fit data from [{infile}]")
    #文字列infileの終わりにある数値をreを使って、取り出す
    m = re.search(r'\d+$', os.path.splitext(infile)[0])
    if m:
        num = int(m.group())
    else:
        num = None
    print(f"  hisotry #: {num}")

    df = pd.read_excel(infile)#, sheet_name = None)
    labels = df.columns.tolist()
    VG_list = df.iloc[:, 0].values.tolist()
    VD_list = df.iloc[:, 1].values.tolist()
    data_list = df.iloc[:, 2:].to_numpy().T.tolist()
    ndata = len(data_list)
    print("labels=", labels)

    if num is not None:
        cfg.historyfile = app.replace_path(cfg.infile, template = ["{dirname}", "{filebody}-history.xlsx"])
        print()
        print(f"Read parameters for #{num} from [{cfg.historyfile}]")
        df_h = pd.read_excel(cfg.historyfile)#, sheet_name = None)
        #df_hの '-iter'列から、値がnumの行を取り出す
        df_num = df_h[df_h['-iter'] == num]
 #       print("df_num=", df_num)
        labels_h = df_h.columns.tolist()
        for i, label in enumerate(labels_h):
            val = df_num.iloc[0, i]
            print(f"  {label:15}: {val}")

# Plot graphs
    fig, axes = plt.subplots(1, 1, figsize = (8, 6))
    ax = axes

    ax.tick_params(labelsize = cfg.fontsize)
    if ndata >= 1:
        ax.plot(VG_list, data_list[0], label = labels[2], linestyle = '', marker = 'o', markersize = 3.0)
    if ndata >= 2:
        ax.plot(VG_list, data_list[1], label = labels[3], linestyle = '-', color = 'blue', linewidth = 0.5)
    if ndata >= 3:
        ax.plot(VG_list, data_list[2], label = labels[4], linestyle = '-', color = 'red', linewidth = 0.5)
    ax.set_xlabel(r"$V_g$ (V)", fontsize = cfg.fontsize)
    ax.set_ylabel(r"$I_d$ (A)", fontsize = cfg.fontsize)
    ax.set_yscale('log')
    ax.legend(fontsize = cfg.legend_fontsize)

    plt.tight_layout()

    plt.pause(0.0001)
    app.terminate(pause = True)

def plot_mu_T(app, cfg, mf):
    print()
    print("Plot mobility vs T")
    fit = mf.init_fit(app, cfg)

    Ts_T = np.arange(250.0, 360.0, 20.0) #[50.0, 100.0, 150.0, 200.0, 250.0, 300.0]
    Ns_T = [1.0e16, 1.0e17, 1.0e18, 1.0e19, 1.0e20]

    cmu.update_variables(*fit.pk)

    cmu.Ts_T = Ts_T
    cmu.Ns_T = Ns_T
    cmu.plot_T()

def plot_weight(app, cfg, mf):
    print()
    print("Plot mu-T/weight-T data")
    fit = mf.init_fit(app, cfg)

    xlabels, ylabels, xdata_list, ydata_list, _ = mf.read_input_data(app, cfg.infile, cfg = cfg)
    T_list = xdata_list[0]
    mu_list = ydata_list[0]
#    cmu.plot(T_list, mu_list, Tcal_list = None, mucal_list = None)

    Eb, s_phi, Eop, aop, pi, ai = cmu.convert_xk(fit.pk)
    mucal_list, mu_KGB_list, mu_ingrain_list, mu_op_list, mu_pi_list = cmu.cal_mu_list(T_list, Eb, s_phi, Eop, aop, pi, ai, rettype = 'all')
    print("")
    print(f"{'T(K)':6} {'mu(obs)':10} {'KGB':10} {'mu,ingrain':10} {'mu,op':10} ", end = '')
    for i, mu in enumerate(mu_pi_list): 
        s = f'mu({pi[i]})'
        print(f"{s:10} ", end = '')
    print()

    for iT, T in enumerate(T_list):
        muobs = mu_list[iT]
        KGB   = mu_KGB_list[iT]
        muing = mu_ingrain_list[iT]
        muop  = mu_op_list[iT]

        print(f"{T:6.3g} {muobs:10.4g} {KGB:10.4g} {muing:10.4g} {muop:10.4g} ", end = '')
        for i, mu in enumerate(mu_pi_list): print(f"{mu[iT]:10.4g} ", end = '')
        print()

    ywmu_gb_list, ywmu_op_list, ywmu_pi_list = cmu.cal_weight_list(T_list, pi, mu_KGB_list, mu_op_list, mu_pi_list)

    out_path = f"mu_components.xlsx"
#    print(f"  Save to [{out_path}]")
    labels = ['T(K)', 'mu(obs)', 'KGB', 'mu,ingrain', 'mu,op'] + [f'mu({p})' for p in pi]
    data_list = [T_list, mu_list, mu_KGB_list, mu_ingrain_list, mu_op_list] + mu_pi_list
    tkVariousData().to_excel(out_path, labels = labels, data_list = data_list)


#====
# plot
    print()
    print("plot")
    figsize = (10, 8)
    fontsize = 16
    legend_fontsize = 12
    colors = ['black', 'red', 'blue', 'green', 'orange', 'darkgreen', 'darkorange', 'navy', 
          'slategray', 'hotpink', 'olive', 'chocolate', 'magenta', 
          'green', 'yellow', 'cyan']

    fig, axes = plt.subplots(1, 2, figsize = figsize)
    axes[0].tick_params(labelsize = fontsize)
    axes[1].tick_params(labelsize = fontsize)

#    plot(T_list, mu_list, T_list, mucal_list)
    cmu.plot_muT_decomposed(axes[0], pi, T_list, mu_list, ymucal = mucal_list, 
                                ymu_ingrain = mu_ingrain_list, ymuop = mu_op_list, ymupi = mu_pi_list)
    cmu.plot_muT_weight(axes[1], pi, T_list, wmugb = ywmu_gb_list, wmu_ingrain = None, wmuop = ywmu_op_list, wmupi = ywmu_pi_list)

    plt.tight_layout()

    plt.pause(0.0001)
    app.terminate(pause = True)

def plot(app, cfg, mf):
    arg1 = cfg.file1.lower()
    print(f"target property: {arg1}")

    if arg1 == '':
        omodel.plot_input(app, cfg, mf)
#        plot_input(app, cfg, mf)
    elif arg1 == 'fit':
        plot_fit(app, cfg, mf)
    elif arg1 == 'weight':
        plot_weight(app, cfg, mf)
    elif arg1 == 'history':
        plot_history(app, cfg)
    else:
        plot(app, cfg)


def main():
    app = tkApplication()

    if len(sys.argv) > 1 and sys.argv[1] == 'mk_config':
        app.make_config_files(arg_config_file = 'arg_config.xlsx', fit_config_file = 'fit_config.xlsx')
        app.terminate("\n", pause = True)


#初期値の確定順序
#intialize(): 起動時引数と初期値の設定
    cfg = mf.initialize(app)

#--infile引数をチェックし、読み込みファイル名を最初に確定する
    cfg.infile = app.check_arg('--infile', defval = None, vartype = 'str')

#cfg.infileからログファイル名を作り、console出力をredirectする
    cfg.logfile = app.replace_path(cfg.infile)
    app.redirect(heading = f"Open logfile [{cfg.logfile}]",
                targets = ["stdout", cfg.logfile], mode = 'w')
    app.set_exception()

    try:
        app.print_title(f"#  {ProgramName}")
    except:
        pass

#cfg.infileが入力ファイルか設定ファイルであるかによって、infileとconfig_pathを返す
#  設定ファイルである場合、infile_storedをNoneにして、後で設定ファイルからcfg.infileを読み込む
    cfg.infile, cfg.parameterfile, infile_stored \
            = app.analyze_infile(cfg.infile, infile_ext = [".xlsx"], config_ext = [".in", ".ini", ".prm"], print_level = 1)

#起動時引数で与えられたパラメータをcfgに設定
    print()
    print("Update parameters from command line arguments")
    app.update_vars(cfg, apply_default = True)
#    cfg.print_parameters()

    mf.initialize_minimize_func(app)

#設定ファイルを読み込み、cfgに設定。#cfg.infileも書き換えられる
    print()
    print(f"Read configration file [{cfg.parameterfile}]")
    mf.read_parameters(cfg.parameterfile, cfg)
#    cfg.print_parameters()

#infile_stored = Noneの場合、設定ファイルのcfg.infileをinfile_storedに再保存
#infile_storedが最終的な読み込みファイル名になる
    if infile_stored is None: infile_stored = cfg.infile

#起動時引数で与えられたパラメータを最優先にするため、再設定
    app.force_given_args(cfg)
#    cfg.print_parameters()
#    exit()

# --daemonフラグが立っていたらdaemonize()
    if cfg.get("daemon", 0): 
        print("\ndaemon flag is set:")
        print("  fplot flag is set to -1")
        cfg.fplot = -1
        app.daemonize(stdout_path = 'daemon.log', stderr_path = 'daemon.log', print_level = 1)
# --nohupフラグが立っていたらHUPシグナルを無視
    if cfg.get("nohup", 0): app.nohup()

#cfg.infileを infile_storedに最終設定
    cfg.infile = infile_stored
    cfg.print_parameters()
#    exit()

    cfg.historyfile   = app.replace_path(cfg.infile, template = ["{dirname}", "{filebody}-history.xlsx"])
    cfg.stopfile      = app.replace_path(cfg.infile, template = ["{dirname}", "stop"])
    print(f"infile: {cfg.infile}")
    print(f"parameterfile: {cfg.parameterfile}")
    print(f"stop file: {cfg.stopfile}")
    if os.path.isfile(cfg.stopfile):
        print(f"\nWarning: stop file [{cfg.stopfile}] exisits")
        print("  remove it")
        os.remove(cfg.stopfile)

    if cfg.file1 in ['clean', 'inf', 'kill', 'plot', 'init', 'search', 'lfit', 'lasso', 'fit', 'update', 'sim', "error"]:
        cfg.mode = cfg.file1
        if cfg.file2 == '':
            cfg.file1 = ''
        else:
            cfg.file1 = cfg.file2
            cfg.file2 = ''

    if cfg.mode == 'clean':
        opt.clean(app, cfg, mf)
    if cfg.mode == 'kill':
        opt.kill(app, 'python')
    elif cfg.mode == 'init':
        init(app, cfg, mf)
    elif cfg.mode == 'inf':
        opt.inf(app, cfg, mf)
    elif cfg.mode == 'lfit':
        opt.linear_fit_func(app, cfg, mf, omodel)
    elif cfg.mode == 'lasso':
        opt.lasso(app, cfg, mf, omodel)
    elif cfg.mode == 'sim':
        opt.sim(app, cfg, mf)
    elif cfg.mode == 'scan':
        opt.scan(app, cfg, mf)
    elif cfg.mode == 'mapping':
        opt.mapping(app, cfg, mf)
    elif cfg.mode == 'sampling':
        opt.sampling(app, cfg, mf)
    elif cfg.mode == 'fit1':
        opt.fit1(app, cfg, mf)
    elif cfg.mode == 'fitlp':
        opt.fit(app, cfg, mf, optid_mode = 'fitlp')
    elif cfg.mode == 'fit':
        if cfg.get("daemon", 0):
            opt.fit_simple(app, cfg, mf)
        else:
            opt.fit(app, cfg, mf)
    elif cfg.mode == 'error':
        opt.error(app, cfg, mf)
    elif cfg.mode == 'mlr':
        opt.mlr(app, cfg, mf)
    elif cfg.mode == 'plot':
        plot(app, cfg, mf)
#        omodel.plot_input(app, cfg, mf)
    elif cfg.mode == 'plot_weight':
        plot_weight(app, cfg, mf)
    elif cfg.mode == 'plot_history':
        plot_history(app, cfg, mf)
    elif cfg.mode == 'plot_fit':
        plot_fit(app, cfg, mf)
    elif cfg.mode == 'plot_mu_T':
        plot_mu_T(app, cfg, mf)
#    elif cfg.mode == 'make_input_data':
#        atlas.make_input_data(app, cfg, mf)
#    elif cfg.mode == 'update':
#        atlas.update_files_func(app, cfg, mf)
#    elif cfg.mode == 'convert':
#        opt.convert(app, cfg, mf, cfg.get('input_template', None), cfg.get('converted_file', None))
    elif cfg.mode in globals() and callable(globals()[cfg.mode]):
# cfg.modeで与えられた関数が定義されていれば、呼び出す
        globals()[cfg.mode](app, cfg)
    else:
        app.terminate(message = "\n", 
                        usage = app.usage,
                        post_message = f"Error in main: Invalide mode [{cfg.mode}]", 
                        pause = True,
                     )


if __name__ == "__main__":
    main()