#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Thu Jan 30 16:31:33 2020

@author: M. A. Bessa (M.A.Bessa@tudelft.nl)
         Assistant Professor at Delft University of Technology

Simple script adapting some examples provided by the scikit-learn user guide.

-> Kernel regression adapted from example created by:
          Mathieu Blondel
          Jake Vanderplas
-> Gaussian Process Regression adapted from example created by:
          Vincent Dubourg <vincent.dubourg@gmail.com>
          Jake Vanderplas <vanderplas@astro.washington.edu>
          Jan Hendrik Metzen <jhm@informatik.uni-bremen.de>

License: BSD 3 clause
"""

print(__doc__)

import sys
import numpy as np
import matplotlib.pyplot as plt
#
from sklearn.linear_model import Ridge
from sklearn.preprocessing import PolynomialFeatures
from sklearn.pipeline import make_pipeline
from sklearn.metrics import mean_squared_error, r2_score
#
#
def f(x):
    """ function to approximate by polynomial interpolation"""
    return x * np.sin(x)

seed = 1234 # set a random seed to replicate results
np.random.seed(seed)

# Consider a noiseless dataset (no uncertainty)
noisy_data = 0 # noisy_data = 0 -> noiseless dataset (no uncertainty)
               # noisy_data = 1 -> dataset has known uncertainty (Gaussian noise)
#
use_KERAS = 1 # = 1 -> Also use a Neural Network estimator (requires KERAS)
gridsearch = 0 # Do you want to look for the best parameters for the Neural Network? (slower)
#
#
n_train_points = 8 # number of training points (use at least 6)
n_test_points = 20 # number of test points

# generate points used to plot
x_plot = np.linspace(0, 10, 1000)

# generate points and keep a subset of them
x = np.linspace(0, 10, 1000)
rng = np.random.RandomState(seed)
rng.shuffle(x)

x_train = np.sort(x[:n_train_points])
x_test = np.sort(x[n_train_points+1:n_train_points+n_test_points+1])
# Calculate the y values for train and test sets:
y_train = f(x_train) # training data
dy_train = noisy_data*(0.5 + 1.0 * np.random.random(y_train.shape))
noise = np.random.normal(0, dy_train)
y_train += noise

y_test = f(x_test) # test data (unseen)
dy_test = noisy_data*(0.5 + 1.0 * np.random.random(y_test.shape))
noise = np.random.normal(0, dy_test)
y_test += noise

# create matrix versions of these arrays
X_train = x_train[:, np.newaxis]
X_test = x_test[:, np.newaxis]
X_plot = x_plot[:, np.newaxis]

colors = ['teal', 'yellowgreen', 'blue']
lw = 2
fig = plt.figure()
plt.plot(x_plot, f(x_plot), color='red', linewidth=lw, linestyle=':',
         label=u'ground truth: $f(x) = x\,\sin(x)$')

# Plot training points
if noisy_data==0: #noiseless data
    plt.plot(x_train, y_train, 'ro', markersize=6, label="training points")
else: # noisy data!
    plt.errorbar(x_train, y_train, dy_train, fmt='ro', markersize=6, label=u'training points inc. uncertainty')

# Plot test points
if noisy_data==0: #noiseless data
    plt.plot(x_test, y_test, 'kx', markersize=6, label="testing points")
else: # noisy data!
    plt.errorbar(x_test, y_test, dy_test, fmt='kx', markersize=6, label=u'testing points inc. uncertainty')

# Plot polynomial approximations (alpha \= 0 will become non-interpolatory)
for count, degree in enumerate([3, 4, n_train_points-1]):
    model = make_pipeline(PolynomialFeatures(min(degree,20)), Ridge(alpha=0.0))
    model.fit(X_train, y_train)
    y_plot = model.predict(X_plot)
    plt.plot(x_plot, y_plot, color=colors[count], linewidth=lw, linestyle='-',
             label="Polynomial of degree %d prediction" % min(degree,20))

plt.xlabel('$x$')
plt.ylabel('$f(x)$')
plt.ylim(-20, 20)
plt.legend(loc='lower left')

# Compute the Mean Squared Error for the polynomial of highest degree
y_test_pred = model.predict(X_test)
mse_value = mean_squared_error(y_test, y_test_pred)
print('MSE for polynomial of degree '+str(min(n_train_points-1,20))+' = ', mse_value)
print('R2 score for polynomial of degree '+str(min(n_train_points-1,20))+' = ', r2_score(y_test, y_test_pred))
#plt.title(r'MSE $\approx$ '+str(round(mse_value,2))+' when using '+str(n_train_points)+' training points')
plt.title(r'Polynomial fitting')

plt.show()

# ----------------------------------------------------------------------------
# Now with a Gaussian Process

from sklearn.gaussian_process import GaussianProcessRegressor
from sklearn.gaussian_process.kernels import RBF, ConstantKernel as C

# Instanciate a Gaussian Process model
kernel = C(1.0, (1e-3, 1e3)) * RBF(10, (1e-2, 1e2))
if noisy_data==0: #noiseless data
    gp = GaussianProcessRegressor(kernel=kernel, alpha=1e-3, n_restarts_optimizer=9) # using a small nugget in case the training points are very close to each other
else: # noisy data!
    gp = GaussianProcessRegressor(kernel, alpha=(dy_train) ** 2,
                                  n_restarts_optimizer=10)
#    gp = GaussianProcessRegressor(kernel=kernel, alpha=(dy_train/y_train) ** 2,
#                                  n_restarts_optimizer=10)

# Fit to data using Maximum Likelihood Estimation of the parameters
gp.fit(X_train, y_train)


# Make the prediction on the meshed x-axis (ask for MSE as well)
y_pred, sigma = gp.predict(X_plot, return_std=True)

# Plot the function, the prediction and the 95% confidence interval based on
# the MSE
fig = plt.figure()
plt.plot(x_plot, f(x_plot), 'r:', label=u'ground truth: $f(x) = x\,\sin(x)$')

# Plot training points
if noisy_data==0: #noiseless data
    plt.plot(x_train, y_train, 'ro', markersize=6, label="training points")
else: # noisy data!
    plt.errorbar(x_train, y_train, dy_train, fmt='ro', markersize=6, label=u'training points inc. uncertainty')

# Plot test points
if noisy_data==0: #noiseless data
    plt.plot(x_test, y_test, 'kx', markersize=6, label="testing points")
else: # noisy data!
    plt.errorbar(x_test, y_test, dy_test, fmt='kx', markersize=6, label=u'testing points inc. uncertainty')

plt.plot(x_plot, y_pred, 'b-', label="GPR prediction")
plt.fill(np.concatenate([x_plot, x_plot[::-1]]),
         np.concatenate([y_pred - 1.9600 * sigma,
                        (y_pred + 1.9600 * sigma)[::-1]]),
         alpha=.5, fc='b', ec='None', label='95% confidence interval')
plt.xlabel('$x$')
plt.ylabel('$f(x)$')
plt.ylim(-20, 20)
plt.legend(loc='upper left')

# Finally, assess the mean square error:
y_test_pred, test_sigma = gp.predict(X_test, return_std=True)
mse_value = mean_squared_error(y_test, y_test_pred)
print('MSE for GPR = ', mse_value)
print('R2 score for GPR = ', r2_score(y_test, y_test_pred))
#plt.title(r'MSE $\approx$ '+str(round(mse_value,2))+' when using '+str(n_train_points)+' training points')
plt.title(r'Gaussian Process Regression')

#
if use_KERAS != 1: # if not asked to estimate with Neural Network using Keras
    sys.exit('Exit because you do not want to fit a Neural Network') # exit code

# -----------------------------------------------------------
# Neural Networks example
# KERAS:
from keras.models import Sequential
from keras.layers.core import Dense
from keras.wrappers.scikit_learn import KerasRegressor
from keras.callbacks import EarlyStopping
from sklearn.preprocessing import StandardScaler
from sklearn.model_selection import GridSearchCV
#
# Function to create model, required for KerasClassifier when SPECIFYING INPUTS
def create_model(input_dimensions=1,neurons1=10,neurons2=10,neurons3=10,neurons4=10,activation='relu',optimizer='adam'):
    # create model
    model = Sequential()
    model.add(Dense(neurons1, input_dim=input_dimensions, activation=activation)) # first hidden layer
    model.add(Dense(neurons2, activation=activation)) # second hidden layer
#        model.add(Dense(neurons3, activation=activation)) # thrid hidden layer
#        model.add(Dense(neurons4, activation=activation)) # fourth hidden layer, etc.
    model.add(Dense(1))
    model.compile(loss='mse', optimizer=optimizer)
    return model
# -----------------------------------------------------------------------------
#

scaler = StandardScaler().fit(X_train)
#
X_train_scaled=scaler.transform(X_train)
X_test_scaled=scaler.transform(X_test)
X_plot_scaled=scaler.transform(X_plot)
#


if gridsearch==1:
    # create model
    early_stopping = EarlyStopping(monitor='val_loss', min_delta=0.0, patience=30, mode='min')
#    NN_model = KerasRegressor(build_fn=create_model(input_dimensions=1,neurons=20),
#                              callbacks=[early_stopping], validation_data=(scaler.transform(X_test), y_test))
    # define the grid search parameters
    neurons1 = [5,20,200] # number of neurons in hidden layer 1
    neurons2 = [5,10] # number of neurons in hidden layer 2 (if present; uncomment in create_model function)
    neurons3 = [10] # number of neurons in hidden layer 3 (if present; uncomment in create_model function)
    neurons4 = [10] # number of neurons in hidden layer 4 (if present; uncomment in create_model function)
    #
    batch_size = [n_train_points]
    #
    epochs = [1000]
    #
    optimizer = ['adam']
#    optimizer = ['SGD', 'RMSprop', 'Adagrad', 'Adadelta', 'Adam', 'Adamax', 'Nadam']
#    init_mode = ['uniform', 'lecun_uniform', 'normal', 'orthogonal', 'zero', 'one', 'glorot_normal', 'glorot_uniform', 'he_normal', 'he_uniform']    
    #
    param_grid = dict(batch_size=batch_size, epochs=epochs,neurons1=neurons1,
                      neurons2=neurons2,
                      #neurons3=neurons3,neurons4=neurons4, # commented out because I am not using them
                      optimizer=optimizer)
    NN_model = KerasRegressor(build_fn=create_model)
    grid = GridSearchCV(estimator=NN_model, param_grid=param_grid, n_jobs=1, cv=3, iid=False)
    grid_result = grid.fit(X_train_scaled, y_train, callbacks=[early_stopping], validation_data=(scaler.transform(X_test), y_test))
    history = grid_result.best_estimator_.fit(X_train_scaled, y_train,callbacks=[early_stopping], validation_data=(scaler.transform(X_test), y_test))
    # summarize results
    print("Best: %f using %s" % (grid_result.best_score_, grid_result.best_params_))
    means = grid_result.cv_results_['mean_test_score']
    stds = grid_result.cv_results_['std_test_score']
    params = grid_result.cv_results_['params']
    for mean, stdev, param in zip(means, stds, params):
        print("%f (%f) with: %r" % (mean, stdev, param))
else: # just use a particular Neural Network of choice
    # Define early stopping:
    early_stopping = EarlyStopping(monitor='val_loss', min_delta=0.0, patience=30, mode='min')
    neurons1=200
    neurons2=10
    NN_model = KerasRegressor(build_fn=create_model, neurons1=neurons1, neurons2=neurons2,
                              batch_size=n_train_points, epochs=2000, optimizer='adam',
                              callbacks=[early_stopping],
                              validation_data=(scaler.transform(X_test), y_test))
    #
    history = NN_model.fit(X_train_scaled, y_train)

#
plt.figure()
plt.plot(history.history['loss'])
plt.plot(history.history['val_loss'])
plt.title('Training and testing loss')
plt.ylabel('loss')
plt.xlabel('epoch')
plt.legend(['training', 'testing'], loc='upper right')
plt.show()

# Plot the function
fig = plt.figure()
plt.plot(x_plot, f(x_plot), 'r:', label=u'ground truth: $f(x) = x\,\sin(x)$')

# Plot training points
if noisy_data==0: #noiseless data
    plt.plot(x_train, y_train, 'ro', markersize=6, label="training points")
else: # noisy data!
    plt.errorbar(x_train, y_train, dy_train, fmt='ro', markersize=6, label=u'training points inc. uncertainty')

# Plot test points
if noisy_data==0: #noiseless data
    plt.plot(x_test, y_test, 'kx', markersize=6, label="testing points")
else: # noisy data!
    plt.errorbar(x_test, y_test, dy_test, fmt='kx', markersize=6, label=u'testing points inc. uncertainty')

y_pred = history.model.predict(X_plot_scaled)
plt.plot(x_plot, y_pred, 'b-', label="Neural Network prediction")
plt.title(r'NN with '+str(neurons1)+' neurons in the 1st hidden layer, and '+str(neurons2)+' in the 2nd')

plt.xlabel('$x$')
plt.ylabel('$f(x)$')
plt.ylim(-20, 20)
plt.legend(loc='upper left')