from sklearn import linear_model, svm, kernel_ridge, preprocessing
from sklearn.metrics import mean_squared_error, r2_score, mean_absolute_error, make_scorer
from sklearn.model_selection import GridSearchCV, train_test_split, ShuffleSplit, cross_val_score
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import Normalizer, normalize, FunctionTransformer
import matplotlib.pyplot as plt
import numpy as np
import h5py
import time


class MyTransform_euc_mnh : #class for KRR additional normalisation
    def __init__(self,mode):
        self.mode=mode
    def fit (self, array, y=None): 
        if self.mode=="euc":
            self.denom=np.sum(np.square(array), axis=0)
        else:
            print ("Manhattan norm")
            self.denom=np.sum(abs(array), axis=0)
        return self
    def transform (self, array, y=None):
        arr2=np.divide(array, self.denom)
        return np.nan_to_num(arr2)


f=h5py.File('your_file.hdf5', 'r')
features=np.array(f['your_coulomb_matrixes'],  dtype=np.float64)
features=features.reshape(len(features),30*30)
f.close()

prop=[]
with open('your_file', 'r') as f2:
    for line in f2:
        prop.append(line.strip().split()[0])
prop=np.array(prop,  dtype=np.float64)


#merge features with properties and randomly shuffle the array
arr=np.append(prop[:,None], features, axis=1)
np.random.shuffle(arr)
data=arr[:,1:]
prop=arr[:,0]
data_train=np.delete(data,np.arange(0,len(data), 10),0)
data_test=data[::10]
prop_train=np.delete(prop,np.arange(0,len(data), 10),0)
prop_test=prop[::10]


########################### ELASTIC NET
###grid_search
alpha=np.power(10, np.arange(-10,1,dtype=float))
params=dict(estim__alpha=alpha)
print (params)
model1=linear_model.ElasticNet()
scaler=preprocessing.StandardScaler()
pipe=Pipeline([('scaler', scaler), ('estim', model1)])
grid=GridSearchCV(pipe, params, scoring='neg_mean_absolute_error', cv=*, n_jobs=-1, return_train_score=False, refit=False)
grid.fit(data, prop)
print ("The best parameters are {} with MAE {:.2f}". format(grid.best_params_, abs(grid.best_score_)))

###cross-validation 
scaler = preprocessing.StandardScaler()
model1 = linear_model.ElasticNet(alpha=your_alpha, normalize=False, fit_intercept=True, max_iter=1000)
pipe = Pipeline([("scaler", scaler), ("Est", model1)])
print(cross_val_score(pipe, data, prop, cv=ShuffleSplit(*,train_size=*,test_size=*), n_jobs=-1, 
                      scoring=make_scorer(mean_absolute_error)))


################## KRR
### grid search RBF
gamma=np.power(2,np.arange(-4,12,dtype=float))
params=dict(estim__gamma=gamma)
model2=kernel_ridge.KernelRidge(alpha=1e-09, kernel='rbf')
scaler=preprocessing.StandardScaler()
norm=MyTransform_euc_mnh(mode="euc")
pipe=Pipeline([('scaler', scaler),('norm', norm),('estim', model2)])
grid=GridSearchCV(pipe, params, scoring='neg_mean_absolute_error', cv=*, n_jobs=-1, return_train_score=False, refit=False)
grid.fit(data_test, prop_test)
print ("The best parameters are {} with MAE {:.2f}". format(grid.best_params_, abs(grid.best_score_)))

### grid search LPS
gamma=np.power(2,np.arange(-4,13,dtype=float))
params=dict(estim__gamma=gamma)
model2=kernel_ridge.KernelRidge(alpha=1e-09, kernel='laplacian')
scaler=preprocessing.StandardScaler()
norm=MyTransform_euc_mnh(mode="mnh")
pipe=Pipeline([('scaler', scaler),('norm', norm),('estim', model2)])
grid=GridSearchCV(pipe, params, scoring='neg_mean_absolute_error', cv=*, n_jobs=-1, return_train_score=False, refit=False)
grid.fit(data_test, prop_test)
print ("The best parameters are {} with MAE {:.2f}". format(grid.best_params_, abs(grid.best_score_)))

### cross-validation 
model2=kernel_ridge.KernelRidge(alpha=1e-09, kernel=*, gamma=*)
scaler=preprocessing.StandardScaler()
norm=MyTransform_euc_mnh(mode="*")
pipe=Pipeline([('scaler', scaler),('norm', norm),('estim', model2)])
print(cross_val_score(pipe, data, prop, cv=ShuffleSplit(*,train_size=*, test_size=*), n_jobs=-1, 
                      scoring=make_scorer(mean_absolute_error)))