from sklearn.gaussian_process import GaussianProcessRegressor
from sklearn.gaussian_process.kernels import RBF, WhiteKernel
import numpy as np
import pandas as pd
from joblib import dump


np.random.seed(152)
observation = pd.read_csv("logs/cbf/states_15bins.csv")
observation = np.asarray(observation)
# be very careful the state and action is from current time step
state_x = observation[0:5000:2, 0:15]  # 1
action_x = observation[1:5001:2, 28:30]  # 0
state_action_x = np.concatenate((state_x, action_x), axis=1)
state_y = observation[1:5001:2, 0:15]  # 2
index = state_action_x[0:2500, 15:17] != 0
index = index[0:2500, 0] * index[0:2500, 1]
index = np.asarray(index)
state_action_x = state_action_x[index, :]
state_y = state_y[index, :]
kern = 1 ** 2 * RBF(length_scale=10) + WhiteKernel(noise_level=.001)
gp_dynamics = GaussianProcessRegressor(kernel=kern, alpha=.1, n_restarts_optimizer=10)
gp_dynamics.fit(state_action_x, state_y)
print('R² of GP dynamics: ' + str(gp_dynamics.score(state_action_x, state_y)))
dump(gp_dynamics, 'logs/cbf/gp_15bins_5k_092.joblib')


np.random.seed(152)
observation = pd.read_csv("logs/cbf/states_15bins.csv")
observation = np.asarray(observation)
# be very careful the state and action is from current time step
state_x = observation[0:2000:2, 0:15]  # 1
action_x = observation[1:2001:2, 28:30]  # 0
state_action_x = np.concatenate((state_x, action_x), axis=1)
state_y = observation[1:2001:2, 0:15]  # 2
index = state_action_x[0:1000, 15:17] != 0
index = index[0:1000, 0] * index[0:1000, 1]
index = np.asarray(index)
state_action_x = state_action_x[index, :]
state_y = state_y[index, :]
kern = 1 ** 2 * RBF(length_scale=10) + WhiteKernel(noise_level=.001)
#kern = RBF(length_scale=10) + WhiteKernel(noise_level=.001)
gp_dynamics = GaussianProcessRegressor(kernel=kern, alpha=.1, n_restarts_optimizer=10)
gp_dynamics.fit(state_action_x, state_y)
print('R² of GP dynamics: ' + str(gp_dynamics.score(state_action_x, state_y)))
# dump(gp_dynamics, 'logs/cbf/gp_15bins_2k_090.joblib')


np.random.seed(152)
observation = pd.read_csv("logs/cbf/states_15bins.csv")
observation = np.asarray(observation)
# be very careful the state and action is from current time step
state_x = observation[0:1000:2, 0:15]  # 1
action_x = observation[1:1001:2, 28:30]  # 0
state_action_x = np.concatenate((state_x, action_x), axis=1)
state_y = observation[1:1001:2, 0:15]  # 2
index = state_action_x[0:500, 15:17] != 0
index = index[0:500, 0] * index[0:500, 1]
index = np.asarray(index)
state_action_x = state_action_x[index, :]
state_y = state_y[index, :]
kern = 1 ** 2 * RBF(length_scale=10) + WhiteKernel(noise_level=.001)
gp_dynamics = GaussianProcessRegressor(kernel=kern, alpha=.1, n_restarts_optimizer=10)
gp_dynamics.fit(state_action_x, state_y)
print('R² of GP dynamics: ' + str(gp_dynamics.score(state_action_x, state_y)))
dump(gp_dynamics, 'logs/cbf/gp_15bins_1k_087.joblib')









from datetime import datetime
# evluate
np.random.seed(152)
observation = pd.read_csv("logs/cbf/states_15bins_eval.csv")
observation = np.asarray(observation)
# be very careful the state and action is from current time step
state_x = observation[0:71616:2, 0:15]  # 1
action_x = observation[1:71617:2, 28:30]  # 0
state_action_x = np.concatenate((state_x, action_x), axis=1)
state_y = observation[1:71617:2, 0:15]  # 2
index = state_action_x[0:71616, 15:17] != 0
index = index[0:35808, 0] * index[0:35808, 1]
index = np.asarray(index)
state_action_x = state_action_x[index, :]
state_y = state_y[index, :]
start = datetime.now()
gp_dynamics.predict(state_action_x[0:250], return_std=True)
end = datetime.now()
print(str(end-start))
start = datetime.now()
for i in range(0, len(state_action_x[0:45])):
    gp_dynamics.predict(state_action_x[i].reshape(1, -1), return_std=True)
end = datetime.now()
print(str(end-start))
print('R² of GP dynamics: ' + str(gp_dynamics.score(state_action_x, state_y)))