import sys
import os
import builtins
import numpy as np
from numpy import sqrt
import openpyxl
import pandas as pd
from sklearn.preprocessing import StandardScaler
from sklearn.metrics import mean_absolute_error, mean_squared_error, r2_score
from sklearn.model_selection import train_test_split

# sklearnの回帰モデルを片っ端から試す
# https://qiita.com/futakuchi0117/items/72ce4afae9adcccd6e18
from sklearn.linear_model import LinearRegression, Ridge, Lasso, ElasticNet, SGDRegressor
from sklearn.linear_model import PassiveAggressiveRegressor, ARDRegression, RidgeCV
from sklearn.linear_model import TheilSenRegressor, RANSACRegressor, HuberRegressor
from sklearn.neural_network import MLPRegressor
from sklearn.svm import SVR, LinearSVR
from sklearn.neighbors import KNeighborsRegressor
import sklearn.gaussian_process as gp
from sklearn.gaussian_process import GaussianProcessRegressor
from sklearn.tree import DecisionTreeRegressor
from sklearn.ensemble import RandomForestRegressor, AdaBoostRegressor, ExtraTreesRegressor, HistGradientBoostingRegressor
from sklearn.ensemble import BaggingRegressor, GradientBoostingRegressor, VotingRegressor, StackingRegressor
from sklearn.preprocessing import PolynomialFeatures
from sklearn.pipeline import Pipeline
from sklearn.cross_decomposition import PLSRegression

import matplotlib.pyplot as plt


from tklib.tkutils import replace_path
from tklib.tkapplication import tkApplication


#=======================================
# プログラム開始
#=======================================
infile = 'random-poly-ML.xlsx'
method = 'linear'
f_test = 0.7
random_state = None

alpha = 1.0
nmaxiter = 1000

plot_variation = 1


figsize = (6, 6)
fontsize   = 16
fontsize_s = 12


argv = sys.argv
narg = len(argv)
if narg >= 2:
    infile = argv[1]
if narg >= 3:
    method = argv[2]
if narg >= 4:
    f_test = float(argv[3])
if narg >= 5:
    alpha = float(argv[4])
if narg >= 6:
    nmaxiter = int(argv[5])

def main():
    app = tkApplication()
    app.log_path = app.replace_path(infile)
    app.redict(["stdout", app.log_path], "w")

    print("Regression by scikit-learn functions")
    print(f"infile={infile}")
    print(f"method={method}")
    print(f"Fraction of test data={f_test}")
    print(f"alpha={alpha}")
    print(f"nmaxiter={nmaxiter}")

    print("")
    print(f"Read [{infile}]")
    df = pd.read_excel(infile, engine = 'openpyxl')
    labels = df.columns.tolist()
    t_label  = labels[1]
    x_labels = labels[2:]
# 記述子
    x = df[labels[2:]]
# 目的関数
    y = df[labels[1]]
# プロット用X
    x_plot = df[labels[0]]

    ndata = len(df.index)
    ndescriptors = len(labels) - 2
    print("ndata=", ndata)
    print("ndescriptors=", ndescriptors)
    print("  x_labels=", x_labels[2:])

    print("Split to training and test data")
    x_train, x_test, y_train, y_test, x_plot_train, x_plot_test = \
            train_test_split(x, y, x_plot, test_size = f_test, random_state = random_state)
    print(" Number of training data:", len(x_train))
    print(" Number of test data:", len(x_test))

    print("")
    print(f"Execute regression")
    print(f"Standaridization")
    scaler = StandardScaler()
    scaler.fit(x_train)
    x_scaled_train = scaler.transform(x_train)
    x_scaled_test = scaler.transform(x_test)

    if method == 'linear':
        model = LinearRegression()
    elif method == 'ridge':
        model = Ridge(alpha = alpha)
    elif method == 'lasso':
        model = Lasso(alpha = alpha, max_iter = nmaxiter, tol = 1.0e-4)
    elif method == 'poly2':
        model = Pipeline([('poly', PolynomialFeatures(degree = 2)),('linear', LinearRegression())])
    elif method == 'poly3':
        model = Pipeline([('poly', PolynomialFeatures(degree = 3)),('linear', LinearRegression())])
    elif method == 'poly4':
        model = Pipeline([('poly', PolynomialFeatures(degree = 4)),('linear', LinearRegression())])
    elif method == 'poly5':
        model = Pipeline([('poly', PolynomialFeatures(degree = 5)),('linear', LinearRegression())])
    elif method == 'mlp':
# Define Neural Neowork model
#   Activation: ‘identity’, ‘logistic’, ‘tanh’, ‘relu’
#   Solver: ‘lbfgs’, ‘sgd’, ‘adam’
        model = MLPRegressor(hidden_layer_sizes = (3, 3, 3, 3), learning_rate = 'constant', learning_rate_init = 0.001,
                    momentum = 0.9, max_iter = nmaxiter, tol = 1.0e-7,
                    nesterovs_momentum = True, power_t = 0.5, random_state = None, shuffle=True,
                    validation_fraction = 0.1,
                    solver = 'lbfgs', activation = 'relu', 
                    alpha = 1.0e-4, batch_size = 'auto', beta_1 = 0.9, beta_2 = 0.999,
                    verbose = True, warm_start = False, early_stopping = True, n_iter_no_change = 5)
    elif method == 'rfr':
        model =RandomForestRegressor(bootstrap = True, criterion = 'squared_error', max_depth = None,
                    max_features = 'auto', max_leaf_nodes = None,
                    min_impurity_decrease = 0.0, 
                    min_samples_leaf = 1, min_samples_split = 2,
                    min_weight_fraction_leaf = 0.0, n_estimators = 10, n_jobs = -1,
                    oob_score = False, random_state = 0, verbose = True, warm_start = False)
    elif method == 'svr':
        model = SVR(kernel = 'rbf', C = 1.0e3, gamma = 0.1, epsilon = 0.1)
    elif method == 'gpr':
        model = GaussianProcessRegressor(kernel = gp.kernels.RBF(length_scale = 1.0), alpha = alpha)
    elif method == 'sgdr':
        model = SGDRegressor()

    print("")
    print(f"Fit")
    model.fit(x_scaled_train, y_train)

    print(f"Calculate")
    y_cal_train = model.predict(x_scaled_train)
    if method == 'gpr':
        y_cal_train, y_std_train = model.predict(x_scaled_train, return_std = True)
        y_cal_test, y_std_test   = model.predict(x_scaled_test, return_std = True)
    else:
        y_cal_test = model.predict(x_scaled_test)

    mae_train  = mean_absolute_error(y_train, y_cal_train)
    mse_train  = mean_squared_error(y_train, y_cal_train)
    rmse_train = sqrt(mse_train)
    r2_train   = r2_score(y_train, y_cal_train)
    mae_test   = mean_absolute_error(y_test, y_cal_test)
    mse_test   = mean_squared_error(y_test, y_cal_test)
    rmse_test  = sqrt(mse_test)
    r2_test    = r2_score(y_test, y_cal_test)

    print(f"Mean absolute error (MAE): training {mae_train:10.3g}  test: {mae_test:10.3g}")
    print(f"Mean squared error (MSE) : training {mse_train:10.3g}  test: {mse_test:10.3g}")
    print(f"Root MSE (RMSE)          : training {rmse_train:10.3g}  test: {rmse_test:10.3g}")
    print(f"R^2 score                : training {r2_train:10.3g}  test: {r2_test:10.3g}")

# modelオブジェクトにintercept, coef_があるか確認するため、
# オブジェクトの辞書型変数 __dict__ に直接アクセスし、get()で確認する
    if model.__dict__.get('intercept', None) is not None:
        print(f"    intercept: {model.intercept_}")
    if model.__dict__.get('coef_', None) is not None:
        for iv in range(len(x_labels)):
            print(f"    {x_labels[iv]:>10}: {model.coef_[iv]:12.4g}")


#================================================================
# Plot
#================================================================
# プロット用のx, yの値
    print("")
    print("plot")

    ymin = min(y)
    ymax = max(y)

    fig, axis = plt.subplots(1, 1, figsize = (8, 6))
    axis.scatter(y_train, y_cal_train, label = 'train')
    axis.scatter(y_test,  y_cal_test, label = 'test')
    axis.set_xlabel('input data', fontsize = fontsize)
    axis.set_ylabel('predict', fontsize = fontsize)
    axis.plot([ymin, ymax], [ymin, ymax], linestyle = 'dashed', color = 'red', linewidth = 0.5)
    axis.legend()
    plt.pause(0.1)

# 実線プロットのため、データをx順にソートした配列を作る
    x_train_l = x_plot_train.tolist()
    y_train_l = y_cal_train.tolist()
    x_test_l = x_plot_test.tolist()
    y_test_l = y_cal_test.tolist()
    n_train = len(x_train_l)
    n_test  = len(x_test_l)
    n = n_train + n_test
# GPR以外はstdデータがないので、ダミーを作る
    if method != 'gpr':
        y_std_train = [0.0] * n_train
        y_std_test  = [0.0] * n_test

    data = []
    for i in range(n_train):
        data.append([x_train_l[i], y_train_l[i], y_std_train[i]])
    for i in range(n_test):
        data.append([x_test_l[i], y_test_l[i], y_std_test[i]])
    data = sorted(data, key = lambda x: x[0])

    x_sorted = np.empty(n)
    y_sorted = np.empty(n)
    std = np.empty(n)
    for i in range(n):
        x_sorted[i]   = data[i][0]
        y_sorted[i]   = data[i][1]
        std[i] = data[i][2]

    fig, axis = plt.subplots(1, 1, figsize = (8, 6))
    axis.scatter(x_plot_train, y_train, label = 'train')
    axis.scatter(x_plot_test,  y_test, label = 'test')
#    axis.plot(x_plot_train, y_cal_train, label = 'fit', linestyle = '', marker = '.', markerfacecolor = 'green', markeredgecolor = 'green')
    if method == 'gpr':
        axis.fill_between(x_sorted, y_sorted - std, y_sorted + std, label = 'mean+-std', color = 'red', alpha = 0.3)
    else:
        axis.plot(x_sorted, y_sorted, linewidth = 0.5, color = 'black')

    axis.set_xlabel('$x$', fontsize = fontsize)
    axis.set_ylabel('$y$', fontsize = fontsize)
    axis.legend()
    plt.pause(0.1)

# 各記述子と目的関数のプロット
    if plot_variation:
        print("")
        print("Plot variasns of objective function vs each descriptor:")
        ncol = int(sqrt(ndescriptors) + 1)
        nrow = int(ndescriptors / ncol + 1.0001)
        nmesh = 101
        print("ncol, nrow=", ncol, nrow)
#        zero_list = np.zeros(nmesh, dtype = float)

        plt.rcParams['font.size'] = fontsize_s
        fig, ax = plt.subplots(nrow, ncol, figsize = figsize)
        for i in range(ndescriptors):
            ix = int(i / ncol)
            iy = i % ncol
            axis = ax[ix, iy]

            x_label = x_labels[i]
            x_list  = x[x_label]
#            print("i=", i, x_label, x_train[x_label])

            x0 = min([min(x_train[x_label]), min(x_train[x_label])])
            x1 = max([max(x_train[x_label]), max(x_train[x_label])])
            xstep = (x1 - x0) / (nmesh - 1)
            x_sim_list = [x0 + i * xstep for i in range(nmesh)]

# make descriptor
            avg = scaler.mean_[i]
            std = scaler.scale_[i]
            y_sim = []
            for idata in range(nmesh):
                x_std = np.zeros(ndescriptors, dtype = float)
# 他の記述子の値を０にする
#                x_std[i] = x_list[idata]
#                x_std = scaler.transform(pd.DataFrame([x_std], columns = x_labels))
# 他の記述子の値を平均値にする
                x_std[i] = (x_list[idata] - avg) / std

                y_sim.append(model.predict([x_std])[0])

            axis.scatter(x_train[x_label], y_train, s = 3.0)
            axis.scatter(x_test[x_label],  y_test, s = 3.0)
            axis.plot(x_sim_list, y_sim, linestyle = '-', color = 'black', linewidth = 0.5)
            axis.set_xlabel(x_label, fontsize = fontsize_s)
            axis.set_ylabel(t_label, fontsize = fontsize_s)

        plt.tight_layout()
        plt.pause(0.1)

    print("")
    print("Press ENTER to terminate")
    input()


if __name__ == "__main__":
    main()