#####################################################################################
# A. Dutta 2021
# 
# This script trains the SVM regressor with the 'rbf' kernel to predict moment value 
# on the magnetic compounds. The dataset 'data-long1.csv' containing 836 
# stable 2D magnetic materials is loaded, out of which 80% is used for 'GridSearchCV'
# based 5-fold cross-validation and the rest 20% is used as the test set.
# ===================================================================================
# The training and evaluation of the model perfomance is done for 100
# different splits of the original dataset and the results are stored in the
# file './stats/mag-rbfSVR-sp0.20-stat.txt'.
#####################################################################################
#import required packages
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from matplotlib import rc
#rc('text', usetex=True) # If latex causes problems, comment this line 
from sklearn.metrics import mean_absolute_error, mean_squared_error, r2_score, explained_variance_score

MAE_data = pd.read_csv('../../datasets/magnetic/data-long2.csv',sep=',')

#print(MAE_data['smeig12'].values)
print('Shape= ', MAE_data.shape)
# Add derived features

# mom is mag_mom per metal atom
MAE_data['mom'] = MAE_data['mag_mom']/MAE_data['n_metal']
MAE_data['mom'] = abs(MAE_data['mom'].values)

print('Shape= ', MAE_data.shape)

# Select desired features
#MAE_data = MAE_data[['mom', 'mean_Z','del_Z','mode_Z','mean_group','del_group','mode_group',\
#'mean_period','del_period','mode_period','mean_val','del_val','mode_val','mean_electroneg','del_electroneg',\
#'mode_electroneg', 'entropy', 'cell_area', 'L2_norm', 'L3_norm', \
#'acsf1','acsf2','acsf3','acsf4','acsf5','acsf6','acsf7','acsf8','acsf9','acsf10','acsf11','acsf12','acsf13','acsf14','acsf15','acsf16',\
#'acsf17','acsf18','acsf19','acsf20','acsf21','acsf22','acsf23','acsf24','acsf25','acsf26','acsf27','acsf28','acsf29','acsf30','acsf31','acsf32',\
#'smeig1','smeig2','smeig3','smeig4','smeig5','smeig6','smeig7','smeig8','smeig9','smeig10','smeig11','smeig12'\
#]]

MAE_data.drop(labels=['id','uu','dd','ud','du','mag_mom','cell_area','E_zx', 'E_zy',\
'H','He','Li','Be','B','C','N','O','F','Ne',\
'Na','Mg','Al','Si','P','S','Cl','Ar',\
'K','Ca','Sc','Ti','V','Cr','Mn','Fe','Co','Ni','Cu','Zn','Ga','Ge','As','Se','Br','Kr',\
'Rb','Sr','Y','Zr','Nb','Mo','Tc','Ru','Rh','Pd','Ag','Cd','In','Sn','Sb','Te','I','Xe',\
'Cs','Ba','La','Ce','Pr','Nd','Pm','Sm','Eu','Gd','Tb','Dy','Ho','Er','Tm','Yb','Lu',\
'Hf','Ta','W','Re','Os','Ir','Pt','Au','Hg','Tl','Pb','Bi','Po','At'], axis=1, inplace=True)
#df.drop(labels=[ 'E_zx', 'E_zy'], axis=1, inplace=True)
print(MAE_data.shape)

#sb.distplot(MAE_data['log_MAE'])
#plt.show()

# Select rows after setting cutoff for the target variable 
y_min = 0.0
#y_max = 4
MAE_data = MAE_data.loc[MAE_data['mom'] >= y_min]

Nsamples = MAE_data.shape[0]
#print('log(MAE) selected between [', y_min, ':', y_max, ']')
print('No. of samples = ', Nsamples)

# open a file to append parameter values and scores
file1 = open("./stats/mag-rbfSVR-sp0.2-stat.txt", "a")
file1.write("# C,\tepsilon,\tgamma,\tr2_train,\tr2_test,\tpearson,\texp_var_train,\texp_var_test,\tMAE_train,\tMAE_test,\tMAE_train_inv,\tMAE_test_inv #\n\n")
Nruns = 100

# Import function to create training and test set splits
from sklearn.model_selection import train_test_split
# Scale the features and the target
from sklearn.preprocessing import QuantileTransformer, RobustScaler, MinMaxScaler
# SVM regerssor and GridSearchCV
from sklearn.svm import SVR
from sklearn.model_selection import GridSearchCV

for cnt in range(Nruns):
	print("Run count : ", cnt+1)
	#Shuffle rows randomly
	MAE_data = MAE_data.sample(frac=1, axis=0).reset_index(drop=True)

	y = MAE_data[['mom']].values
	X = MAE_data.drop(labels=['mom'], axis=1, inplace=False).to_numpy()
	
	# Test/train split
	X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=1)
	print('Shape ',X_train.shape)

	y_test_data = y_test[:,0]
	Ntrain = np.shape(y_train)[0] # no. of training examples

	# Scaling feature data
	Xscaler = QuantileTransformer(n_quantiles=Ntrain, output_distribution='normal').fit(X_train) # Handles outliers better
#Xscaler = MinMaxScaler(feature_range=(-1,1)).fit(X_train)
	X_train = Xscaler.transform(X_train)
	X_test = Xscaler.transform(X_test)

	# Scaling target data
	y_train = y_train.reshape(-1,1)
	y_test = y_test.reshape(-1,1)


	Yscaler = QuantileTransformer(n_quantiles=Ntrain, output_distribution='normal').fit(y_train)
	#Yscaler = MinMaxScaler(feature_range=(-1, 1)).fit(y_train)
	y_train = Yscaler.transform(y_train)
	y_test = Yscaler.transform(y_test)
	#######################

	#Create SVM regressor
	N_CV = 10
	C_vals = np.logspace(-2,2,10, base=10, endpoint=True)
	eps_vals = np.logspace(-3,1,10, base=10, endpoint=True)
	gamma_vals = np.logspace(-3,1,12, base=10, endpoint=True)

	parameters = {'kernel':['rbf'], 'C':C_vals, 'gamma':gamma_vals, 'epsilon':eps_vals}
	# Perform parallelised grid search over parmeter space (uses all available cores)
	reg = GridSearchCV(SVR(), parameters, scoring='neg_mean_absolute_error', cv=N_CV, n_jobs=-1, verbose=0)
	opt = reg.fit(X_train, y_train.ravel())

	print(opt.best_params_.items())
	# append to file
	row = []
	for k,v in opt.best_params_.items():
		if k != 'kernel':
			row.append(v)

	y_pred = opt.best_estimator_.predict(X_test)
	y_pred_train = opt.best_estimator_.predict(X_train)

	# R^2 for training set
	row.append(r2_score(y_train, y_pred_train))
	# R^2 for test set 
	row.append(r2_score(y_test, y_pred))

	y_train_data = y_train[:,0]

	# Pearson's correlation coefficient
	r = np.corrcoef(y_train_data, y_pred_train)
	row.append(r[0,1])
	r = np.corrcoef(y_test_data, y_pred)
	row.append(r[0,1])

	# Explained variance
	row.append(explained_variance_score(y_train,y_pred_train))
	row.append(explained_variance_score(y_test,y_pred))

	# The mean squared error
	#print('RMSE for train y: %.4f'
	#        % np.sqrt(mean_squared_error(y_train, y_pred_train)))
	#print('RMSE for test y: %.4f '
	#        % np.sqrt(mean_squared_error(y_test, y_pred)))
	#print()

	#Mean absolute error
	row.append(mean_absolute_error(y_train, y_pred_train))
	row.append(mean_absolute_error(y_test, y_pred))

	# Perform the inverse scaling transormations
	X_test_inv = Xscaler.inverse_transform(X_test)
	y_test_inv = Yscaler.inverse_transform(y_test.reshape(-1,1))
	y_train_inv = Yscaler.inverse_transform(y_train.reshape(-1,1))
	y_pred_inv = Yscaler.inverse_transform(y_pred.reshape(-1,1))
	y_pred_train_inv = Yscaler.inverse_transform(y_pred_train.reshape(-1,1))

	#Mean absolute error
	row.append(mean_absolute_error(y_train_inv, y_pred_train_inv))
	row.append(mean_absolute_error(y_test_inv, y_pred_inv))

	file1.write(','.join(map(str,row)))
	file1.write('\n')
	# Visualise the quality of fit
	#fig1, ax1 = plt.subplots(1,2,figsize=(6,3))
	#ax1[0].scatter(y_test_inv, y_pred_inv)
	#ax1[0].tick_params(direction='in', labelsize=18)
	#ax1[0].set_xlabel('$y_{test-inv}$', fontsize=18)
	#ax1[0].set_ylabel('$y_{pred-inv}$', fontsize=18)
	#ax1[0].set_xlim(-2,0)
	#ax1[0].set_ylim(-2,0)
	#ax1[1].scatter(y_test,y_pred)
	#ax1[1].set_xlabel('$y_{test}$', fontsize=18)
	#ax1[1].set_ylabel('$y_{pred}$', fontsize=18)
	#ax1[1].tick_params(direction='in', labelsize=18)
	#ax1[1].set_xlim(-2,0)
	#ax1[1].set_ylim(-2,0)
	#fig1.tight_layout()
	#plt.show()

	#N_features = np.shape(X_test)[1]
	#fig2, ax2 = plt.subplots(2, N_features, figsize=(16,4))
	#
	#for c in range(N_features):
	#	ax2[0,c].scatter(X_test[:,c], y_test, marker='o', s=16, label='$test$')
	#	ax2[0,c].scatter(X_test[:,c], y_pred, marker='s', s=16, label='$pred$')
	#	ax2[1,c].scatter(X_test_inv[:,c], y_test_inv, marker='o', s=16, label='$test$')
	#	ax2[1,c].scatter(X_test_inv[:,c], y_pred_inv, marker='s', s=16, label='$pred$')
	#	ax2[0,c].tick_params(direction='in', labelsize=16)
	#	ax2[1,c].tick_params(direction='in', labelsize=16)
	#	if c > 0:
	#		ax2[0,c].yaxis.set_visible(False)
	#		ax2[1,c].yaxis.set_visible(False)
	#	#ax2[1,c].set_xlabel(MAE_data.columns[c], fontsize=18)
	#	ax2[0,c].locator_params(axis='x', nbins=4)
	#	ax2[0,c].locator_params(axis='y', nbins=4)
	#	ax2[1,c].locator_params(axis='x', nbins=4)
	#	ax2[1,c].locator_params(axis='y', nbins=4)
	#
	#handles, labels = ax2[-1,-1].get_legend_handles_labels()
	#fig2.legend(handles, labels, fontsize=18, loc='upper center', bbox_to_anchor=(0.16,1))
	##ax2[0,0].set_ylabel('$scaled\\,\\, y$', fontsize=18)
	#ax2[0,0].set_ylabel('scaled y', fontsize=18)
	##ax2[1,0].set_ylabel('$log(MAE)$', fontsize=18)
	#ax2[1,0].set_ylabel('log(MAE)', fontsize=18)
	##fig2.tight_layout()
	#plt.show()

