import pandas as pd
import numpy as np
import sklearn as sk
import os
import matplotlib.pyplot as plt
from pandas.plotting import scatter_matrix
from sklearn.model_selection import StratifiedShuffleSplit
from sklearn.base import BaseEstimator, TransformerMixin
from sklearn.pipeline import Pipeline, FeatureUnion
from sklearn.preprocessing import OneHotEncoder, LabelBinarizer, StandardScaler
from data_preparation import data_reader, save_data, load_data
import periodictable as pt
import matplotlib
import re
from mendeleev import elements as el
from numpy.linalg import norm
from mendeleev import get_table
# from mendeleev.fetch import fetch_table


pd.options.mode.chained_assignment = None

matplotlib.use('TKAgg')
pd.set_option('display.max_rows', 500)
pd.set_option('display.max_columns', 500)
pd.set_option('display.width', 1000)

# To initialize the periodic elemental table for the further calculation
element_dict = {}
for i, j in enumerate(pt.elements):
    element_dict.update({'{}'.format(i): j})

element_df = get_table('elements')
e_conf = element_df['electronic_configuration']


def stratified_data(raw_df):

    # data = raw_df.iloc[:, [i for i in range(28, 108)] + [24, 20, 21, 22, 23]]
    data = raw_df
    data.loc[:, 'YEAR_cat'] = np.ceil(
        (data.loc[:, 'Year']-data.loc[:, 'Year'].min()+1)/
        (data.loc[:, 'Year'].max()-data.loc[:, 'Year'].min()+1)*5
                                      )
    data.loc[:, 'Na_cat'] = np.ceil(data.loc[:, '11'] / data.loc[:, '11'].max()*5)
    return data


def stratify_split(data_cat_added):
    split = StratifiedShuffleSplit(n_splits=1, test_size=0.2, random_state=42)
    for train_index, test_index in split.split(data_cat_added, data_cat_added['Na_cat']):
        strat_train_data = data_cat_added.loc[train_index]
        strat_test_data = data_cat_added.loc[test_index]

    # for set in (strat_train_data, strat_test_data):
    #     set.drop(['Na_cat'], axis=1, inplace=True)
    #     set.drop(['Year'], axis=1, inplace=True)

    return strat_train_data, strat_test_data


class DataframeSelector(BaseEstimator, TransformerMixin):   # transform the data from pd.df to array
    def __init__(self, attribute_names):
        self.attribute_names = attribute_names

    def fit(self, X, y=None):
        return self

    def transform(self, X, y=None):
        return X[self.attribute_names].values


class CombinedAttributesAdderNBO(BaseEstimator, TransformerMixin):  # ADD NBO ratio

    def __init__(self, X_list, add_total_alkali=True, ):
        self.add_total_alkali = add_total_alkali
        self.X_list = np.array(X_list).astype('object')

    def fit(self, X, y=None):
        return self

    def transform(self, X, y=None):
        Li =  np.where(self.X_list == '3')[0][0]  # (array([i], dtype=int64))[0][0]
        Na = np.where(self.X_list == '11')[0][0]
        K =  np.where(self.X_list == '19')[0][0]
        Rb = np.where(self.X_list == '37')[0][0]
        Cs = np.where(self.X_list == '55')[0][0]
        Mg = np.where(self.X_list == '12')[0][0]
        Ca = np.where(self.X_list == '20')[0][0]
        Sr = np.where(self.X_list == '38')[0][0]
        Ba = np.where(self.X_list == '56')[0][0]
        Si = np.where(self.X_list == '14')[0][0]
        Al = np.where(self.X_list == '13')[0][0]
        B =   np.where(self.X_list == '5')[0][0]
        if self.add_total_alkali:
            NBO_T_temp = (X[:, Li] + X[:, Na] + X[:, K] + X[:, Rb] +X[:, Cs] + 2*X[:, Mg] + 2*X[:, Ca] + 2*X[:, Sr] +
            2*X[:, Ba] -X[:, Al]-X[:, B])/(X[:, Si]+X[:, Al]+X[:, B])
            NBO_T = np.where(NBO_T_temp < 0, 0, NBO_T_temp)
            return np.c_[X, NBO_T]
        else:
            return X


class CombinedAttributesAdderAtomicNumber(BaseEstimator, TransformerMixin):
    # the order of features added is max, min, mean, std, mode1, mode2

    def __init__(self, X_list, el_list, add_atomic_number=True):
        self.add_atomic_number = add_atomic_number
        self.X_list = np.array(X_list).astype('object')
        self.el_list = el_list
        self.atomic_number_list = np.array(
            [
                self.el_list[i].atomic_number for i in range(len(self.el_list))
            ]
        )

    def fit(self, X, y=None):
        return self

    def transform(self, X, y=None):
        if self.add_atomic_number:   # x[i, :80] select all elements except other features
            atomic_number_max = np.array(
                [
                    int(self.X_list[max(np.where(X[i, :80] > 0)[0])]) for i in range(X.shape[0])
                ]
            )
            atomic_number_min = np.array(
                [
                    int(self.X_list[min(np.where(X[i, :80] > 0)[0])]) for i in range(X.shape[0])
                ]
            )
            atomic_number_mean = np.matmul(X[:, :80], self.atomic_number_list)
            atomic_number_std = np.array(
                [
                    np.matmul(X[i, :80], abs(self.atomic_number_list-atomic_number_mean[i]))
                    for i in range(X.shape[0])
                ]
            )
            atomic_mode_1 = np.array(
                [
                    self.atomic_number_list[np.where(X[i, :80]==np.sort(X[i, :80])[-1])[0][0]]
                    for i in range(X.shape[0])
                ]
            )
            atomic_mode_2 = np.array(
                [
                    self.atomic_number_list[np.where(X[i, :80]==np.sort(X[i, :80])[-2])[0][0]]
                    for i in range(X.shape[0])
                ]
            )
            return np.c_[X, atomic_number_max, atomic_number_min, atomic_number_mean,
                         atomic_number_std, atomic_mode_1, atomic_mode_2]
        else:
            return X


'''
the order of atomic property we used is:
'atomic_number', 'atomic_radius', 'mendeleev_number', 'atomic_weight', 'melting_point', 'period',
'series_id', 'covalent_radius_cordero', 'en_allen', 
'''


class CombinedAttributesAdderAtomicProperty(BaseEstimator, TransformerMixin):
    # the order of features added is max, min, mean, std, mode1, mode2

    def __init__(self, X_list, property_id, mendeleev_df, add_atomic_property=True):
        self.add_atomic_property = add_atomic_property
        self.X_list = np.array(X_list).astype('object')  # in X_list, the number of H is '1'
        self.mendeleev_df = mendeleev_df
        self.property_id = property_id
        self.atomic_property_list_full = self.mendeleev_df[self.property_id].fillna(method='pad')
        self.atomic_property_list = self.atomic_property_list_full.iloc[
            self.X_list.astype('int')-1    # the row start with 0, e.g. row number of H is zero
        ]

    def fit(self, X, y=None):
        return self

    def transform(self, X, y=None):
        if self.add_atomic_property:  # x[i, :80] select all elements except other features
            atomic_property_max = np.array(
                [
                    self.atomic_property_list.iloc[
                        np.where(X[i, :80] > 0)[0]
                    ].max() for i in range(X.shape[0])
                ]
            )
            atomic_property_min = np.array(
                [
                    self.atomic_property_list.iloc[
                        np.where(X[i, :80] > 0)[0]
                    ].min() for i in range(X.shape[0])
                ]
            )
            atomic_property_mean = np.matmul(X[:, :80], self.atomic_property_list)
            atomic_property_std = np.array(
                [
                    np.matmul(X[i, :80], abs(self.atomic_property_list - atomic_property_mean[i]))
                    for i in range(X.shape[0])
                ]
            )
            atomic_mode_1 = np.array(
                [
                    self.atomic_property_list.iloc[np.where(X[i, :80] == np.sort(X[i, :80])[-1])[0][0]]
                    for i in range(X.shape[0])
                ]
            )
            atomic_mode_2 = np.array(
                [
                    self.atomic_property_list.iloc[np.where(X[i, :80] == np.sort(X[i, :80])[-2])[0][0]]
                    for i in range(X.shape[0])
                ]
            )
            return np.c_[X, atomic_property_max, atomic_property_min, atomic_property_mean,
                         atomic_property_std, atomic_mode_1, atomic_mode_2]
        else:
            return X


'''
CombinedAttributesAdderElectronProperty
property id for electron configuration is 's', 'd', 'p', 'f', if filled=True, return valence electrons,
if filled=False, return unfilled states
'''


class CombinedAttributesAdderElectronProperty(BaseEstimator, TransformerMixin):

    def __init__(self, X_list, property_id, mendeleev_df, add_electron_property=True, filled=True):
        self.add_electron_property = add_electron_property
        self.filled = filled
        self.X_list = np.array(X_list).astype('object') # in X_list, the number of H is '1'
        self.mendeleev_df = mendeleev_df
        self.e_conf = mendeleev_df['electronic_configuration']
        self.property_id = property_id
        self.electron_dict = {'s': 2, 'd': 10, 'p': 6, 'f': 14}
        self.electron_list_full = []
        for element in self.e_conf:
            if re.findall(r'{}'.format(self.property_id), element):
                if re.findall(r'{}\d+'.format(self.property_id), element):
                    self.electron_list_full.append(
                        int(re.findall(r'{}\d+'.format(self.property_id), element)[0][1:])
                    )
                else:
                    self.electron_list_full.append(1)
            else:
                self.electron_list_full.append(0)
        self.electron_list = np.array(self.electron_list_full)[self.X_list.astype('int')-1]
        self.unfilled_electron_list = self.electron_dict.get(self.property_id)-self.electron_list

    def fit(self, X, y=None):
        return self

    def transform(self, X, y=None):
        if self.add_electron_property:  # x[i, :80] select all elements except other features
            if self.filled:
                electron_property_max = np.array(
                    [
                        self.electron_list[
                            np.where(X[i, :80] > 0)[0]
                        ].max() for i in range(X.shape[0])
                    ]
                )
                electron_property_min = np.array(
                    [
                        self.electron_list[
                            np.where(X[i, :80] > 0)[0]
                        ].min() for i in range(X.shape[0])
                    ]
                )
                electron_property_mean = np.matmul(X[:, :80], self.electron_list)
                electron_property_std = np.array(
                    [
                        np.matmul(X[i, :80], abs(self.electron_list - electron_property_mean[i]))
                        for i in range(X.shape[0])
                    ]
                )
                electron_mode_1 = np.array(
                    [
                        self.electron_list[np.where(X[i, :80] == np.sort(X[i, :80])[-1])[0][0]]
                        for i in range(X.shape[0])
                    ]
                )
                electron_mode_2 = np.array(
                    [
                        self.electron_list[np.where(X[i, :80] == np.sort(X[i, :80])[-2])[0][0]]
                        for i in range(X.shape[0])
                    ]
                )
                return np.c_[X, electron_property_max, electron_property_min, electron_property_mean,
                             electron_property_std, electron_mode_1, electron_mode_2]
            else:
                electron_property_max = np.array(
                    [
                        self.unfilled_electron_list[
                            np.where(X[i, :80] > 0)[0]
                        ].max() for i in range(X.shape[0])
                    ]
                )
                electron_property_min = np.array(
                    [
                        self.unfilled_electron_list[
                            np.where(X[i, :80] > 0)[0]
                        ].min() for i in range(X.shape[0])
                    ]
                )
                electron_property_mean = np.matmul(X[:, :80], self.unfilled_electron_list)
                electron_property_std = np.array(
                    [
                        np.matmul(X[i, :80], abs(self.unfilled_electron_list - electron_property_mean[i]))
                        for i in range(X.shape[0])
                    ]
                )
                electron_mode_1 = np.array(
                    [
                        self.unfilled_electron_list[np.where(X[i, :80] == np.sort(X[i, :80])[-1])[0][0]]
                        for i in range(X.shape[0])
                    ]
                )
                electron_mode_2 = np.array(
                    [
                        self.unfilled_electron_list[np.where(X[i, :80] == np.sort(X[i, :80])[-2])[0][0]]
                        for i in range(X.shape[0])
                    ]
                )
                return np.c_[X, electron_property_max, electron_property_min, electron_property_mean,
                             electron_property_std, electron_mode_1, electron_mode_2]
        else:
            return X


'''
add Lp norm features using np.linarg.norm(). p_list = [0, 2, 3, 5, 7, 10]
'''


class CombinedAttributesAdderLpNorm(BaseEstimator, TransformerMixin):

    def __init__(self, X_list, p_list, add_Lp_Norm=True):
        self.X_list = np.array(X_list).astype('object') # in X_list, the number of H is '1'
        self.p_list = p_list
        self.add_Lp_Norm = add_Lp_Norm

    def fit(self, X, y=None):
        return self

    def transform(self, X, y=None):
        if self.add_Lp_Norm:
            Lp_Norm = np.zeros([X.shape[0], len(self.p_list)])
            for num, i in enumerate(self.p_list):
                Lp_Norm[:, num] = norm(X[:, :80], i, axis=1)
            return np.c_[X, Lp_Norm]
        else:
            return X


class CombinedAttributesAdderOrbitalOccupation(BaseEstimator, TransformerMixin):

    def __init__(self, X_list, mendeleev_df, add_orbital_occupation=True):
        self.X_list = np.array(X_list).astype('object') # in X_list, the number of H is '1'
        self.add_orbital_occupation = add_orbital_occupation
        self.mendeleev_df = mendeleev_df
        self.property_ids = ['s', 'p', 'd', 'f']
        self.e_conf = mendeleev_df['electronic_configuration']
        self.electron_list_s_full = []
        self.electron_list_p_full = []
        self.electron_list_d_full = []
        self.electron_list_f_full = []
        for list, id in zip(
                (
                    self.electron_list_s_full,
                    self.electron_list_p_full,
                    self.electron_list_d_full,
                    self.electron_list_f_full
                ), self.property_ids
        ):
            for element in self.e_conf:
                if re.findall(r'{}'.format(id), element):
                    if re.findall(r'{}\d+'.format(id), element):
                        list.append(
                            int(re.findall(r'{}\d+'.format(id), element)[0][1:])
                        )
                    else:
                        list.append(1)
                else:
                    list.append(0)
        self.electron_list_s = np.array(self.electron_list_s_full)[self.X_list.astype('int')-1]
        self.electron_list_p = np.array(self.electron_list_p_full)[self.X_list.astype('int')-1]
        self.electron_list_d = np.array(self.electron_list_d_full)[self.X_list.astype('int')-1]
        self.electron_list_f = np.array(self.electron_list_f_full)[self.X_list.astype('int')-1]
        self.electron_list_total = self.electron_list_s + self.electron_list_p + \
                                   self.electron_list_d + self.electron_list_f

    def fit(self, X, y=None):
        return self

    def transform(self, X, y=None):
        if self.add_orbital_occupation:
            s = np.matmul(X[:, :80], self.electron_list_s)/\
                np.matmul(X[:, :80], self.electron_list_total)
            p = np.matmul(X[:, :80], self.electron_list_p)/\
                np.matmul(X[:, :80], self.electron_list_total)
            d = np.matmul(X[:, :80], self.electron_list_d)/\
                np.matmul(X[:, :80], self.electron_list_total)
            f = np.matmul(X[:, :80], self.electron_list_f)/\
                np.matmul(X[:, :80], self.electron_list_total)
            return np.c_[X, s, p, d, f]
        else:
            return X


class NumAttributesDropper(BaseEstimator, TransformerMixin):

    def __init__(self, X_list, drop_element_conc=True):
        self.X_list = np.array(X_list).astype('object') # in X_list, the number of H is '1'
        self.drop_element_conc = drop_element_conc

    def fit(self, X, y=None):
        return self

    def transform(self, X, y=None):
        if self.drop_element_conc:
            return X[:, 80:]
        else:
            return X


class MyLabelBinarizer(BaseEstimator, TransformerMixin):
    def __init__(self,):
        self.encoder = LabelBinarizer()

    def fit(self, X, y=0):
        return self.encoder.fit(X)

    def transform(self, X, y=0):
        return self.encoder.transform(X)


def main_element_base(DEG):
    raw_df = data_reader(path)
    name_list = list(raw_df.columns[28:108])
    drop_list = raw_df[raw_df.duplicated(name_list+['RO{}'.format(DEG)])].index.to_list()
    raw_df = raw_df.drop(drop_list, axis=0)
    name_list_ahead =list(raw_df.columns[0:28])
    cleaned_df = raw_df.groupby(name_list).mean().reset_index()
    cleaned_df = cleaned_df[name_list_ahead+name_list]
    data = stratified_data(cleaned_df)
    strat_train_data, strat_test_data = stratify_split(data)
    X_train_raw = strat_train_data.drop(['RO20', 'RO100', 'RO150', 'RO300'], axis=1)
    X_test_raw = strat_test_data.drop(['RO20', 'RO100', 'RO150', 'RO300'], axis=1)
    y_train = strat_train_data['RO{}'.format(DEG)].copy()
    y_test =strat_test_data['RO{}'.format(DEG)].copy()
    num_attribs = list(X_train_raw)[:-1]    # return a list of columns, str type. 80 elements
    print(len(num_attribs))
    cat_attribs = 'YEAR_cat'
    num_pipeline = Pipeline([
        ('selector', DataframeSelector(num_attribs)),
        # ('attribs_adder', CombinedAttributesAdderNBO(num_attribs)),
    ])
    cat_pipeline = Pipeline([
        ('selector', DataframeSelector(cat_attribs)),
        ('label_binarizer', MyLabelBinarizer()),
    ])
    full_pipeline = FeatureUnion([
        ('num_pipeline', num_pipeline),
        # ('cat_pipline', cat_pipeline),
    ])
    X_train = full_pipeline.fit_transform(X_train_raw)
    X_test = full_pipeline.fit_transform(X_test_raw)
    save_data(X_test,    'X_test_ele_based_clean_{}.pkl'.format(DEG)   )
    save_data(X_train,  'X_train_ele_based_clean_{}.pkl'.format(DEG))
    save_data(y_test,    'y_test_ele_based_clean_{}.pkl'.format(DEG)  )
    save_data(y_train,  'y_train_ele_based_clean_{}.pkl'.format(DEG) )

    df = get_table('elements')
    symbols_list = df['symbol'][[int(i)-1 for i in num_attribs]].to_list()

    return symbols_list


## raw=True for nomral data, and (raw=True AND test =True) for testing data
def main_property_base(DEG, path_name, raw=True, test=False):

    raw_df = data_reader(path_name)

    if raw:
        name_list = list(raw_df.columns[28:108])
        drop_list = raw_df[raw_df.duplicated(name_list + ['RO{}'.format(DEG)])].index.to_list()
        raw_df = raw_df.drop(drop_list, axis=0)
        name_list_ahead = list(raw_df.columns[0:28])
        cleaned_df = raw_df.groupby(name_list).mean().reset_index()
        cleaned_df = cleaned_df[name_list_ahead + name_list]
        if test:
            data = raw_df.iloc[:, [i for i in range(28, 108)] + [24, 20, 21, 22, 23]]
            X_test_raw = data.drop(['RO20', 'RO100', 'RO150', 'RO300'], axis=1)
            y_test = data['RO{}'.format(DEG)].copy()
            num_attribs= list(X_test_raw)[:-1]    # return a list of columns, str type. 80 elements
            print(len(num_attribs))
            print(num_attribs)
        else:
            data = stratified_data(cleaned_df)
            strat_train_data, strat_test_data = stratify_split(data)
            X_train_oxides = strat_train_data.iloc[:, 2:20]
            X_test_oxides = strat_test_data.iloc[:, 2:20]

            # X_train_raw = strat_train_data.drop(['RO20', 'RO100', 'RO150', 'RO300'], axis=1)
            # X_test_raw = strat_test_data.drop(['RO20', 'RO100', 'RO150', 'RO300'], axis=1)

            X_train_raw = strat_train_data.iloc[
                          :, [i for i in range(28, 108)] + [24, 20, 21, 22, 23]
                          ].drop(['RO20', 'RO100', 'RO150', 'RO300'], axis=1)
            print(X_train_raw.columns)
            X_test_raw = strat_test_data.iloc[
                         :, [i for i in range(28, 108)] + [24, 20, 21, 22, 23]
                         ].drop(['RO20', 'RO100', 'RO150', 'RO300'], axis=1)

            y_train = strat_train_data['RO{}'.format(DEG)].copy()
            y_test = strat_test_data['RO{}'.format(DEG)].copy()
            num_attribs = list(X_train_raw)[:-1]    # return a list of columns, str type. 80 elements
            print(len(num_attribs))
    else:
        data = raw_df.iloc[:, [i for i in range(28, 108)] + [24, 20, 21, 22, 23]]
        X_train_raw = data.drop(['RO20', 'RO100', 'RO150', 'RO300'], axis=1)
        X_test_raw = data.drop(['RO20', 'RO100', 'RO150', 'RO300'], axis=1)
        y_train = data['RO{}'.format(DEG)].copy()
        y_test = data['RO{}'.format(DEG)].copy()
        num_attribs = list(X_train_raw)[:-1]

    cat_attribs = 'YEAR_cat'
    num_pipeline = Pipeline([
        ('selector', DataframeSelector(num_attribs)),  # 80
        ('Lp_norm', CombinedAttributesAdderLpNorm(
            num_attribs, [0, 2, 3, 5, 7, 10],                            add_Lp_Norm=False )),    # 6 attributes
        ('attribs_adder', CombinedAttributesAdderNBO(num_attribs,   add_total_alkali=False)),  # 1 attributes
        ('atomic_num_adder', CombinedAttributesAdderAtomicProperty(  # 6 attributes
            num_attribs, 'atomic_number', element_df,            add_atomic_property=False)),
        ('atomic_radius', CombinedAttributesAdderAtomicProperty(
            num_attribs, 'atomic_radius', element_df,            add_atomic_property=False )),       # 6 attributes
        ('mendeleev_number', CombinedAttributesAdderAtomicProperty(
            num_attribs, 'mendeleev_number', element_df,         add_atomic_property=True)),   # 6 attributes
        ('atomic_weight', CombinedAttributesAdderAtomicProperty(
            num_attribs, 'atomic_weight', element_df,            add_atomic_property=False)),      # 6 attributes
        ('melting_point', CombinedAttributesAdderAtomicProperty(
            num_attribs, 'melting_point', element_df,            add_atomic_property=False)),      # 6 attributes
        ('period', CombinedAttributesAdderAtomicProperty(
            num_attribs, 'period', element_df,                   add_atomic_property=False)),      # 6 attributes
        ('series_id', CombinedAttributesAdderAtomicProperty(
            num_attribs, 'series_id', element_df,                add_atomic_property=False)),    # 6 attributes
        ('covalent_radius_cordero', CombinedAttributesAdderAtomicProperty(
            num_attribs, 'covalent_radius_cordero', element_df,  add_atomic_property=False)),   # 6 attributes
        ('en_allen', CombinedAttributesAdderAtomicProperty(
            num_attribs, 'en_allen', element_df,                 add_atomic_property=False)),    # 6 attributes
        ('s_valence', CombinedAttributesAdderElectronProperty(
            num_attribs, 's', element_df,                      add_electron_property=True)), # 6 attributes
        ('p_valence', CombinedAttributesAdderElectronProperty(
            num_attribs, 'p', element_df,                      add_electron_property=False)), # 6 attributes
        ('d_valence', CombinedAttributesAdderElectronProperty(
            num_attribs, 'd', element_df,                      add_electron_property=False)), # 6 attributes
        ('f_valence', CombinedAttributesAdderElectronProperty(
            num_attribs, 'f', element_df,                      add_electron_property=False)), # 6 attributes
        ('s_unfilled', CombinedAttributesAdderElectronProperty(
            num_attribs, 's', element_df, filled=False,        add_electron_property=True )), # 6 attributes
        ('p_unfilled', CombinedAttributesAdderElectronProperty(
            num_attribs, 'p', element_df, filled=False,        add_electron_property=False)), # 6 attributes
        ('d_unfilled', CombinedAttributesAdderElectronProperty(
            num_attribs, 'd', element_df, filled=False,        add_electron_property=False)), # 6 attributes
        ('f_unfilled', CombinedAttributesAdderElectronProperty(
            num_attribs, 'f', element_df, filled=False,        add_electron_property=False)), # 6 attributes
        ('orbital_occupation', CombinedAttributesAdderOrbitalOccupation(
            num_attribs, element_df,                          add_orbital_occupation=True )),    # 4 attributes
        ('num_dropper', NumAttributesDropper(num_attribs,          drop_element_conc=False)),
        # ('std_scaler', StandardScaler()),
    ])

    cat_pipeline = Pipeline([
        ('selector', DataframeSelector(cat_attribs)),
        ('label_binarizer', MyLabelBinarizer()),   # 5 attributes
    ])
    full_pipeline = FeatureUnion([
        ('num_pipeline', num_pipeline),  # 80 + 113
        # ('cat_pipline', cat_pipeline),  # 5
    ])
    if test:
        X_test = full_pipeline.fit_transform(X_test_raw)
        save_data(X_test, 'X_test_{}.pkl'.format(DEG))
    else:
        X_train = full_pipeline.fit_transform(X_train_raw)
        X_test = full_pipeline.fit_transform(X_test_raw)
        X_train_final = np.c_[X_train_oxides.values, X_train[:, 80:]]
        X_test_final = np.c_[X_test_oxides.values, X_test[:, 80:]]
        if raw:
            save_data(X_test,    'X_test_oxides_property_{}.pkl'.format(DEG)   )
            save_data(X_train,  'X_train_oxides_property_{}.pkl'.format(DEG))
            save_data(y_test,    'y_test_oxides_property_{}.pkl'.format(DEG)  )
            save_data(y_train,  'y_train_oxides_property_{}.pkl'.format(DEG) )
        else:
            save_data(X_test,    'X_test_product.pkl'.format(DEG)   )
            save_data(X_train,  'X_train_product.pkl'.format(DEG))
            save_data(y_test,    'y_test_product.pkl'.format(DEG)  )
            save_data(y_train,  'y_train_product.pkl'.format(DEG) )

    df = get_table('elements')
    symbols_list = df['symbol'][[int(i)-1 for i in num_attribs]].to_list()

    return symbols_list


def oxide_property_data_prep(target_DEG, raw_data_path):

    DEGs = np.array([100, 150, 300])

    path_300 = os.path.join(raw_data_path, 'worksheet_{}.xlsx'.format(300))
    path_150 = os.path.join(raw_data_path, 'worksheet_{}.xlsx'.format(150))
    path_100 = os.path.join(raw_data_path, 'worksheet_{}.xlsx'.format(100))

    clean_df_list = []
    X_train_list = []
    X_test_list = []
    y_train_list = []
    y_test_list = []

    raw_df_list = [data_reader(path_100), data_reader(path_150), data_reader(path_300)]

    num_attribs = list(raw_df_list[0].iloc[:, 28:108].columns)
    cat_attribs = 'YEAR_cat'
    num_pipeline = Pipeline([
        ('selector', DataframeSelector(num_attribs)),  # 80
        ('Lp_norm', CombinedAttributesAdderLpNorm(
            num_attribs, [0, 2, 3, 5, 7, 10],                            add_Lp_Norm=False )),    # 6 attributes
        ('attribs_adder', CombinedAttributesAdderNBO(num_attribs,   add_total_alkali=False)),  # 1 attributes
        ('atomic_num_adder', CombinedAttributesAdderAtomicProperty(  # 6 attributes
            num_attribs, 'atomic_number', element_df,            add_atomic_property=False)),
        ('atomic_radius', CombinedAttributesAdderAtomicProperty(
            num_attribs, 'atomic_radius', element_df,            add_atomic_property=False )),       # 6 attributes
        ('mendeleev_number', CombinedAttributesAdderAtomicProperty(
            num_attribs, 'mendeleev_number', element_df,         add_atomic_property=True)),   # 6 attributes
        ('atomic_weight', CombinedAttributesAdderAtomicProperty(
            num_attribs, 'atomic_weight', element_df,            add_atomic_property=False)),      # 6 attributes
        ('melting_point', CombinedAttributesAdderAtomicProperty(
            num_attribs, 'melting_point', element_df,            add_atomic_property=False)),      # 6 attributes
        ('period', CombinedAttributesAdderAtomicProperty(
            num_attribs, 'period', element_df,                   add_atomic_property=False)),      # 6 attributes
        ('series_id', CombinedAttributesAdderAtomicProperty(
            num_attribs, 'series_id', element_df,                add_atomic_property=False)),    # 6 attributes
        ('covalent_radius_cordero', CombinedAttributesAdderAtomicProperty(
            num_attribs, 'covalent_radius_cordero', element_df,  add_atomic_property=False)),   # 6 attributes
        ('en_allen', CombinedAttributesAdderAtomicProperty(
            num_attribs, 'en_allen', element_df,                 add_atomic_property=False)),    # 6 attributes
        ('s_valence', CombinedAttributesAdderElectronProperty(
            num_attribs, 's', element_df,                      add_electron_property=True)), # 6 attributes
        ('p_valence', CombinedAttributesAdderElectronProperty(
            num_attribs, 'p', element_df,                      add_electron_property=False)), # 6 attributes
        ('d_valence', CombinedAttributesAdderElectronProperty(
            num_attribs, 'd', element_df,                      add_electron_property=False)), # 6 attributes
        ('f_valence', CombinedAttributesAdderElectronProperty(
            num_attribs, 'f', element_df,                      add_electron_property=False)), # 6 attributes
        ('s_unfilled', CombinedAttributesAdderElectronProperty(
            num_attribs, 's', element_df, filled=False,        add_electron_property=True )), # 6 attributes
        ('p_unfilled', CombinedAttributesAdderElectronProperty(
            num_attribs, 'p', element_df, filled=False,        add_electron_property=False)), # 6 attributes
        ('d_unfilled', CombinedAttributesAdderElectronProperty(
            num_attribs, 'd', element_df, filled=False,        add_electron_property=False)), # 6 attributes
        ('f_unfilled', CombinedAttributesAdderElectronProperty(
            num_attribs, 'f', element_df, filled=False,        add_electron_property=False)), # 6 attributes
        ('orbital_occupation', CombinedAttributesAdderOrbitalOccupation(
            num_attribs, element_df,                          add_orbital_occupation=True )),    # 4 attributes
        ('num_dropper', NumAttributesDropper(num_attribs,          drop_element_conc=False)),
        # ('std_scaler', StandardScaler()),
    ])

    cat_pipeline = Pipeline([
        ('selector', DataframeSelector(cat_attribs)),
        ('label_binarizer', MyLabelBinarizer()),   # 5 attributes
    ])
    full_pipeline = FeatureUnion([
        ('num_pipeline', num_pipeline),  # 80 + 113
        # ('cat_pipline', cat_pipeline),  # 5
    ])

    for i, raw_df in enumerate(raw_df_list):
        element_name_list = list(raw_df.columns[28:108])
        drop_list = raw_df[raw_df.duplicated(element_name_list + ['RO{}'.format(DEGs[i])])].index.to_list()
        raw_df = raw_df.drop(drop_list, axis=0)
        oxides_list_ahead = list(raw_df.columns[0:28])
        cleaned_df = raw_df.groupby(element_name_list).mean().reset_index()
        cleaned_df = cleaned_df[oxides_list_ahead + element_name_list]
        clean_df_list.append(cleaned_df)

        data = stratified_data(cleaned_df)
        strat_train_data, strat_test_data = stratify_split(data)
        X_train_oxides = strat_train_data.iloc[:, 2:20]
        X_test_oxides = strat_test_data.iloc[:, 2:20]


        X_train_raw = strat_train_data.iloc[
                      :, [i for i in range(28, 108)] + [24, 20, 21, 22, 23]
                      ].drop(['RO20', 'RO100', 'RO150', 'RO300'], axis=1)
        print(X_train_raw.columns)
        X_test_raw = strat_test_data.iloc[
                     :, [i for i in range(28, 108)] + [24, 20, 21, 22, 23]
                     ].drop(['RO20', 'RO100', 'RO150', 'RO300'], axis=1)

        y_train = strat_train_data['RO{}'.format(DEGs[i])].copy()
        y_test = strat_test_data['RO{}'.format(DEGs[i])].copy()

        X_train = full_pipeline.fit_transform(X_train_raw)
        X_test = full_pipeline.fit_transform(X_test_raw)
        X_train_final = np.c_[X_train_oxides.values, X_train[:, 80:]]
        X_test_final = np.c_[X_test_oxides.values, X_test[:, 80:]]

        X_train_list.append(X_train_final)
        X_test_list.append(X_test_final)
        y_train_list.append(y_train)
        y_test_list.append(y_test)


    ## avoid data leakage

    X_test_target_deg = X_test_list[np.where(DEGs==target_DEG)[0][0]]

    for num in np.where(DEGs!=target_DEG)[0]:
        mask = []
        for row in X_train_list[num]:
            mask.append(any(np.equal(X_test_target_deg, row).all(1)))

        (bool_val, counts) = np.unique(mask, return_counts=True)

        print('{} leakage points found for dataset at {} DEG'.format(counts[1], DEGs[num]))
        X_train_list[num] = X_train_list[num][~np.array(mask)]
        y_train_list[num] = y_train_list[num][~np.array(mask)]

    ## output the datasets

    for num, DEG in enumerate(DEGs):
        save_data(
            X_train_list[num],
            './datasets/target_{}/X_train_oxides_property_{}.pkl'.format(target_DEG, DEG)
        )
        save_data(
            X_test_list[num],
            './datasets/target_{}/X_test_oxides_property_{}.pkl'.format(target_DEG, DEG)
        )

        save_data(
            y_train_list[num],
            './datasets/target_{}/y_train_oxides_property_{}.pkl'.format(target_DEG, DEG)
        )
        save_data(
            y_test_list[num],
            './datasets/target_{}/y_test_oxides_property_{}.pkl'.format(target_DEG, DEG)
        )


# function to create raw data composition file
def test_raw_data_preparation(component_list, prime_component_molar_fraction):
    ## e.g. [SIO2, NA2O, K2O], molar fraction in %
    column_list = np.array(['GLASNO', 'A_WT', 'SIO2', 'AL2O3', 'B2O3', 'CAO', 'K2O', 'NA2O', 'PBO', 'Li2O', 'MgO',
                   'SRO', 'BAO', 'ZNO', 'P2O5', 'GEO2', 'ZRO2', 'TIO2', 'TEO2', 'FemOn', 'RO20', 'RO100',
                   'RO150', 'RO300', 'Year', 'Refer_Id', 'Patent', 'SLeg'])
    if len(component_list)==3:
        temp_data = np.zeros([200, len(column_list)])
        temp_data[:, np.where(column_list==component_list[0])[0][0]] = prime_component_molar_fraction
        temp_data[:, np.where(column_list==component_list[1])[0][0]]=np.linspace(0,
                                                                           100-prime_component_molar_fraction,
                                                                           temp_data.shape[0])
        temp_data[:, np.where(column_list==component_list[2])] = 100 - \
                                                                 temp_data[:, np.where(column_list==component_list[1])]-\
                                                                 prime_component_molar_fraction
        test_raw_data_df = pd.DataFrame(temp_data, columns=column_list)
        test_raw_data_df.to_csv(r'./data/testing data/mixed alkali effect/raw_data_{0}_{1}.csv'
                                .format(component_list, prime_component_molar_fraction))
    else:
        print('error! please check the inputs!')
    return test_raw_data_df


def product_data_preparation(path, DEG):
    raw_df = pd.read_excel(path, engine='openpyxl')

    num_attribs = list(raw_df.iloc[:, 28:108].columns)
    cat_attribs = 'YEAR_cat'
    num_pipeline = Pipeline([
        ('selector', DataframeSelector(num_attribs)),  # 80
        ('Lp_norm', CombinedAttributesAdderLpNorm(
            num_attribs, [0, 2, 3, 5, 7, 10],                            add_Lp_Norm=False )),    # 6 attributes
        ('attribs_adder', CombinedAttributesAdderNBO(num_attribs,   add_total_alkali=False)),  # 1 attributes
        ('atomic_num_adder', CombinedAttributesAdderAtomicProperty(  # 6 attributes
            num_attribs, 'atomic_number', element_df,            add_atomic_property=False)),
        ('atomic_radius', CombinedAttributesAdderAtomicProperty(
            num_attribs, 'atomic_radius', element_df,            add_atomic_property=False )),       # 6 attributes
        ('mendeleev_number', CombinedAttributesAdderAtomicProperty(
            num_attribs, 'mendeleev_number', element_df,         add_atomic_property=True)),   # 6 attributes
        ('atomic_weight', CombinedAttributesAdderAtomicProperty(
            num_attribs, 'atomic_weight', element_df,            add_atomic_property=False)),      # 6 attributes
        ('melting_point', CombinedAttributesAdderAtomicProperty(
            num_attribs, 'melting_point', element_df,            add_atomic_property=False)),      # 6 attributes
        ('period', CombinedAttributesAdderAtomicProperty(
            num_attribs, 'period', element_df,                   add_atomic_property=False)),      # 6 attributes
        ('series_id', CombinedAttributesAdderAtomicProperty(
            num_attribs, 'series_id', element_df,                add_atomic_property=False)),    # 6 attributes
        ('covalent_radius_cordero', CombinedAttributesAdderAtomicProperty(
            num_attribs, 'covalent_radius_cordero', element_df,  add_atomic_property=False)),   # 6 attributes
        ('en_allen', CombinedAttributesAdderAtomicProperty(
            num_attribs, 'en_allen', element_df,                 add_atomic_property=False)),    # 6 attributes
        ('s_valence', CombinedAttributesAdderElectronProperty(
            num_attribs, 's', element_df,                      add_electron_property=True)), # 6 attributes
        ('p_valence', CombinedAttributesAdderElectronProperty(
            num_attribs, 'p', element_df,                      add_electron_property=False)), # 6 attributes
        ('d_valence', CombinedAttributesAdderElectronProperty(
            num_attribs, 'd', element_df,                      add_electron_property=False)), # 6 attributes
        ('f_valence', CombinedAttributesAdderElectronProperty(
            num_attribs, 'f', element_df,                      add_electron_property=False)), # 6 attributes
        ('s_unfilled', CombinedAttributesAdderElectronProperty(
            num_attribs, 's', element_df, filled=False,        add_electron_property=True )), # 6 attributes
        ('p_unfilled', CombinedAttributesAdderElectronProperty(
            num_attribs, 'p', element_df, filled=False,        add_electron_property=False)), # 6 attributes
        ('d_unfilled', CombinedAttributesAdderElectronProperty(
            num_attribs, 'd', element_df, filled=False,        add_electron_property=False)), # 6 attributes
        ('f_unfilled', CombinedAttributesAdderElectronProperty(
            num_attribs, 'f', element_df, filled=False,        add_electron_property=False)), # 6 attributes
        ('orbital_occupation', CombinedAttributesAdderOrbitalOccupation(
            num_attribs, element_df,                          add_orbital_occupation=True )),    # 4 attributes
        ('num_dropper', NumAttributesDropper(num_attribs,          drop_element_conc=False)),
        # ('std_scaler', StandardScaler()),
    ])

    full_pipeline = FeatureUnion([
        ('num_pipeline', num_pipeline),  # 80 + 113
    ])

    X_product_elements = raw_df.iloc[:, [i for i in range(28, 108)] + [24, 20, 21, 22, 23]
                      ].drop(['RO20', 'RO100', 'RO150', 'RO300'], axis=1)
    X_product_oxides = raw_df.iloc[:, 2:20]
    y_product = raw_df['RO{}'.format(DEG)].copy()
    X_product = np.c_[
        X_product_oxides,
        full_pipeline.fit_transform(X_product_elements)[:, 80:]
    ]

    save_data(X_product, './datasets/X_product.pkl')
    save_data(y_product, './datasets/y_product.pkl')

    return X_product, y_product

property_list = ['Lp_0', 'Lp_2', 'Lp_3', 'Lp_5', 'Lp_7', 'Lp_10',
                 'NBO',
                 'a_num_max', 'a_num_min', 'a_num_mean', 'a_num_std', 'a_num_1', 'a_num_2',
                 'a_radii_max', 'a_radii_min', 'a_radii_mean', 'a_radii_std', 'a_radii_1', 'a_radii_2',
                 'me_num_max', 'me_num_min', 'me_num_mean', 'me_num_std', 'me_num_1', 'me_num_2',
                 'a_wt_max', 'a_wt_min', 'a_wt_mean', 'a_wt_std', 'a_wt_1', 'a_wt_2',
                 'a_mel_max', 'a_mel_min', 'a_mel_mean', 'a_mel_std', 'a_mel_1', 'a_mel_2',
                 'a_per_max', 'a_per_min', 'a_per_mean', 'a_per_std', 'a_per_1', 'a_per_2',
                 'a_sid_max', 'a_sid_min', 'a_sid_mean', 'a_sid_std', 'a_sid_1', 'a_sid_2',
                 'a_cra_max', 'a_cra_min', 'a_cra_mean', 'a_cra_std', 'a_cra_1', 'a_cra_2',
                 'a_ale_max', 'a_ale_min', 'a_ale_mean', 'a_ale_std', 'a_ale_1', 'a_ale_2',
                 'e_sf_max', 'e_sf_min', 'e_sf_mean', 'e_sf_std', 'e_sf_1', 'e_sf_2',
                 'e_pf_max', 'e_pf_min', 'e_pf_mean', 'e_pf_std', 'e_pf_1', 'e_pf_2',
                 'e_df_max', 'e_df_min', 'e_df_mean', 'e_df_std', 'e_df_1', 'e_df_2',
                 'e_ff_max', 'e_ff_min', 'e_ff_mean', 'e_ff_std', 'e_ff_1', 'e_ff_2',
                 'e_suf_max', 'e_suf_min', 'e_suf_mean', 'e_suf_std', 'e_suf_1', 'e_suf_2',
                 'e_puf_max', 'e_puf_min', 'e_puf_mean', 'e_puf_std', 'e_puf_1', 'e_puf_2',
                 'e_duf_max', 'e_duf_min', 'e_duf_mean', 'e_duf_std', 'e_duf_1', 'e_duf_2',
                 'e_fuf_max', 'e_fuf_min', 'e_fuf_mean', 'e_fuf_std', 'e_fuf_1', 'e_fuf_2',
                 'oc_s', 'oc_p', 'oc_d', 'oc_f'
                 ]


'''
the order of atomic property we used is:
'atomic_number', 'atomic_radius', 'mendeleev_number', 'atomic_weight', 'melting_point', 'period',
'series_id', 'covalent_radius_cordero', 'en_allen', 
'''

if __name__ == '__main__':

    # na =main_element_base(300)
    Temperature = 300
    # path_product = r'./data/worksheet_product_{}.xlsx'.format(Temperature)
    path = r'./data/worksheet_product.xlsx'

    # test_raw_data_preparation(['SIO2', 'K2O', 'Li2O'], 67)
    X, y = product_data_preparation(path, Temperature)
import pandas as pd
import numpy as np
import sklearn as sk
import os
import matplotlib.pyplot as plt
from pandas.plotting import scatter_matrix
from sklearn.model_selection import StratifiedShuffleSplit
from sklearn.base import BaseEstimator, TransformerMixin
from sklearn.pipeline import Pipeline, FeatureUnion
from sklearn.preprocessing import OneHotEncoder, LabelBinarizer, StandardScaler
from data_preparation import data_reader, save_data, load_data
import periodictable as pt
import matplotlib
import re
from mendeleev import elements as el
from numpy.linalg import norm
from mendeleev import get_table
# from mendeleev.fetch import fetch_table


pd.options.mode.chained_assignment = None

matplotlib.use('TKAgg')
pd.set_option('display.max_rows', 500)
pd.set_option('display.max_columns', 500)
pd.set_option('display.width', 1000)

# To initialize the periodic elemental table for the further calculation
element_dict = {}
for i, j in enumerate(pt.elements):
    element_dict.update({'{}'.format(i): j})

element_df = get_table('elements')
e_conf = element_df['electronic_configuration']


def stratified_data(raw_df):

    # data = raw_df.iloc[:, [i for i in range(28, 108)] + [24, 20, 21, 22, 23]]
    data = raw_df
    data.loc[:, 'YEAR_cat'] = np.ceil(
        (data.loc[:, 'Year']-data.loc[:, 'Year'].min()+1)/
        (data.loc[:, 'Year'].max()-data.loc[:, 'Year'].min()+1)*5
                                      )
    data.loc[:, 'Na_cat'] = np.ceil(data.loc[:, '11'] / data.loc[:, '11'].max()*5)
    return data


def stratify_split(data_cat_added):
    split = StratifiedShuffleSplit(n_splits=1, test_size=0.2, random_state=42)
    for train_index, test_index in split.split(data_cat_added, data_cat_added['Na_cat']):
        strat_train_data = data_cat_added.loc[train_index]
        strat_test_data = data_cat_added.loc[test_index]

    # for set in (strat_train_data, strat_test_data):
    #     set.drop(['Na_cat'], axis=1, inplace=True)
    #     set.drop(['Year'], axis=1, inplace=True)

    return strat_train_data, strat_test_data


class DataframeSelector(BaseEstimator, TransformerMixin):   # transform the data from pd.df to array
    def __init__(self, attribute_names):
        self.attribute_names = attribute_names

    def fit(self, X, y=None):
        return self

    def transform(self, X, y=None):
        return X[self.attribute_names].values


class CombinedAttributesAdderNBO(BaseEstimator, TransformerMixin):  # ADD NBO ratio

    def __init__(self, X_list, add_total_alkali=True, ):
        self.add_total_alkali = add_total_alkali
        self.X_list = np.array(X_list).astype('object')

    def fit(self, X, y=None):
        return self

    def transform(self, X, y=None):
        Li =  np.where(self.X_list == '3')[0][0]  # (array([i], dtype=int64))[0][0]
        Na = np.where(self.X_list == '11')[0][0]
        K =  np.where(self.X_list == '19')[0][0]
        Rb = np.where(self.X_list == '37')[0][0]
        Cs = np.where(self.X_list == '55')[0][0]
        Mg = np.where(self.X_list == '12')[0][0]
        Ca = np.where(self.X_list == '20')[0][0]
        Sr = np.where(self.X_list == '38')[0][0]
        Ba = np.where(self.X_list == '56')[0][0]
        Si = np.where(self.X_list == '14')[0][0]
        Al = np.where(self.X_list == '13')[0][0]
        B =   np.where(self.X_list == '5')[0][0]
        if self.add_total_alkali:
            NBO_T_temp = (X[:, Li] + X[:, Na] + X[:, K] + X[:, Rb] +X[:, Cs] + 2*X[:, Mg] + 2*X[:, Ca] + 2*X[:, Sr] +
            2*X[:, Ba] -X[:, Al]-X[:, B])/(X[:, Si]+X[:, Al]+X[:, B])
            NBO_T = np.where(NBO_T_temp < 0, 0, NBO_T_temp)
            return np.c_[X, NBO_T]
        else:
            return X


class CombinedAttributesAdderAtomicNumber(BaseEstimator, TransformerMixin):
    # the order of features added is max, min, mean, std, mode1, mode2

    def __init__(self, X_list, el_list, add_atomic_number=True):
        self.add_atomic_number = add_atomic_number
        self.X_list = np.array(X_list).astype('object')
        self.el_list = el_list
        self.atomic_number_list = np.array(
            [
                self.el_list[i].atomic_number for i in range(len(self.el_list))
            ]
        )

    def fit(self, X, y=None):
        return self

    def transform(self, X, y=None):
        if self.add_atomic_number:   # x[i, :80] select all elements except other features
            atomic_number_max = np.array(
                [
                    int(self.X_list[max(np.where(X[i, :80] > 0)[0])]) for i in range(X.shape[0])
                ]
            )
            atomic_number_min = np.array(
                [
                    int(self.X_list[min(np.where(X[i, :80] > 0)[0])]) for i in range(X.shape[0])
                ]
            )
            atomic_number_mean = np.matmul(X[:, :80], self.atomic_number_list)
            atomic_number_std = np.array(
                [
                    np.matmul(X[i, :80], abs(self.atomic_number_list-atomic_number_mean[i]))
                    for i in range(X.shape[0])
                ]
            )
            atomic_mode_1 = np.array(
                [
                    self.atomic_number_list[np.where(X[i, :80]==np.sort(X[i, :80])[-1])[0][0]]
                    for i in range(X.shape[0])
                ]
            )
            atomic_mode_2 = np.array(
                [
                    self.atomic_number_list[np.where(X[i, :80]==np.sort(X[i, :80])[-2])[0][0]]
                    for i in range(X.shape[0])
                ]
            )
            return np.c_[X, atomic_number_max, atomic_number_min, atomic_number_mean,
                         atomic_number_std, atomic_mode_1, atomic_mode_2]
        else:
            return X


'''
the order of atomic property we used is:
'atomic_number', 'atomic_radius', 'mendeleev_number', 'atomic_weight', 'melting_point', 'period',
'series_id', 'covalent_radius_cordero', 'en_allen', 
'''


class CombinedAttributesAdderAtomicProperty(BaseEstimator, TransformerMixin):
    # the order of features added is max, min, mean, std, mode1, mode2

    def __init__(self, X_list, property_id, mendeleev_df, add_atomic_property=True):
        self.add_atomic_property = add_atomic_property
        self.X_list = np.array(X_list).astype('object')  # in X_list, the number of H is '1'
        self.mendeleev_df = mendeleev_df
        self.property_id = property_id
        self.atomic_property_list_full = self.mendeleev_df[self.property_id].fillna(method='pad')
        self.atomic_property_list = self.atomic_property_list_full.iloc[
            self.X_list.astype('int')-1    # the row start with 0, e.g. row number of H is zero
        ]

    def fit(self, X, y=None):
        return self

    def transform(self, X, y=None):
        if self.add_atomic_property:  # x[i, :80] select all elements except other features
            atomic_property_max = np.array(
                [
                    self.atomic_property_list.iloc[
                        np.where(X[i, :80] > 0)[0]
                    ].max() for i in range(X.shape[0])
                ]
            )
            atomic_property_min = np.array(
                [
                    self.atomic_property_list.iloc[
                        np.where(X[i, :80] > 0)[0]
                    ].min() for i in range(X.shape[0])
                ]
            )
            atomic_property_mean = np.matmul(X[:, :80], self.atomic_property_list)
            atomic_property_std = np.array(
                [
                    np.matmul(X[i, :80], abs(self.atomic_property_list - atomic_property_mean[i]))
                    for i in range(X.shape[0])
                ]
            )
            atomic_mode_1 = np.array(
                [
                    self.atomic_property_list.iloc[np.where(X[i, :80] == np.sort(X[i, :80])[-1])[0][0]]
                    for i in range(X.shape[0])
                ]
            )
            atomic_mode_2 = np.array(
                [
                    self.atomic_property_list.iloc[np.where(X[i, :80] == np.sort(X[i, :80])[-2])[0][0]]
                    for i in range(X.shape[0])
                ]
            )
            return np.c_[X, atomic_property_max, atomic_property_min, atomic_property_mean,
                         atomic_property_std, atomic_mode_1, atomic_mode_2]
        else:
            return X


'''
CombinedAttributesAdderElectronProperty
property id for electron configuration is 's', 'd', 'p', 'f', if filled=True, return valence electrons,
if filled=False, return unfilled states
'''


class CombinedAttributesAdderElectronProperty(BaseEstimator, TransformerMixin):

    def __init__(self, X_list, property_id, mendeleev_df, add_electron_property=True, filled=True):
        self.add_electron_property = add_electron_property
        self.filled = filled
        self.X_list = np.array(X_list).astype('object') # in X_list, the number of H is '1'
        self.mendeleev_df = mendeleev_df
        self.e_conf = mendeleev_df['electronic_configuration']
        self.property_id = property_id
        self.electron_dict = {'s': 2, 'd': 10, 'p': 6, 'f': 14}
        self.electron_list_full = []
        for element in self.e_conf:
            if re.findall(r'{}'.format(self.property_id), element):
                if re.findall(r'{}\d+'.format(self.property_id), element):
                    self.electron_list_full.append(
                        int(re.findall(r'{}\d+'.format(self.property_id), element)[0][1:])
                    )
                else:
                    self.electron_list_full.append(1)
            else:
                self.electron_list_full.append(0)
        self.electron_list = np.array(self.electron_list_full)[self.X_list.astype('int')-1]
        self.unfilled_electron_list = self.electron_dict.get(self.property_id)-self.electron_list

    def fit(self, X, y=None):
        return self

    def transform(self, X, y=None):
        if self.add_electron_property:  # x[i, :80] select all elements except other features
            if self.filled:
                electron_property_max = np.array(
                    [
                        self.electron_list[
                            np.where(X[i, :80] > 0)[0]
                        ].max() for i in range(X.shape[0])
                    ]
                )
                electron_property_min = np.array(
                    [
                        self.electron_list[
                            np.where(X[i, :80] > 0)[0]
                        ].min() for i in range(X.shape[0])
                    ]
                )
                electron_property_mean = np.matmul(X[:, :80], self.electron_list)
                electron_property_std = np.array(
                    [
                        np.matmul(X[i, :80], abs(self.electron_list - electron_property_mean[i]))
                        for i in range(X.shape[0])
                    ]
                )
                electron_mode_1 = np.array(
                    [
                        self.electron_list[np.where(X[i, :80] == np.sort(X[i, :80])[-1])[0][0]]
                        for i in range(X.shape[0])
                    ]
                )
                electron_mode_2 = np.array(
                    [
                        self.electron_list[np.where(X[i, :80] == np.sort(X[i, :80])[-2])[0][0]]
                        for i in range(X.shape[0])
                    ]
                )
                return np.c_[X, electron_property_max, electron_property_min, electron_property_mean,
                             electron_property_std, electron_mode_1, electron_mode_2]
            else:
                electron_property_max = np.array(
                    [
                        self.unfilled_electron_list[
                            np.where(X[i, :80] > 0)[0]
                        ].max() for i in range(X.shape[0])
                    ]
                )
                electron_property_min = np.array(
                    [
                        self.unfilled_electron_list[
                            np.where(X[i, :80] > 0)[0]
                        ].min() for i in range(X.shape[0])
                    ]
                )
                electron_property_mean = np.matmul(X[:, :80], self.unfilled_electron_list)
                electron_property_std = np.array(
                    [
                        np.matmul(X[i, :80], abs(self.unfilled_electron_list - electron_property_mean[i]))
                        for i in range(X.shape[0])
                    ]
                )
                electron_mode_1 = np.array(
                    [
                        self.unfilled_electron_list[np.where(X[i, :80] == np.sort(X[i, :80])[-1])[0][0]]
                        for i in range(X.shape[0])
                    ]
                )
                electron_mode_2 = np.array(
                    [
                        self.unfilled_electron_list[np.where(X[i, :80] == np.sort(X[i, :80])[-2])[0][0]]
                        for i in range(X.shape[0])
                    ]
                )
                return np.c_[X, electron_property_max, electron_property_min, electron_property_mean,
                             electron_property_std, electron_mode_1, electron_mode_2]
        else:
            return X


'''
add Lp norm features using np.linarg.norm(). p_list = [0, 2, 3, 5, 7, 10]
'''


class CombinedAttributesAdderLpNorm(BaseEstimator, TransformerMixin):

    def __init__(self, X_list, p_list, add_Lp_Norm=True):
        self.X_list = np.array(X_list).astype('object') # in X_list, the number of H is '1'
        self.p_list = p_list
        self.add_Lp_Norm = add_Lp_Norm

    def fit(self, X, y=None):
        return self

    def transform(self, X, y=None):
        if self.add_Lp_Norm:
            Lp_Norm = np.zeros([X.shape[0], len(self.p_list)])
            for num, i in enumerate(self.p_list):
                Lp_Norm[:, num] = norm(X[:, :80], i, axis=1)
            return np.c_[X, Lp_Norm]
        else:
            return X


class CombinedAttributesAdderOrbitalOccupation(BaseEstimator, TransformerMixin):

    def __init__(self, X_list, mendeleev_df, add_orbital_occupation=True):
        self.X_list = np.array(X_list).astype('object') # in X_list, the number of H is '1'
        self.add_orbital_occupation = add_orbital_occupation
        self.mendeleev_df = mendeleev_df
        self.property_ids = ['s', 'p', 'd', 'f']
        self.e_conf = mendeleev_df['electronic_configuration']
        self.electron_list_s_full = []
        self.electron_list_p_full = []
        self.electron_list_d_full = []
        self.electron_list_f_full = []
        for list, id in zip(
                (
                    self.electron_list_s_full,
                    self.electron_list_p_full,
                    self.electron_list_d_full,
                    self.electron_list_f_full
                ), self.property_ids
        ):
            for element in self.e_conf:
                if re.findall(r'{}'.format(id), element):
                    if re.findall(r'{}\d+'.format(id), element):
                        list.append(
                            int(re.findall(r'{}\d+'.format(id), element)[0][1:])
                        )
                    else:
                        list.append(1)
                else:
                    list.append(0)
        self.electron_list_s = np.array(self.electron_list_s_full)[self.X_list.astype('int')-1]
        self.electron_list_p = np.array(self.electron_list_p_full)[self.X_list.astype('int')-1]
        self.electron_list_d = np.array(self.electron_list_d_full)[self.X_list.astype('int')-1]
        self.electron_list_f = np.array(self.electron_list_f_full)[self.X_list.astype('int')-1]
        self.electron_list_total = self.electron_list_s + self.electron_list_p + \
                                   self.electron_list_d + self.electron_list_f

    def fit(self, X, y=None):
        return self

    def transform(self, X, y=None):
        if self.add_orbital_occupation:
            s = np.matmul(X[:, :80], self.electron_list_s)/\
                np.matmul(X[:, :80], self.electron_list_total)
            p = np.matmul(X[:, :80], self.electron_list_p)/\
                np.matmul(X[:, :80], self.electron_list_total)
            d = np.matmul(X[:, :80], self.electron_list_d)/\
                np.matmul(X[:, :80], self.electron_list_total)
            f = np.matmul(X[:, :80], self.electron_list_f)/\
                np.matmul(X[:, :80], self.electron_list_total)
            return np.c_[X, s, p, d, f]
        else:
            return X


class NumAttributesDropper(BaseEstimator, TransformerMixin):

    def __init__(self, X_list, drop_element_conc=True):
        self.X_list = np.array(X_list).astype('object') # in X_list, the number of H is '1'
        self.drop_element_conc = drop_element_conc

    def fit(self, X, y=None):
        return self

    def transform(self, X, y=None):
        if self.drop_element_conc:
            return X[:, 80:]
        else:
            return X


class MyLabelBinarizer(BaseEstimator, TransformerMixin):
    def __init__(self,):
        self.encoder = LabelBinarizer()

    def fit(self, X, y=0):
        return self.encoder.fit(X)

    def transform(self, X, y=0):
        return self.encoder.transform(X)


def main_element_base(DEG):
    raw_df = data_reader(path)
    name_list = list(raw_df.columns[28:108])
    drop_list = raw_df[raw_df.duplicated(name_list+['RO{}'.format(DEG)])].index.to_list()
    raw_df = raw_df.drop(drop_list, axis=0)
    name_list_ahead =list(raw_df.columns[0:28])
    cleaned_df = raw_df.groupby(name_list).mean().reset_index()
    cleaned_df = cleaned_df[name_list_ahead+name_list]
    data = stratified_data(cleaned_df)
    strat_train_data, strat_test_data = stratify_split(data)
    X_train_raw = strat_train_data.drop(['RO20', 'RO100', 'RO150', 'RO300'], axis=1)
    X_test_raw = strat_test_data.drop(['RO20', 'RO100', 'RO150', 'RO300'], axis=1)
    y_train = strat_train_data['RO{}'.format(DEG)].copy()
    y_test =strat_test_data['RO{}'.format(DEG)].copy()
    num_attribs = list(X_train_raw)[:-1]    # return a list of columns, str type. 80 elements
    print(len(num_attribs))
    cat_attribs = 'YEAR_cat'
    num_pipeline = Pipeline([
        ('selector', DataframeSelector(num_attribs)),
        # ('attribs_adder', CombinedAttributesAdderNBO(num_attribs)),
    ])
    cat_pipeline = Pipeline([
        ('selector', DataframeSelector(cat_attribs)),
        ('label_binarizer', MyLabelBinarizer()),
    ])
    full_pipeline = FeatureUnion([
        ('num_pipeline', num_pipeline),
        # ('cat_pipline', cat_pipeline),
    ])
    X_train = full_pipeline.fit_transform(X_train_raw)
    X_test = full_pipeline.fit_transform(X_test_raw)
    save_data(X_test,    'X_test_ele_based_clean_{}.pkl'.format(DEG)   )
    save_data(X_train,  'X_train_ele_based_clean_{}.pkl'.format(DEG))
    save_data(y_test,    'y_test_ele_based_clean_{}.pkl'.format(DEG)  )
    save_data(y_train,  'y_train_ele_based_clean_{}.pkl'.format(DEG) )

    df = get_table('elements')
    symbols_list = df['symbol'][[int(i)-1 for i in num_attribs]].to_list()

    return symbols_list


## raw=True for nomral data, and (raw=True AND test =True) for testing data
def main_property_base(DEG, path_name, raw=True, test=False):

    raw_df = data_reader(path_name)

    if raw:
        name_list = list(raw_df.columns[28:108])
        drop_list = raw_df[raw_df.duplicated(name_list + ['RO{}'.format(DEG)])].index.to_list()
        raw_df = raw_df.drop(drop_list, axis=0)
        name_list_ahead = list(raw_df.columns[0:28])
        cleaned_df = raw_df.groupby(name_list).mean().reset_index()
        cleaned_df = cleaned_df[name_list_ahead + name_list]
        if test:
            data = raw_df.iloc[:, [i for i in range(28, 108)] + [24, 20, 21, 22, 23]]
            X_test_raw = data.drop(['RO20', 'RO100', 'RO150', 'RO300'], axis=1)
            y_test = data['RO{}'.format(DEG)].copy()
            num_attribs= list(X_test_raw)[:-1]    # return a list of columns, str type. 80 elements
            print(len(num_attribs))
            print(num_attribs)
        else:
            data = stratified_data(cleaned_df)
            strat_train_data, strat_test_data = stratify_split(data)
            X_train_oxides = strat_train_data.iloc[:, 2:20]
            X_test_oxides = strat_test_data.iloc[:, 2:20]

            # X_train_raw = strat_train_data.drop(['RO20', 'RO100', 'RO150', 'RO300'], axis=1)
            # X_test_raw = strat_test_data.drop(['RO20', 'RO100', 'RO150', 'RO300'], axis=1)

            X_train_raw = strat_train_data.iloc[
                          :, [i for i in range(28, 108)] + [24, 20, 21, 22, 23]
                          ].drop(['RO20', 'RO100', 'RO150', 'RO300'], axis=1)
            print(X_train_raw.columns)
            X_test_raw = strat_test_data.iloc[
                         :, [i for i in range(28, 108)] + [24, 20, 21, 22, 23]
                         ].drop(['RO20', 'RO100', 'RO150', 'RO300'], axis=1)

            y_train = strat_train_data['RO{}'.format(DEG)].copy()
            y_test = strat_test_data['RO{}'.format(DEG)].copy()
            num_attribs = list(X_train_raw)[:-1]    # return a list of columns, str type. 80 elements
            print(len(num_attribs))
    else:
        data = raw_df.iloc[:, [i for i in range(28, 108)] + [24, 20, 21, 22, 23]]
        X_train_raw = data.drop(['RO20', 'RO100', 'RO150', 'RO300'], axis=1)
        X_test_raw = data.drop(['RO20', 'RO100', 'RO150', 'RO300'], axis=1)
        y_train = data['RO{}'.format(DEG)].copy()
        y_test = data['RO{}'.format(DEG)].copy()
        num_attribs = list(X_train_raw)[:-1]

    cat_attribs = 'YEAR_cat'
    num_pipeline = Pipeline([
        ('selector', DataframeSelector(num_attribs)),  # 80
        ('Lp_norm', CombinedAttributesAdderLpNorm(
            num_attribs, [0, 2, 3, 5, 7, 10],                            add_Lp_Norm=False )),    # 6 attributes
        ('attribs_adder', CombinedAttributesAdderNBO(num_attribs,   add_total_alkali=False)),  # 1 attributes
        ('atomic_num_adder', CombinedAttributesAdderAtomicProperty(  # 6 attributes
            num_attribs, 'atomic_number', element_df,            add_atomic_property=False)),
        ('atomic_radius', CombinedAttributesAdderAtomicProperty(
            num_attribs, 'atomic_radius', element_df,            add_atomic_property=False )),       # 6 attributes
        ('mendeleev_number', CombinedAttributesAdderAtomicProperty(
            num_attribs, 'mendeleev_number', element_df,         add_atomic_property=True)),   # 6 attributes
        ('atomic_weight', CombinedAttributesAdderAtomicProperty(
            num_attribs, 'atomic_weight', element_df,            add_atomic_property=False)),      # 6 attributes
        ('melting_point', CombinedAttributesAdderAtomicProperty(
            num_attribs, 'melting_point', element_df,            add_atomic_property=False)),      # 6 attributes
        ('period', CombinedAttributesAdderAtomicProperty(
            num_attribs, 'period', element_df,                   add_atomic_property=False)),      # 6 attributes
        ('series_id', CombinedAttributesAdderAtomicProperty(
            num_attribs, 'series_id', element_df,                add_atomic_property=False)),    # 6 attributes
        ('covalent_radius_cordero', CombinedAttributesAdderAtomicProperty(
            num_attribs, 'covalent_radius_cordero', element_df,  add_atomic_property=False)),   # 6 attributes
        ('en_allen', CombinedAttributesAdderAtomicProperty(
            num_attribs, 'en_allen', element_df,                 add_atomic_property=False)),    # 6 attributes
        ('s_valence', CombinedAttributesAdderElectronProperty(
            num_attribs, 's', element_df,                      add_electron_property=True)), # 6 attributes
        ('p_valence', CombinedAttributesAdderElectronProperty(
            num_attribs, 'p', element_df,                      add_electron_property=False)), # 6 attributes
        ('d_valence', CombinedAttributesAdderElectronProperty(
            num_attribs, 'd', element_df,                      add_electron_property=False)), # 6 attributes
        ('f_valence', CombinedAttributesAdderElectronProperty(
            num_attribs, 'f', element_df,                      add_electron_property=False)), # 6 attributes
        ('s_unfilled', CombinedAttributesAdderElectronProperty(
            num_attribs, 's', element_df, filled=False,        add_electron_property=True )), # 6 attributes
        ('p_unfilled', CombinedAttributesAdderElectronProperty(
            num_attribs, 'p', element_df, filled=False,        add_electron_property=False)), # 6 attributes
        ('d_unfilled', CombinedAttributesAdderElectronProperty(
            num_attribs, 'd', element_df, filled=False,        add_electron_property=False)), # 6 attributes
        ('f_unfilled', CombinedAttributesAdderElectronProperty(
            num_attribs, 'f', element_df, filled=False,        add_electron_property=False)), # 6 attributes
        ('orbital_occupation', CombinedAttributesAdderOrbitalOccupation(
            num_attribs, element_df,                          add_orbital_occupation=True )),    # 4 attributes
        ('num_dropper', NumAttributesDropper(num_attribs,          drop_element_conc=False)),
        # ('std_scaler', StandardScaler()),
    ])

    cat_pipeline = Pipeline([
        ('selector', DataframeSelector(cat_attribs)),
        ('label_binarizer', MyLabelBinarizer()),   # 5 attributes
    ])
    full_pipeline = FeatureUnion([
        ('num_pipeline', num_pipeline),  # 80 + 113
        # ('cat_pipline', cat_pipeline),  # 5
    ])
    if test:
        X_test = full_pipeline.fit_transform(X_test_raw)
        save_data(X_test, 'X_test_{}.pkl'.format(DEG))
    else:
        X_train = full_pipeline.fit_transform(X_train_raw)
        X_test = full_pipeline.fit_transform(X_test_raw)
        X_train_final = np.c_[X_train_oxides.values, X_train[:, 80:]]
        X_test_final = np.c_[X_test_oxides.values, X_test[:, 80:]]
        if raw:
            save_data(X_test,    'X_test_oxides_property_{}.pkl'.format(DEG)   )
            save_data(X_train,  'X_train_oxides_property_{}.pkl'.format(DEG))
            save_data(y_test,    'y_test_oxides_property_{}.pkl'.format(DEG)  )
            save_data(y_train,  'y_train_oxides_property_{}.pkl'.format(DEG) )
        else:
            save_data(X_test,    'X_test_product.pkl'.format(DEG)   )
            save_data(X_train,  'X_train_product.pkl'.format(DEG))
            save_data(y_test,    'y_test_product.pkl'.format(DEG)  )
            save_data(y_train,  'y_train_product.pkl'.format(DEG) )

    df = get_table('elements')
    symbols_list = df['symbol'][[int(i)-1 for i in num_attribs]].to_list()

    return symbols_list


def oxide_property_data_prep(target_DEG, raw_data_path):

    DEGs = np.array([100, 150, 300])

    path_300 = os.path.join(raw_data_path, 'worksheet_{}.xlsx'.format(300))
    path_150 = os.path.join(raw_data_path, 'worksheet_{}.xlsx'.format(150))
    path_100 = os.path.join(raw_data_path, 'worksheet_{}.xlsx'.format(100))

    clean_df_list = []
    X_train_list = []
    X_test_list = []
    y_train_list = []
    y_test_list = []

    raw_df_list = [data_reader(path_100), data_reader(path_150), data_reader(path_300)]

    num_attribs = list(raw_df_list[0].iloc[:, 28:108].columns)
    cat_attribs = 'YEAR_cat'
    num_pipeline = Pipeline([
        ('selector', DataframeSelector(num_attribs)),  # 80
        ('Lp_norm', CombinedAttributesAdderLpNorm(
            num_attribs, [0, 2, 3, 5, 7, 10],                            add_Lp_Norm=False )),    # 6 attributes
        ('attribs_adder', CombinedAttributesAdderNBO(num_attribs,   add_total_alkali=False)),  # 1 attributes
        ('atomic_num_adder', CombinedAttributesAdderAtomicProperty(  # 6 attributes
            num_attribs, 'atomic_number', element_df,            add_atomic_property=False)),
        ('atomic_radius', CombinedAttributesAdderAtomicProperty(
            num_attribs, 'atomic_radius', element_df,            add_atomic_property=False )),       # 6 attributes
        ('mendeleev_number', CombinedAttributesAdderAtomicProperty(
            num_attribs, 'mendeleev_number', element_df,         add_atomic_property=True)),   # 6 attributes
        ('atomic_weight', CombinedAttributesAdderAtomicProperty(
            num_attribs, 'atomic_weight', element_df,            add_atomic_property=False)),      # 6 attributes
        ('melting_point', CombinedAttributesAdderAtomicProperty(
            num_attribs, 'melting_point', element_df,            add_atomic_property=False)),      # 6 attributes
        ('period', CombinedAttributesAdderAtomicProperty(
            num_attribs, 'period', element_df,                   add_atomic_property=False)),      # 6 attributes
        ('series_id', CombinedAttributesAdderAtomicProperty(
            num_attribs, 'series_id', element_df,                add_atomic_property=False)),    # 6 attributes
        ('covalent_radius_cordero', CombinedAttributesAdderAtomicProperty(
            num_attribs, 'covalent_radius_cordero', element_df,  add_atomic_property=False)),   # 6 attributes
        ('en_allen', CombinedAttributesAdderAtomicProperty(
            num_attribs, 'en_allen', element_df,                 add_atomic_property=False)),    # 6 attributes
        ('s_valence', CombinedAttributesAdderElectronProperty(
            num_attribs, 's', element_df,                      add_electron_property=True)), # 6 attributes
        ('p_valence', CombinedAttributesAdderElectronProperty(
            num_attribs, 'p', element_df,                      add_electron_property=False)), # 6 attributes
        ('d_valence', CombinedAttributesAdderElectronProperty(
            num_attribs, 'd', element_df,                      add_electron_property=False)), # 6 attributes
        ('f_valence', CombinedAttributesAdderElectronProperty(
            num_attribs, 'f', element_df,                      add_electron_property=False)), # 6 attributes
        ('s_unfilled', CombinedAttributesAdderElectronProperty(
            num_attribs, 's', element_df, filled=False,        add_electron_property=True )), # 6 attributes
        ('p_unfilled', CombinedAttributesAdderElectronProperty(
            num_attribs, 'p', element_df, filled=False,        add_electron_property=False)), # 6 attributes
        ('d_unfilled', CombinedAttributesAdderElectronProperty(
            num_attribs, 'd', element_df, filled=False,        add_electron_property=False)), # 6 attributes
        ('f_unfilled', CombinedAttributesAdderElectronProperty(
            num_attribs, 'f', element_df, filled=False,        add_electron_property=False)), # 6 attributes
        ('orbital_occupation', CombinedAttributesAdderOrbitalOccupation(
            num_attribs, element_df,                          add_orbital_occupation=True )),    # 4 attributes
        ('num_dropper', NumAttributesDropper(num_attribs,          drop_element_conc=False)),
        # ('std_scaler', StandardScaler()),
    ])

    cat_pipeline = Pipeline([
        ('selector', DataframeSelector(cat_attribs)),
        ('label_binarizer', MyLabelBinarizer()),   # 5 attributes
    ])
    full_pipeline = FeatureUnion([
        ('num_pipeline', num_pipeline),  # 80 + 113
        # ('cat_pipline', cat_pipeline),  # 5
    ])

    for i, raw_df in enumerate(raw_df_list):
        element_name_list = list(raw_df.columns[28:108])
        drop_list = raw_df[raw_df.duplicated(element_name_list + ['RO{}'.format(DEGs[i])])].index.to_list()
        raw_df = raw_df.drop(drop_list, axis=0)
        oxides_list_ahead = list(raw_df.columns[0:28])
        cleaned_df = raw_df.groupby(element_name_list).mean().reset_index()
        cleaned_df = cleaned_df[oxides_list_ahead + element_name_list]
        clean_df_list.append(cleaned_df)

        data = stratified_data(cleaned_df)
        strat_train_data, strat_test_data = stratify_split(data)
        X_train_oxides = strat_train_data.iloc[:, 2:20]
        X_test_oxides = strat_test_data.iloc[:, 2:20]


        X_train_raw = strat_train_data.iloc[
                      :, [i for i in range(28, 108)] + [24, 20, 21, 22, 23]
                      ].drop(['RO20', 'RO100', 'RO150', 'RO300'], axis=1)
        print(X_train_raw.columns)
        X_test_raw = strat_test_data.iloc[
                     :, [i for i in range(28, 108)] + [24, 20, 21, 22, 23]
                     ].drop(['RO20', 'RO100', 'RO150', 'RO300'], axis=1)

        y_train = strat_train_data['RO{}'.format(DEGs[i])].copy()
        y_test = strat_test_data['RO{}'.format(DEGs[i])].copy()

        X_train = full_pipeline.fit_transform(X_train_raw)
        X_test = full_pipeline.fit_transform(X_test_raw)
        X_train_final = np.c_[X_train_oxides.values, X_train[:, 80:]]
        X_test_final = np.c_[X_test_oxides.values, X_test[:, 80:]]

        X_train_list.append(X_train_final)
        X_test_list.append(X_test_final)
        y_train_list.append(y_train)
        y_test_list.append(y_test)


    ## avoid data leakage

    X_test_target_deg = X_test_list[np.where(DEGs==target_DEG)[0][0]]

    for num in np.where(DEGs!=target_DEG)[0]:
        mask = []
        for row in X_train_list[num]:
            mask.append(any(np.equal(X_test_target_deg, row).all(1)))

        (bool_val, counts) = np.unique(mask, return_counts=True)

        print('{} leakage points found for dataset at {} DEG'.format(counts[1], DEGs[num]))
        X_train_list[num] = X_train_list[num][~np.array(mask)]
        y_train_list[num] = y_train_list[num][~np.array(mask)]

    ## output the datasets

    for num, DEG in enumerate(DEGs):
        save_data(
            X_train_list[num],
            './datasets/target_{}/X_train_oxides_property_{}.pkl'.format(target_DEG, DEG)
        )
        save_data(
            X_test_list[num],
            './datasets/target_{}/X_test_oxides_property_{}.pkl'.format(target_DEG, DEG)
        )

        save_data(
            y_train_list[num],
            './datasets/target_{}/y_train_oxides_property_{}.pkl'.format(target_DEG, DEG)
        )
        save_data(
            y_test_list[num],
            './datasets/target_{}/y_test_oxides_property_{}.pkl'.format(target_DEG, DEG)
        )


# function to create raw data composition file
def test_raw_data_preparation(component_list, prime_component_molar_fraction):
    ## e.g. [SIO2, NA2O, K2O], molar fraction in %
    column_list = np.array(['GLASNO', 'A_WT', 'SIO2', 'AL2O3', 'B2O3', 'CAO', 'K2O', 'NA2O', 'PBO', 'Li2O', 'MgO',
                   'SRO', 'BAO', 'ZNO', 'P2O5', 'GEO2', 'ZRO2', 'TIO2', 'TEO2', 'FemOn', 'RO20', 'RO100',
                   'RO150', 'RO300', 'Year', 'Refer_Id', 'Patent', 'SLeg'])
    if len(component_list)==3:
        temp_data = np.zeros([200, len(column_list)])
        temp_data[:, np.where(column_list==component_list[0])[0][0]] = prime_component_molar_fraction
        temp_data[:, np.where(column_list==component_list[1])[0][0]]=np.linspace(0,
                                                                           100-prime_component_molar_fraction,
                                                                           temp_data.shape[0])
        temp_data[:, np.where(column_list==component_list[2])] = 100 - \
                                                                 temp_data[:, np.where(column_list==component_list[1])]-\
                                                                 prime_component_molar_fraction
        test_raw_data_df = pd.DataFrame(temp_data, columns=column_list)
        test_raw_data_df.to_csv(r'./data/testing data/mixed alkali effect/raw_data_{0}_{1}.csv'
                                .format(component_list, prime_component_molar_fraction))
    else:
        print('error! please check the inputs!')
    return test_raw_data_df


def product_data_preparation(path, DEG):
    raw_df = pd.read_excel(path, engine='openpyxl')

    num_attribs = list(raw_df.iloc[:, 28:108].columns)
    cat_attribs = 'YEAR_cat'
    num_pipeline = Pipeline([
        ('selector', DataframeSelector(num_attribs)),  # 80
        ('Lp_norm', CombinedAttributesAdderLpNorm(
            num_attribs, [0, 2, 3, 5, 7, 10],                            add_Lp_Norm=False )),    # 6 attributes
        ('attribs_adder', CombinedAttributesAdderNBO(num_attribs,   add_total_alkali=False)),  # 1 attributes
        ('atomic_num_adder', CombinedAttributesAdderAtomicProperty(  # 6 attributes
            num_attribs, 'atomic_number', element_df,            add_atomic_property=False)),
        ('atomic_radius', CombinedAttributesAdderAtomicProperty(
            num_attribs, 'atomic_radius', element_df,            add_atomic_property=False )),       # 6 attributes
        ('mendeleev_number', CombinedAttributesAdderAtomicProperty(
            num_attribs, 'mendeleev_number', element_df,         add_atomic_property=True)),   # 6 attributes
        ('atomic_weight', CombinedAttributesAdderAtomicProperty(
            num_attribs, 'atomic_weight', element_df,            add_atomic_property=False)),      # 6 attributes
        ('melting_point', CombinedAttributesAdderAtomicProperty(
            num_attribs, 'melting_point', element_df,            add_atomic_property=False)),      # 6 attributes
        ('period', CombinedAttributesAdderAtomicProperty(
            num_attribs, 'period', element_df,                   add_atomic_property=False)),      # 6 attributes
        ('series_id', CombinedAttributesAdderAtomicProperty(
            num_attribs, 'series_id', element_df,                add_atomic_property=False)),    # 6 attributes
        ('covalent_radius_cordero', CombinedAttributesAdderAtomicProperty(
            num_attribs, 'covalent_radius_cordero', element_df,  add_atomic_property=False)),   # 6 attributes
        ('en_allen', CombinedAttributesAdderAtomicProperty(
            num_attribs, 'en_allen', element_df,                 add_atomic_property=False)),    # 6 attributes
        ('s_valence', CombinedAttributesAdderElectronProperty(
            num_attribs, 's', element_df,                      add_electron_property=True)), # 6 attributes
        ('p_valence', CombinedAttributesAdderElectronProperty(
            num_attribs, 'p', element_df,                      add_electron_property=False)), # 6 attributes
        ('d_valence', CombinedAttributesAdderElectronProperty(
            num_attribs, 'd', element_df,                      add_electron_property=False)), # 6 attributes
        ('f_valence', CombinedAttributesAdderElectronProperty(
            num_attribs, 'f', element_df,                      add_electron_property=False)), # 6 attributes
        ('s_unfilled', CombinedAttributesAdderElectronProperty(
            num_attribs, 's', element_df, filled=False,        add_electron_property=True )), # 6 attributes
        ('p_unfilled', CombinedAttributesAdderElectronProperty(
            num_attribs, 'p', element_df, filled=False,        add_electron_property=False)), # 6 attributes
        ('d_unfilled', CombinedAttributesAdderElectronProperty(
            num_attribs, 'd', element_df, filled=False,        add_electron_property=False)), # 6 attributes
        ('f_unfilled', CombinedAttributesAdderElectronProperty(
            num_attribs, 'f', element_df, filled=False,        add_electron_property=False)), # 6 attributes
        ('orbital_occupation', CombinedAttributesAdderOrbitalOccupation(
            num_attribs, element_df,                          add_orbital_occupation=True )),    # 4 attributes
        ('num_dropper', NumAttributesDropper(num_attribs,          drop_element_conc=False)),
        # ('std_scaler', StandardScaler()),
    ])

    full_pipeline = FeatureUnion([
        ('num_pipeline', num_pipeline),  # 80 + 113
    ])

    X_product_elements = raw_df.iloc[:, [i for i in range(28, 108)] + [24, 20, 21, 22, 23]
                      ].drop(['RO20', 'RO100', 'RO150', 'RO300'], axis=1)
    X_product_oxides = raw_df.iloc[:, 2:20]
    y_product = raw_df['RO{}'.format(DEG)].copy()
    X_product = np.c_[
        X_product_oxides,
        full_pipeline.fit_transform(X_product_elements)[:, 80:]
    ]

    save_data(X_product, './datasets/X_product.pkl')
    save_data(y_product, './datasets/y_product.pkl')

    return X_product, y_product

property_list = ['Lp_0', 'Lp_2', 'Lp_3', 'Lp_5', 'Lp_7', 'Lp_10',
                 'NBO',
                 'a_num_max', 'a_num_min', 'a_num_mean', 'a_num_std', 'a_num_1', 'a_num_2',
                 'a_radii_max', 'a_radii_min', 'a_radii_mean', 'a_radii_std', 'a_radii_1', 'a_radii_2',
                 'me_num_max', 'me_num_min', 'me_num_mean', 'me_num_std', 'me_num_1', 'me_num_2',
                 'a_wt_max', 'a_wt_min', 'a_wt_mean', 'a_wt_std', 'a_wt_1', 'a_wt_2',
                 'a_mel_max', 'a_mel_min', 'a_mel_mean', 'a_mel_std', 'a_mel_1', 'a_mel_2',
                 'a_per_max', 'a_per_min', 'a_per_mean', 'a_per_std', 'a_per_1', 'a_per_2',
                 'a_sid_max', 'a_sid_min', 'a_sid_mean', 'a_sid_std', 'a_sid_1', 'a_sid_2',
                 'a_cra_max', 'a_cra_min', 'a_cra_mean', 'a_cra_std', 'a_cra_1', 'a_cra_2',
                 'a_ale_max', 'a_ale_min', 'a_ale_mean', 'a_ale_std', 'a_ale_1', 'a_ale_2',
                 'e_sf_max', 'e_sf_min', 'e_sf_mean', 'e_sf_std', 'e_sf_1', 'e_sf_2',
                 'e_pf_max', 'e_pf_min', 'e_pf_mean', 'e_pf_std', 'e_pf_1', 'e_pf_2',
                 'e_df_max', 'e_df_min', 'e_df_mean', 'e_df_std', 'e_df_1', 'e_df_2',
                 'e_ff_max', 'e_ff_min', 'e_ff_mean', 'e_ff_std', 'e_ff_1', 'e_ff_2',
                 'e_suf_max', 'e_suf_min', 'e_suf_mean', 'e_suf_std', 'e_suf_1', 'e_suf_2',
                 'e_puf_max', 'e_puf_min', 'e_puf_mean', 'e_puf_std', 'e_puf_1', 'e_puf_2',
                 'e_duf_max', 'e_duf_min', 'e_duf_mean', 'e_duf_std', 'e_duf_1', 'e_duf_2',
                 'e_fuf_max', 'e_fuf_min', 'e_fuf_mean', 'e_fuf_std', 'e_fuf_1', 'e_fuf_2',
                 'oc_s', 'oc_p', 'oc_d', 'oc_f'
                 ]


'''
the order of atomic property we used is:
'atomic_number', 'atomic_radius', 'mendeleev_number', 'atomic_weight', 'melting_point', 'period',
'series_id', 'covalent_radius_cordero', 'en_allen', 
'''

if __name__ == '__main__':

    # na =main_element_base(300)
    Temperature = 300
    # path_product = r'./data/worksheet_product_{}.xlsx'.format(Temperature)
    path = r'./data/worksheet_product.xlsx'

    # test_raw_data_preparation(['SIO2', 'K2O', 'Li2O'], 67)
    X, y = product_data_preparation(path, Temperature)
# from mendeleev.fetch import fetch_table # To initialize the periodic elemental table for the further calculation # data = raw_df.iloc[:, [i for i in range(28, 108)] + [24, 20, 21, 22, 23]] # for set in (strat_train_data, strat_test_data): #     set.drop(['Na_cat'], axis=1, inplace=True) #     set.drop(['Year'], axis=1, inplace=True) # transform the data from pd.df to array # ADD NBO ratio # (array([i], dtype=int64))[0][0] # the order of features added is max, min, mean, std, mode1, mode2 # x[i, :80] select all elements except other features the order of atomic property we used is:
 'atomic_number', 'atomic_radius', 'mendeleev_number', 'atomic_weight', 'melting_point', 'period',
 'series_id', 'covalent_radius_cordero', 'en_allen', # the order of features added is max, min, mean, std, mode1, mode2 # in X_list, the number of H is '1' # the row start with 0, e.g. row number of H is zero # x[i, :80] select all elements except other features CombinedAttributesAdderElectronProperty
 property id for electron configuration is 's', 'd', 'p', 'f', if filled=True, return valence electrons,
 if filled=False, return unfilled states # in X_list, the number of H is '1' # x[i, :80] select all elements except other features add Lp norm features using np.linarg.norm(). p_list = [0, 2, 3, 5, 7, 10] # in X_list, the number of H is '1' # in X_list, the number of H is '1' # in X_list, the number of H is '1' # return a list of columns, str type. 80 elements # ('attribs_adder', CombinedAttributesAdderNBO(num_attribs)), # ('cat_pipline', cat_pipeline), ## raw=True for nomral data, and (raw=True AND test =True) for testing data # return a list of columns, str type. 80 elements # X_train_raw = strat_train_data.drop(['RO20', 'RO100', 'RO150', 'RO300'], axis=1) # X_test_raw = strat_test_data.drop(['RO20', 'RO100', 'RO150', 'RO300'], axis=1) # return a list of columns, str type. 80 elements # 80 # 6 attributes # 1 attributes # 6 attributes # 6 attributes # 6 attributes # 6 attributes # 6 attributes # 6 attributes # 6 attributes # 6 attributes # 6 attributes # 6 attributes # 6 attributes # 6 attributes # 6 attributes # 6 attributes # 6 attributes # 6 attributes # 6 attributes # 4 attributes # ('std_scaler', StandardScaler()), # 5 attributes # 80 + 113 # ('cat_pipline', cat_pipeline),  # 5 # 80 # 6 attributes # 1 attributes # 6 attributes # 6 attributes # 6 attributes # 6 attributes # 6 attributes # 6 attributes # 6 attributes # 6 attributes # 6 attributes # 6 attributes # 6 attributes # 6 attributes # 6 attributes # 6 attributes # 6 attributes # 6 attributes # 6 attributes # 4 attributes # ('std_scaler', StandardScaler()), # 5 attributes # 80 + 113 # ('cat_pipline', cat_pipeline),  # 5 ## avoid data leakage ## output the datasets # function to create raw data composition file ## e.g. [SIO2, NA2O, K2O], molar fraction in % # 80 # 6 attributes # 1 attributes # 6 attributes # 6 attributes # 6 attributes # 6 attributes # 6 attributes # 6 attributes # 6 attributes # 6 attributes # 6 attributes # 6 attributes # 6 attributes # 6 attributes # 6 attributes # 6 attributes # 6 attributes # 6 attributes # 6 attributes # 4 attributes # ('std_scaler', StandardScaler()), # 80 + 113 the order of atomic property we used is:
 'atomic_number', 'atomic_radius', 'mendeleev_number', 'atomic_weight', 'melting_point', 'period',
 'series_id', 'covalent_radius_cordero', 'en_allen', # na =main_element_base(300) # path_product = r'./data/worksheet_product_{}.xlsx'.format(Temperature) # test_raw_data_preparation(['SIO2', 'K2O', 'Li2O'], 67)