import numpy as np
from sklearn.metrics import mean_squared_error, mean_absolute_error
import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
from sklearn.ensemble import RandomForestRegressor
from sklearn.model_selection import GridSearchCV
from sklearn.linear_model import Lasso
from sklearn.linear_model import LassoCV
from sklearn.metrics import r2_score
import matplotlib.pyplot as plt

# read excel file
features = pd.read_excel('rfFeatures2.xlsx')
print('The shape of features is :', features.shape)

# prepare train dataset
X_true = features.iloc[:, 2:].values
y_true = features.iloc[:, 1].values

# Standard
sc = StandardScaler()
X_sc = sc.fit_transform(X_true)

# train set and test set
X_train, X_test, y_train, y_test = train_test_split(X_sc, y_true, test_size=0.15, random_state=0)

# Grid Search for estimators
new_params = {
    "n_estimators": range(100, 1500, 100),
    "max_depth": range(1, 20, 1),
    "min_samples_leaf": range(1, 20, 1),
    "min_samples_split": range(2, 20, 2)
}

forest_reg = RandomForestRegressor(bootstrap=True, oob_score=True)
gs = GridSearchCV(estimator=forest_reg, param_grid=new_params,  scoring='neg_mean_absolute_error')
grid_search = gs.fit(X_train, y_train)
best_params = gs.best_params_
best_estimator = gs.best_estimator_
best_score = gs.best_score_
print('best parameters:', best_params)
print('best estimator:', best_estimator)
print('best score:', best_score)

# RF_0 model training
rf_0 = RandomForestRegressor(bootstrap=True,
                             oob_score=True,
                             n_estimators=1100,
                             max_features='auto',
                             max_depth=18,
                             n_jobs=-1,
                             min_samples_leaf=18,
                             min_samples_split=16)

rf_0.fit(X_train, y_train)
rfy_pred = rf_0.predict(X_sc)

# R2
R2 = r2_score(y_true, rfy_pred)
print("R2:", R2)
# RMSE
RMSE = np.sqrt(mean_squared_error(y_true, rfy_pred))
print('RMSE:', RMSE)
# MAE
MAE = mean_absolute_error(y_true, rfy_pred)
print('MAE:', MAE)

# importance
feat_labels = features.columns[2:]
print(feat_labels)
rf_importance = rf_0.feature_importances_
print(rf_importance)

# RF error
error_rf = rfy_pred - y_true

# Output RF result
a = np.array(error_rf)
np.savetxt('error_rf.txt', a, newline='\n', delimiter=' ')

b = np.array(rfy_pred)
np.savetxt('rf_pre_result.txt', b, newline='\n', delimiter=' ')

c = np.array(rf_importance)
np.savetxt('rf_feature_importance.txt', c, newline='\n', delimiter=' ')

# ML Model
# LassoCV search
lassocv = LassoCV()
lassocv.fit(X_train, y_train)
print(lassocv.alpha_)
alg = Lasso(alpha=lassocv.alpha_)

# fitting
alg.fit(X_train, y_train)

# prediction
ml_y_pred = alg.predict(X_sc)
error_ml = ml_y_pred - y_true

# coefficient
print('Coefficient:', alg.coef_)

# intercept
print('intercept:', alg.intercept_)

# RMSE
RMSE = np.sqrt(mean_squared_error(y_true, ml_y_pred))
print('RMSE:', RMSE)

# MAE
MAE = mean_absolute_error(y_true, y_pred)
print('MAE:', MAE)

# R2
R2 = r2_score(y_true, y_pred)
print("R2:", R2)

# output ML results
d = np.array(ml_y_pred)
np.savetxt('lasso_pre_result.txt', d, newline='\n', delimiter=' ')

# draw plot
# read data
plot_xy = pd.read_excel('TrueandRFresult1214.xlsx')
x_true = plot_xy.iloc[:, 1].values
y_true = plot_xy.iloc[:, 2].values
z_true = plot_xy.iloc[:, 3].values
# plot
x = np.array(x_true)
y = np.array(y_true)
z= np.array(z_true)

xmin = x.min()
xmax = 0.5*x.max()
ymin = 0
ymax = y.max()
zmin = 0
zmax = z.max()

fig, axs = plt.subplots(ncols=1,nrows=2,sharey=True,dpi=600,figsize=(8, 16))
fig.subplots_adjust(hspace=0.5, left=0.07, right=0.93)

ax = axs[0]
hb = ax.hexbin(x, y, gridsize=50, bins='log', cmap='Blues')
ax.set(xlim=(xmin, xmax), ylim=(ymin, ymax))
ax.yaxis.get_ticklocs(minor=True)
ax.minorticks_on()
ax.set_xlabel('True_value')
ax.set_ylabel('Prediction_value')
ax.set_title("RF Result")
ax.plot([0, xmax], [0, xmax])
cb = fig.colorbar(hb, ax=ax)
cb.set_label('Count')

# ML model
xmin = 0
xmax =  0.5*x.max()
zmin = 0
zymax = z.max()
ax = axs[1]
hb = ax.hexbin(x, z, gridsize=50, bins='log', cmap='Blues')
ax.set(xlim=(xmin, xmax), ylim=(zmin, zmax))
ax.yaxis.get_ticklocs(minor=True)
ax.minorticks_on()
ax.set_xlabel('True_value')
ax.set_ylabel('Prediction_value')
ax.set_title("ML Result")
ax.plot([0, xmax], [0, xmax])
cb = fig.colorbar(hb, ax=ax)
cb.set_label('Count')
plt.show()