# -*- coding: utf-8 -*-
"""
Created on Thu Jan  9 09:22:42 2020

@author: hcji
"""


import numpy as np
import tensorflow.keras.backend as K
import matplotlib.pyplot as plt
from sklearn.model_selection import train_test_split
from tensorflow.keras.models import Model
from tensorflow.keras.layers import Dense, Input, Flatten, Conv1D, MaxPooling1D, concatenate
from tensorflow.keras import optimizers
from sklearn.metrics import mean_absolute_error, r2_score
from smiles_to_onehot.encoding import get_dict, one_hot_coding

class multi_CNN:
    def __init__(self, X, Y):
        self.X = X
        self.Y = Y
        self.X_tr, self.X_ts, self.Y_tr, self.Y_ts = train_test_split(X, Y, test_size=0.1)
        
        inp = Input(shape=(X.shape[1:3]))
        n = X.shape[1]

        hidv1 = Conv1D(n, kernel_size=2, activation='relu')(inp)
        # hidv1 = MaxPooling1D(pool_size=2)(hidv1)
        hidv1 = Conv1D(n, kernel_size=2, activation='relu')(hidv1)
        # hidv1 = MaxPooling1D(pool_size=2)(hidv1)
        hidv1 = Flatten()(hidv1)
        
        hidv2 = Conv1D(n, kernel_size=3, activation='relu')(inp)
        # hidv2 = MaxPooling1D(pool_size=3)(hidv2)
        hidv2 = Conv1D(n, kernel_size=3, activation='relu')(hidv2)
        # hidv2 = MaxPooling1D(pool_size=3)(hidv2)
        hidv2 = Flatten()(hidv2)
        
        hidv3 = Conv1D(n, kernel_size=4, activation='relu')(inp)
        # hidv3 = MaxPooling1D(pool_size=4)(hidv3)
        hidv3 = Conv1D(n, kernel_size=4, activation='relu')(hidv3)
        # hidv3 = MaxPooling1D(pool_size=4)(hidv3)
        hidv3 = Flatten()(hidv3)

        hid = concatenate([hidv1, hidv2, hidv3], axis=-1)
        hid = Dense(32, activation="relu")(hid)
        hid = Dense(32, activation="relu")(hid)
        
        prd = Dense(1, activation="linear")(hid)
        opt = optimizers.Adam(lr=0.001)
        model = Model(inp, prd)
        model.compile(optimizer=opt, loss='mse', metrics=['mae'])
        self.model = model
    
    def train(self, epochs=20):
        history = self.model.fit(self.X_tr, self.Y_tr, epochs=epochs, validation_split = 0.1)
        plt.cla()
        plt.plot(history.history['val_loss'], alpha= 0.8)
        plt.plot(history.history['val_mean_absolute_error'], alpha= 0.8)
        plt.legend(['loss', 'accuracy'], loc="lower left")
        plt.xlabel('epoch')
        return history
    
    def test(self):
        Y_test = self.Y_ts
        Y_pred = np.round(self.model.predict(self.X_ts))
        r2 = round(r2_score(Y_pred, Y_test), 4)
        mae = round(mean_absolute_error(Y_pred, Y_test), 4)

        plt.cla()
        plt.plot(Y_test, Y_pred, '.', color = 'blue')
        plt.plot([0,4500], [0,4500], color ='red')
        plt.ylabel('Predicted RI')
        plt.xlabel('Experimental RI')        
        plt.text(0, 4000, 'R2='+str(r2), fontsize=12)
        plt.text(0, 3600, 'MAE='+str(mae), fontsize=12)
        plt.show()
        return r2, mae
    
    def save(self, path):
        self.model.save(path)
        K.clear_session()
        

if __name__ == '__main__':
    
    import json
    
    with open('DeepEI/data/split.json', 'r') as js:
        keep = np.array(json.load(js)['keep'])
        
    smiles = np.array(json.load(open('DeepEI/data/all_smiles.json')))[keep]
    rindex = np.load('DeepEI/data/retention.npy')[keep,:]
    
    words = get_dict(smiles, save_path='DeepEI/data/words.json')
    smiles = [one_hot_coding(smi, words, max_len=100).todense() for smi in smiles]
    smiles = np.array(smiles)
    
    # simipolar
    i = np.where(~ np.isnan(rindex[:,0]))[0]
    mod = multi_CNN(smiles[i], rindex[i,0])
    mod.train()
    mod.test()
    mod.save('Retention/models/SimiStdNP_CNN_multi_model.h5')
    
    # nonpolar
    i = np.where(~ np.isnan(rindex[:,1]))[0]
    mod = multi_CNN(smiles[i], rindex[i,1])
    mod.train()
    mod.test()
    mod.save('Retention/models/StdNP_CNN_multi_model.h5')

    # polar
    i = np.where(~ np.isnan(rindex[:,2]))[0]
    mod = multi_CNN(smiles[i], rindex[i,2])
    mod.train()
    mod.test()
    mod.save('Retention/models/StdPolar_CNN_multi_model.h5')
    
matplotlib.pyplot.legend sklearn.metrics.r2_score sklearn.metrics.mean_absolute_error matplotlib.pyplot.plot tensorflow.keras.backend.clear_session numpy.load tensorflow.keras.layers.Flatten tensorflow.keras.models.Model numpy.isnan tensorflow.keras.layers.Dense sklearn.model_selection.train_test_split numpy.array matplotlib.pyplot.show matplotlib.pyplot.ylabel matplotlib.pyplot.cla tensorflow.keras.layers.Conv1D tensorflow.keras.layers.concatenate tensorflow.keras.optimizers.Adam matplotlib.pyplot.xlabel tensorflow.keras.layers.Input
[(96, 'smiles_to_onehot.encoding.get_dict', 'get_dict', (['smiles'], {'save_path': '"""DeepEI/data/words.json"""'}), False, 'from smiles_to_onehot.encoding import get_dict, one_hot_coding\n'), (98, 'numpy.array', 'np.array', (['smiles'], {}), True, 'import numpy as np\n'), (23, 'sklearn.model_selection.train_test_split', 'train_test_split', (['X', 'Y'], {'test_size': '(0.1)'}), False, 'from sklearn.model_selection import train_test_split\n'), (25, 'tensorflow.keras.layers.Input', 'Input', ([], {'shape': 'X.shape[1:3]'}), False, 'from tensorflow.keras.layers import Dense, Input, Flatten, Conv1D, MaxPooling1D, concatenate\n'), (46, 'tensorflow.keras.layers.concatenate', 'concatenate', (['[hidv1, hidv2, hidv3]'], {'axis': '(-1)'}), False, 'from tensorflow.keras.layers import Dense, Input, Flatten, Conv1D, MaxPooling1D, concatenate\n'), (51, 'tensorflow.keras.optimizers.Adam', 'optimizers.Adam', ([], {'lr': '(0.001)'}), False, 'from tensorflow.keras import optimizers\n'), (52, 'tensorflow.keras.models.Model', 'Model', (['inp', 'prd'], {}), False, 'from tensorflow.keras.models import Model\n'), (58, 'matplotlib.pyplot.cla', 'plt.cla', ([], {}), True, 'import matplotlib.pyplot as plt\n'), (59, 'matplotlib.pyplot.plot', 'plt.plot', (["history.history['val_loss']"], {'alpha': '(0.8)'}), True, 'import matplotlib.pyplot as plt\n'), (60, 'matplotlib.pyplot.plot', 'plt.plot', (["history.history['val_mean_absolute_error']"], {'alpha': '(0.8)'}), True, 'import matplotlib.pyplot as plt\n'), (61, 'matplotlib.pyplot.legend', 'plt.legend', (["['loss', 'accuracy']"], {'loc': '"""lower left"""'}), True, 'import matplotlib.pyplot as plt\n'), (62, 'matplotlib.pyplot.xlabel', 'plt.xlabel', (['"""epoch"""'], {}), True, 'import matplotlib.pyplot as plt\n'), (71, 'matplotlib.pyplot.cla', 'plt.cla', ([], {}), True, 'import matplotlib.pyplot as plt\n'), (72, 'matplotlib.pyplot.plot', 'plt.plot', (['Y_test', 'Y_pred', '"""."""'], {'color': '"""blue"""'}), True, 'import matplotlib.pyplot as plt\n'), (73, 'matplotlib.pyplot.plot', 'plt.plot', (['[0, 4500]', '[0, 4500]'], {'color': '"""red"""'}), True, 'import matplotlib.pyplot as plt\n'), (74, 'matplotlib.pyplot.ylabel', 'plt.ylabel', (['"""Predicted RI"""'], {}), True, 'import matplotlib.pyplot as plt\n'), (75, 'matplotlib.pyplot.xlabel', 'plt.xlabel', (['"""Experimental RI"""'], {}), True, 'import matplotlib.pyplot as plt\n'), (78, 'matplotlib.pyplot.show', 'plt.show', ([], {}), True, 'import matplotlib.pyplot as plt\n'), (83, 'tensorflow.keras.backend.clear_session', 'K.clear_session', ([], {}), True, 'import tensorflow.keras.backend as K\n'), (94, 'numpy.load', 'np.load', (['"""DeepEI/data/retention.npy"""'], {}), True, 'import numpy as np\n'), (28, 'tensorflow.keras.layers.Conv1D', 'Conv1D', (['n'], {'kernel_size': '(2)', 'activation': '"""relu"""'}), False, 'from tensorflow.keras.layers import Dense, Input, Flatten, Conv1D, MaxPooling1D, concatenate\n'), (30, 'tensorflow.keras.layers.Conv1D', 'Conv1D', (['n'], {'kernel_size': '(2)', 'activation': '"""relu"""'}), False, 'from tensorflow.keras.layers import Dense, Input, Flatten, Conv1D, MaxPooling1D, concatenate\n'), (32, 'tensorflow.keras.layers.Flatten', 'Flatten', ([], {}), False, 'from tensorflow.keras.layers import Dense, Input, Flatten, Conv1D, MaxPooling1D, concatenate\n'), (34, 'tensorflow.keras.layers.Conv1D', 'Conv1D', (['n'], {'kernel_size': '(3)', 'activation': '"""relu"""'}), False, 'from tensorflow.keras.layers import Dense, Input, Flatten, Conv1D, MaxPooling1D, concatenate\n'), (36, 'tensorflow.keras.layers.Conv1D', 'Conv1D', (['n'], {'kernel_size': '(3)', 'activation': '"""relu"""'}), False, 'from tensorflow.keras.layers import Dense, Input, Flatten, Conv1D, MaxPooling1D, concatenate\n'), (38, 'tensorflow.keras.layers.Flatten', 'Flatten', ([], {}), False, 'from tensorflow.keras.layers import Dense, Input, Flatten, Conv1D, MaxPooling1D, concatenate\n'), (40, 'tensorflow.keras.layers.Conv1D', 'Conv1D', (['n'], {'kernel_size': '(4)', 'activation': '"""relu"""'}), False, 'from tensorflow.keras.layers import Dense, Input, Flatten, Conv1D, MaxPooling1D, concatenate\n'), (42, 'tensorflow.keras.layers.Conv1D', 'Conv1D', (['n'], {'kernel_size': '(4)', 'activation': '"""relu"""'}), False, 'from tensorflow.keras.layers import Dense, Input, Flatten, Conv1D, MaxPooling1D, concatenate\n'), (44, 'tensorflow.keras.layers.Flatten', 'Flatten', ([], {}), False, 'from tensorflow.keras.layers import Dense, Input, Flatten, Conv1D, MaxPooling1D, concatenate\n'), (47, 'tensorflow.keras.layers.Dense', 'Dense', (['(32)'], {'activation': '"""relu"""'}), False, 'from tensorflow.keras.layers import Dense, Input, Flatten, Conv1D, MaxPooling1D, concatenate\n'), (48, 'tensorflow.keras.layers.Dense', 'Dense', (['(32)'], {'activation': '"""relu"""'}), False, 'from tensorflow.keras.layers import Dense, Input, Flatten, Conv1D, MaxPooling1D, concatenate\n'), (50, 'tensorflow.keras.layers.Dense', 'Dense', (['(1)'], {'activation': '"""linear"""'}), False, 'from tensorflow.keras.layers import Dense, Input, Flatten, Conv1D, MaxPooling1D, concatenate\n'), (68, 'sklearn.metrics.r2_score', 'r2_score', (['Y_pred', 'Y_test'], {}), False, 'from sklearn.metrics import mean_absolute_error, r2_score\n'), (69, 'sklearn.metrics.mean_absolute_error', 'mean_absolute_error', (['Y_pred', 'Y_test'], {}), False, 'from sklearn.metrics import mean_absolute_error, r2_score\n'), (91, 'json.load', 'json.load', (['js'], {}), False, 'import json\n'), (97, 'smiles_to_onehot.encoding.one_hot_coding', 'one_hot_coding', (['smi', 'words'], {'max_len': '(100)'}), False, 'from smiles_to_onehot.encoding import get_dict, one_hot_coding\n'), (101, 'numpy.isnan', 'np.isnan', (['rindex[:, (0)]'], {}), True, 'import numpy as np\n'), (108, 'numpy.isnan', 'np.isnan', (['rindex[:, (1)]'], {}), True, 'import numpy as np\n'), (115, 'numpy.isnan', 'np.isnan', (['rindex[:, (2)]'], {}), True, 'import numpy as np\n')]