from sklearn.metrics import roc_auc_score, f1_score, average_precision_score, precision_score, recall_score, accuracy_score

import matplotlib.pyplot as plt

from copy import deepcopy 

import numpy as np 

from tqdm import tqdm 

import torch 

torch.manual_seed(0)

from torch import nn 

from torch.autograd import Variable

import torch.nn.functional as F

from HINT.module import Highway, GCN 

from functools import reduce 

import pickle

class Interaction(nn.Sequential):

    def __init__(self, molecule_encoder, disease_encoder, protocol_encoder, 

                    device, 

                    global_embed_size, 

                    highway_num_layer,

                    prefix_name, 

                    epoch = 20,

                    lr = 3e-4, 

                    weight_decay = 0, 

                    ):

        super(Interaction, self).__init__()

        self.molecule_encoder = molecule_encoder 

        self.disease_encoder = disease_encoder 

        self.protocol_encoder = protocol_encoder 

        self.global_embed_size = global_embed_size 

        self.highway_num_layer = highway_num_layer 

        self.feature_dim = self.molecule_encoder.embedding_size + self.disease_encoder.embedding_size + self.protocol_encoder.embedding_size

        self.epoch = epoch 

        self.lr = lr 

        self.weight_decay = weight_decay 

        self.save_name = prefix_name + '_interaction'

        self.f = F.relu

        self.loss = nn.BCEWithLogitsLoss()

        ##### NN 

        self.encoder2interaction_fc = nn.Linear(self.feature_dim, self.global_embed_size).to(device)

        self.encoder2interaction_highway = Highway(self.global_embed_size, self.highway_num_layer).to(device)

        self.pred_nn = nn.Linear(self.global_embed_size, 1)

        self.device = device 

        self = self.to(device)

    def feed_lst_of_module(self, input_feature, lst_of_module):

        x = input_feature

        for single_module in lst_of_module:

            x = self.f(single_module(x))

        return x

    def forward_get_three_encoders(self, smiles_lst2, icdcode_lst3, criteria_lst):

        molecule_embed = self.molecule_encoder.forward_smiles_lst_lst(smiles_lst2)

        icd_embed = self.disease_encoder.forward_code_lst3(icdcode_lst3)

        protocol_embed = self.protocol_encoder.forward(criteria_lst)

        return molecule_embed, icd_embed, protocol_embed    

    def forward_encoder_2_interaction(self, molecule_embed, icd_embed, protocol_embed):

        encoder_embedding = torch.cat([molecule_embed, icd_embed, protocol_embed], 1)

        # interaction_embedding = self.feed_lst_of_module(encoder_embedding, [self.encoder2interaction_fc, self.encoder2interaction_highway])

        h = self.encoder2interaction_fc(encoder_embedding)

        h = self.f(h)

        h = self.encoder2interaction_highway(h)

        interaction_embedding = self.f(h)

        return interaction_embedding 

    def forward(self, smiles_lst2, icdcode_lst3, criteria_lst):

        molecule_embed, icd_embed, protocol_embed = self.forward_get_three_encoders(smiles_lst2, icdcode_lst3, criteria_lst)

        interaction_embedding = self.forward_encoder_2_interaction(molecule_embed, icd_embed, protocol_embed)

        output = self.pred_nn(interaction_embedding)

        return output ### 32, 1

    def evaluation(self, predict_all, label_all, threshold = 0.5):

        import pickle, os

        from sklearn.metrics import roc_curve, precision_recall_curve

        with open("predict_label.txt", 'w') as fout:

            for i,j in zip(predict_all, label_all):

                fout.write(str(i)[:6] + '\t' + str(j)[:4]+'\n')

        auc_score = roc_auc_score(label_all, predict_all)

        figure_folder = "figure"

        #### ROC-curve 

        fpr, tpr, thresholds = roc_curve(label_all, predict_all, pos_label=1)

        # roc_curve =plt.figure()

        # plt.plot(fpr,tpr,'-',label=self.save_name + ' ROC Curve ')

        # plt.legend(fontsize = 15)

        # plt.savefig(os.path.join(figure_folder,self.save_name+"_roc_curve.png"))

        #### PR-curve

        precision, recall, thresholds = precision_recall_curve(label_all, predict_all)

        # plt.plot(recall,precision, label = self.save_name + ' PR Curve')

        # plt.legend(fontsize = 15)

        # plt.savefig(os.path.join(figure_folder,self.save_name + "_pr_curve.png"))

        label_all = [int(i) for i in label_all]

        float2binary = lambda x:0 if x < threshold else 1

        predict_all = list(map(float2binary, predict_all))

        f1score = f1_score(label_all, predict_all)

        prauc_score = average_precision_score(label_all, predict_all)

        # print(predict_all)

        precision = precision_score(label_all, predict_all)

        recall = recall_score(label_all, predict_all)

        accuracy = accuracy_score(label_all, predict_all)

        predict_1_ratio = sum(predict_all) / len(predict_all)

        label_1_ratio = sum(label_all) / len(label_all)

        return auc_score, f1score, prauc_score, precision, recall, accuracy, predict_1_ratio, label_1_ratio 

    def testloader_to_lst(self, dataloader):

        nctid_lst, label_lst, smiles_lst2, icdcode_lst3, criteria_lst = [], [], [], [], []

        for nctid, label, smiles, icdcode, criteria in dataloader:

            nctid_lst.extend(nctid)

            label_lst.extend([i.item() for i in label])

            smiles_lst2.extend(smiles)

            icdcode_lst3.extend(icdcode)

            criteria_lst.extend(criteria)

        length = len(nctid_lst)

        assert length == len(smiles_lst2) and length == len(icdcode_lst3)

        return nctid_lst, label_lst, smiles_lst2, icdcode_lst3, criteria_lst, length 

    def generate_predict(self, dataloader):

        whole_loss = 0 

        label_all, predict_all, nctid_all = [], [], []

        for nctid_lst, label_vec, smiles_lst2, icdcode_lst3, criteria_lst in dataloader:

            nctid_all.extend(nctid_lst)

            label_vec = label_vec.to(self.device)

            output = self.forward(smiles_lst2, icdcode_lst3, criteria_lst).view(-1)  

            loss = self.loss(output, label_vec.float())

            whole_loss += loss.item()

            predict_all.extend([i.item() for i in torch.sigmoid(output)])

            label_all.extend([i.item() for i in label_vec])

        return whole_loss, predict_all, label_all, nctid_all

    def bootstrap_test(self, dataloader, valid_loader = None, sample_num = 20):

        best_threshold = 0.5

        # if validloader is not None:

        #   best_threshold = self.select_threshold_for_binary(valid_loader)

        #   print(f"best_threshold: {best_threshold}")

        self.eval()

        whole_loss, predict_all, label_all, nctid_all = self.generate_predict(dataloader)

        from HINT.utils import plot_hist

        plt.clf()

        prefix_name = "./figure/" + self.save_name 

        plot_hist(prefix_name, predict_all, label_all)      

        def bootstrap(length, sample_num):

            idx = [i for i in range(length)]

            from random import choices 

            bootstrap_idx = [choices(idx, k = length) for i in range(sample_num)]

            return bootstrap_idx 

        results_lst = []

        bootstrap_idx_lst = bootstrap(len(predict_all), sample_num = sample_num)

        for bootstrap_idx in bootstrap_idx_lst: 

            bootstrap_label = [label_all[idx] for idx in bootstrap_idx]     

            bootstrap_predict = [predict_all[idx] for idx in bootstrap_idx]

            results = self.evaluation(bootstrap_predict, bootstrap_label, threshold = best_threshold)

            results_lst.append(results)

        self.train() 

        auc = [results[0] for results in results_lst]

        f1score = [results[1] for results in results_lst]

        prauc_score = [results[2] for results in results_lst]

        print("PR-AUC   mean: "+str(np.mean(prauc_score))[:6], "std: "+str(np.std(prauc_score))[:6])

        print("F1       mean: "+str(np.mean(f1score))[:6], "std: "+str(np.std(f1score))[:6])

        print("ROC-AUC  mean: "+str(np.mean(auc))[:6], "std: "+str(np.std(auc))[:6])

        for nctid, label, predict in zip(nctid_all, label_all, predict_all):

            if (predict > 0.5 and label == 0) or (predict < 0.5 and label == 1):

                print(nctid, label, str(predict)[:6])

        nctid2predict = {nctid:predict for nctid, predict in zip(nctid_all, predict_all)} 

        pickle.dump(nctid2predict, open('results/nctid2predict.pkl', 'wb'))

        return nctid_all, predict_all 

    def ongoing_test(self, dataloader, sample_num = 20):

        self.eval()

        best_threshold = 0.5 

        whole_loss, predict_all, label_all, nctid_all = self.generate_predict(dataloader) 

        self.train() 

        return nctid_all, predict_all 

    def test(self, dataloader, return_loss = True, validloader=None):

        # if validloader is not None:

        #   best_threshold = self.select_threshold_for_binary(validloader)

        self.eval()

        best_threshold = 0.5 

        whole_loss, predict_all, label_all, nctid_all = self.generate_predict(dataloader)

        # from HINT.utils import plot_hist

        # plt.clf()

        # prefix_name = "./figure/" + self.save_name 

        # plot_hist(prefix_name, predict_all, label_all)

        self.train()

        if return_loss:

            return whole_loss, predict_all, label_all

        else:

            print_num = 6

            auc_score, f1score, prauc_score, precision, recall, accuracy, \

            predict_1_ratio, label_1_ratio = self.evaluation(predict_all, label_all, threshold = best_threshold)

            print("ROC AUC: " + str(auc_score)[:print_num] + "\nF1: " + str(f1score)[:print_num] \

                 + "\nPR-AUC: " + str(prauc_score)[:print_num] \

                 + "\nPrecision: " + str(precision)[:print_num] \

                 + "\nrecall: "+str(recall)[:print_num] + "\naccuracy: "+str(accuracy)[:print_num] \

                 + "\npredict 1 ratio: " + str(predict_1_ratio)[:print_num] \

                 + "\nlabel 1 ratio: " + str(label_1_ratio)[:print_num])

            return auc_score, f1score, prauc_score, precision, recall, accuracy, predict_1_ratio, label_1_ratio 

    def learn(self, train_loader, valid_loader, test_loader):

        opt = torch.optim.Adam(self.parameters(), lr = self.lr, weight_decay = self.weight_decay)

        train_loss_record = [] 

        valid_loss, valid_predict, valid_label = self.test(valid_loader, return_loss=True)

        valid_loss_record = [valid_loss]

        best_valid_loss = valid_loss

        best_model = deepcopy(self)

        train_output = []

        valid_output = []

        for ep in tqdm(range(self.epoch)):

            for nctid_lst, label_vec, smiles_lst2, icdcode_lst3, criteria_lst in train_loader:

                label_vec = label_vec.to(self.device)

                output = self.forward(smiles_lst2, icdcode_lst3, criteria_lst).view(-1)  #### 32, 1 -> 32, ||  label_vec 32,

                loss = self.loss(output, label_vec.float())

                train_loss_record.append(loss.item())

                train_output.append((loss.item(), output, label_vec))

                opt.zero_grad()

                loss.backward()

                opt.step()

            valid_loss, valid_predict, valid_label = self.test(valid_loader, return_loss=True)

            valid_loss_record.append(valid_loss)

            valid_output.append((valid_loss, valid_predict, valid_label))

            print(f"valid_loss: {valid_loss}")

            print(best_valid_loss)

            if valid_loss < best_valid_loss:

                best_valid_loss = valid_loss 

                best_model = deepcopy(self)

        self.plot_learning_curve(train_loss_record, valid_loss_record)

        self = deepcopy(best_model)

        auc_score, f1score, prauc_score, precision, recall, accuracy, predict_1_ratio, label_1_ratio = self.test(test_loader, return_loss = False, validloader = valid_loader)

        return train_output, valid_output

    def plot_learning_curve(self, train_loss_record, valid_loss_record):

        plt.plot(train_loss_record)

        plt.savefig("./figure/" + self.save_name + '_train_loss.jpg')

        plt.clf() 

        plt.plot(valid_loss_record)

        plt.savefig("./figure/" + self.save_name + '_valid_loss.jpg')

        plt.clf() 

    def select_threshold_for_binary(self, validloader):

        _, prediction, label_all, nctid_all = self.generate_predict(validloader)

        best_f1 = 0

        for threshold in prediction:

            float2binary = lambda x:0 if x<threshold else 1

            predict_all = list(map(float2binary, prediction))

            f1score = precision_score(label_all, predict_all)        

            if f1score > best_f1:

                best_f1 = f1score 

                best_threshold = threshold

        return best_threshold 

class HINTModel_multi(Interaction):

    def __init__(self, molecule_encoder, disease_encoder, protocol_encoder, 

                    device, 

                    global_embed_size, 

                    highway_num_layer,

                    prefix_name, 

                    epoch = 20,

                    lr = 3e-4, 

                    weight_decay = 0, 

                    ):

        super(HINTModel_multi, self).__init__(molecule_encoder = molecule_encoder, 

                                   disease_encoder = disease_encoder, 

                                   protocol_encoder = protocol_encoder, 

                                   device = device, 

                                   prefix_name = prefix_name, 

                                   global_embed_size = global_embed_size, 

                                   highway_num_layer = highway_num_layer,

                                   epoch = epoch,

                                   lr = lr, 

                                   weight_decay = weight_decay)

        self.pred_nn = nn.Linear(self.global_embed_size, 4)

        self.loss = nn.CrossEntropyLoss()

    def forward(self, smiles_lst2, icdcode_lst3, criteria_lst):

        molecule_embed, icd_embed, protocol_embed = self.forward_get_three_encoders(smiles_lst2, icdcode_lst3, criteria_lst)

        interaction_embedding = self.forward_encoder_2_interaction(molecule_embed, icd_embed, protocol_embed)

        output = self.pred_nn(interaction_embedding)

        return output ### 32, 4

    def generate_predict(self, dataloader):

        whole_loss = 0 

        label_all, predict_all = [], []

        for nctid_lst, label_vec, smiles_lst2, icdcode_lst3, criteria_lst in dataloader:

            label_vec = label_vec.to(self.device)

            output = self.forward(smiles_lst2, icdcode_lst3, criteria_lst) 

            loss = self.loss(output, label_vec)

            whole_loss += loss.item()

            predict_all.extend(torch.argmax(output, 1).tolist())

            # predict_all.extend([i.item() for i in torch.sigmoid(output)])

            label_all.extend([i.item() for i in label_vec])

        accuracy = len(list(filter(lambda x:x[0]==x[1], zip(predict_all, label_all)))) / len(label_all)

        return whole_loss, predict_all, label_all, accuracy

    def test(self, dataloader, return_loss = True, validloader=None):

        # if validloader is not None:

        #   best_threshold = self.select_threshold_for_binary(validloader)

        self.eval()

        whole_loss, predict_all, label_all, accuracy = self.generate_predict(dataloader)

        self.train()

        return whole_loss, predict_all, label_all, accuracy

        # # from HINT.utils import plot_hist

        # # plt.clf()

        # # prefix_name = "./figure/" + self.save_name 

        # # plot_hist(prefix_name, predict_all, label_all)

        # self.train()

        # if return_loss:

        #   return whole_loss

        # else:

        #   print_num = 5

        #   auc_score, f1score, prauc_score, precision, recall, accuracy, \

        #   predict_1_ratio, label_1_ratio = self.evaluation(predict_all, label_all, threshold = best_threshold)

        #   print("ROC AUC: " + str(auc_score)[:print_num] + "\nF1: " + str(f1score)[:print_num] \

        #        + "\nPR-AUC: " + str(prauc_score)[:print_num] \

        #        + "\nPrecision: " + str(precision)[:print_num] \

        #        + "\nrecall: "+str(recall)[:print_num] + "\naccuracy: "+str(accuracy)[:print_num] \

        #        + "\npredict 1 ratio: " + str(predict_1_ratio)[:print_num] \

        #        + "\nlabel 1 ratio: " + str(label_1_ratio)[:print_num])

        #   return auc_score, f1score, prauc_score, precision, recall, accuracy, predict_1_ratio, label_1_ratio 

    def learn(self, train_loader, valid_loader, test_loader):

        opt = torch.optim.Adam(self.parameters(), lr = self.lr, weight_decay = self.weight_decay)

        train_loss_record = []

        valid_loss, predict_all, label_all, accuracy = self.test(valid_loader, return_loss=True)

        print('accuracy', accuracy)

        # valid_loss_record = [valid_loss]

        # best_valid_loss = valid_loss

        best_model = deepcopy(self)

        for ep in tqdm(range(self.epoch)):

            self.train() 

            for nctid_lst, label_vec, smiles_lst2, icdcode_lst3, criteria_lst in train_loader:

                label_vec = label_vec.to(self.device)

                output = self.forward(smiles_lst2, icdcode_lst3, criteria_lst)  #### 32, 1 -> 32, ||  label_vec 32,

                # print(label_vec.shape, output.shape, label_vec, output)

                loss = self.loss(output, label_vec)

                train_loss_record.append(loss.item())

                opt.zero_grad() 

                loss.backward() 

                opt.step()

            valid_loss, predict_all, label_all, accuracy = self.test(valid_loader, return_loss=True)

            print('accuracy', accuracy)

        return predict_all, label_all

        #   valid_loss_record.append(valid_loss)

        #   if valid_loss < best_valid_loss:

        #       best_valid_loss = valid_loss 

        #       best_model = deepcopy(self)

        # self.plot_learning_curve(train_loss_record, valid_loss_record)

        # self = deepcopy(best_model)

        # auc_score, f1score, prauc_score, precision, recall, accuracy, predict_1_ratio, label_1_ratio = self.test(test_loader, return_loss = False, validloader = valid_loader)

class HINT_nograph(Interaction):

    def __init__(self, molecule_encoder, disease_encoder, protocol_encoder, device, 

                    global_embed_size, 

                    highway_num_layer,

                    prefix_name, 

                    epoch = 20,

                    lr = 3e-4, 

                    weight_decay = 0, ):

        super(HINT_nograph, self).__init__(molecule_encoder = molecule_encoder, 

                                   disease_encoder = disease_encoder, 

                                   protocol_encoder = protocol_encoder,

                                   device = device,  

                                   global_embed_size = global_embed_size, 

                                   prefix_name = prefix_name, 

                                   highway_num_layer = highway_num_layer,

                                   epoch = epoch,

                                   lr = lr, 

                                   weight_decay = weight_decay, 

                                   ) 

        self.save_name = prefix_name + '_HINT_nograph'

        ''' ### interaction model 

        self.molecule_encoder = molecule_encoder 

        self.disease_encoder = disease_encoder 

        self.protocol_encoder = protocol_encoder 

        self.global_embed_size = global_embed_size 

        self.highway_num_layer = highway_num_layer 

        self.feature_dim = self.molecule_encoder.embedding_size + self.disease_encoder.embedding_size + self.protocol_encoder.embedding_size

        self.epoch = epoch 

        self.lr = lr 

        self.weight_decay = weight_decay 

        self.save_name = save_name

        self.f = F.relu

        self.loss = nn.BCEWithLogitsLoss()

        ##### NN 

        self.encoder2interaction_fc = nn.Linear(self.feature_dim, self.global_embed_size)

        self.encoder2interaction_highway = Highway(self.global_embed_size, self.highway_num_layer)

        self.pred_nn = nn.Linear(self.global_embed_size, 1)

        '''

        #### risk of disease 

        self.risk_disease_fc = nn.Linear(self.disease_encoder.embedding_size, self.global_embed_size)

        self.risk_disease_higway = Highway(self.global_embed_size, self.highway_num_layer)

        #### augment interaction 

        self.augment_interaction_fc = nn.Linear(self.global_embed_size*2, self.global_embed_size)

        self.augment_interaction_highway = Highway(self.global_embed_size, self.highway_num_layer)

        #### ADMET 

        self.admet_model = []

        for i in range(5):

            admet_fc = nn.Linear(self.molecule_encoder.embedding_size, self.global_embed_size).to(device)

            admet_highway = Highway(self.global_embed_size, self.highway_num_layer).to(device)

            self.admet_model.append(nn.ModuleList([admet_fc, admet_highway])) 

        self.admet_model = nn.ModuleList(self.admet_model)

        #### PK 

        self.pk_fc = nn.Linear(self.global_embed_size*5, self.global_embed_size)

        self.pk_highway = Highway(self.global_embed_size, self.highway_num_layer)

        #### trial node 

        self.trial_fc = nn.Linear(self.global_embed_size*2, self.global_embed_size)

        self.trial_highway = Highway(self.global_embed_size, self.highway_num_layer)

        ## self.pred_nn = nn.Linear(self.global_embed_size, 1)

        self.device = device 

        self = self.to(device)

    def forward(self, smiles_lst2, icdcode_lst3, criteria_lst, if_gnn = False):

        ### encoder for molecule, disease and protocol

        molecule_embed, icd_embed, protocol_embed = self.forward_get_three_encoders(smiles_lst2, icdcode_lst3, criteria_lst)

        ### interaction 

        interaction_embedding = self.forward_encoder_2_interaction(molecule_embed, icd_embed, protocol_embed)

        ### risk of disease 

        risk_of_disease_embedding = self.feed_lst_of_module(input_feature = icd_embed, 

                                                            lst_of_module = [self.risk_disease_fc, self.risk_disease_higway])

        ### augment interaction   

        augment_interaction_input = torch.cat([interaction_embedding, risk_of_disease_embedding], 1)

        augment_interaction_embedding = self.feed_lst_of_module(input_feature = augment_interaction_input, 

                                                                lst_of_module = [self.augment_interaction_fc, self.augment_interaction_highway])

        ### admet 

        admet_embedding_lst = []

        for idx in range(5):

            admet_embedding = self.feed_lst_of_module(input_feature = molecule_embed, 

                                                      lst_of_module = self.admet_model[idx])

            admet_embedding_lst.append(admet_embedding)

        ### pk 

        pk_input = torch.cat(admet_embedding_lst, 1)

        pk_embedding = self.feed_lst_of_module(input_feature = pk_input, 

                                               lst_of_module = [self.pk_fc, self.pk_highway])

        ### trial 

        trial_input = torch.cat([pk_embedding, augment_interaction_embedding], 1)

        trial_embedding = self.feed_lst_of_module(input_feature = trial_input, 

                                                  lst_of_module = [self.trial_fc, self.trial_highway])

        output = self.pred_nn(trial_embedding)

        if if_gnn == False:

            return output 

        else:

            embedding_lst = [molecule_embed, icd_embed, protocol_embed, interaction_embedding, risk_of_disease_embedding, \

                             augment_interaction_embedding] + admet_embedding_lst + [pk_embedding, trial_embedding]

            return embedding_lst

class HINTModel(HINT_nograph):

    def __init__(self, molecule_encoder, disease_encoder, protocol_encoder, 

                    device, 

                    global_embed_size, 

                    highway_num_layer,

                    prefix_name, 

                    gnn_hidden_size, 

                    epoch = 20,

                    lr = 3e-4, 

                    weight_decay = 0,):

        super(HINTModel, self).__init__(molecule_encoder = molecule_encoder, 

                                   disease_encoder = disease_encoder, 

                                   protocol_encoder = protocol_encoder, 

                                   device = device, 

                                   prefix_name = prefix_name, 

                                   global_embed_size = global_embed_size, 

                                   highway_num_layer = highway_num_layer,

                                   epoch = epoch,

                                   lr = lr, 

                                   weight_decay = weight_decay)

        self.save_name = prefix_name 

        self.gnn_hidden_size = gnn_hidden_size 

        #### GNN 

        self.adj = self.generate_adj()          

        self.gnn = GCN(

            nfeat = self.global_embed_size,

            nhid = self.gnn_hidden_size,

            nclass = 1,

            dropout = 0.6,

            init = 'uniform') 

        ### gnn's attention         

        self.node_size = self.adj.shape[0]

        '''

        self.graph_attention_model_mat = nn.ModuleList([nn.ModuleList([self.gnn_attention() \

                                        if self.adj[i,j]==1 else None  \

                                        for j in range(self.node_size)]) \

                                        for i in range(self.node_size)])

        '''

        self.graph_attention_model_mat = nn.ModuleList([nn.ModuleList([self.gnn_attention() if self.adj[i,j]==1 else None for j in range(self.node_size)]) for i in range(self.node_size)])

        # self.graph_attention_model_mat = nn.ModuleList([nn.ModuleList([self.gnn_attention() if self.adj[i,j]==1 else None for j in range(self.node_size)]) for i in range(self.node_size)])

        '''

nn.ModuleList([ nn.ModuleList([nn.Linear(3,2) for j in range(5)] + [None]) for i in range(3)])

        '''

        self.device = device 

        self = self.to(device)

    def generate_adj(self):                                     

        ##### consistent with HINT_nograph.forward

        lst = ["molecule", "disease", "criteria", 'INTERACTION', 'risk_disease', 'augment_interaction', 'A', 'D', 'M', 'E', 'T', 'PK', "final"]

        edge_lst = [("disease", "molecule"), ("disease", "criteria"), ("molecule", "criteria"), 

                    ("disease", "INTERACTION"), ("molecule", "INTERACTION"),  ("criteria", "INTERACTION"), 

                    ("disease", "risk_disease"), ('risk_disease', 'augment_interaction'), ('INTERACTION', 'augment_interaction'),

                    ("molecule", "A"), ("molecule", "D"), ("molecule", "M"), ("molecule", "E"), ("molecule", "T"),

                    ('A', 'PK'), ('D', 'PK'), ('M', 'PK'), ('E', 'PK'), ('T', 'PK'), 

                    ('augment_interaction', 'final'), ('PK', 'final')]

        adj = torch.zeros(len(lst), len(lst))

        adj = torch.eye(len(lst)) * len(lst)

        num2str = {k:v for k,v in enumerate(lst)}

        str2num = {v:k for k,v in enumerate(lst)}

        for i,j in edge_lst:

            n1,n2 = str2num[i], str2num[j]

            adj[n1,n2] = 1

            adj[n2,n1] = 1

        return adj.to(self.device) 

    def generate_attention_matrx(self, node_feature_mat):

        attention_mat = torch.zeros(self.node_size, self.node_size).to(self.device)

        for i in range(self.node_size):

            for j in range(self.node_size):

                if self.adj[i,j]!=1:

                    continue 

                feature = torch.cat([node_feature_mat[i].view(1,-1), node_feature_mat[j].view(1,-1)], 1)

                attention_model = self.graph_attention_model_mat[i][j]

                attention_mat[i,j] = torch.sigmoid(self.feed_lst_of_module(input_feature=feature, lst_of_module=attention_model))

        return attention_mat 

    ##### self.global_embed_size*2 -> 1 

    def gnn_attention(self):

        highway_nn = Highway(size = self.global_embed_size*2, num_layers = self.highway_num_layer).to(self.device)

        highway_fc = nn.Linear(self.global_embed_size*2, 1).to(self.device)

        return nn.ModuleList([highway_nn, highway_fc])  

    def forward(self, smiles_lst2, icdcode_lst3, criteria_lst, return_attention_matrix = False):

        embedding_lst = HINT_nograph.forward(self, smiles_lst2, icdcode_lst3, criteria_lst, if_gnn = True)

        ### length is 13, each is 32,50 

        batch_size = embedding_lst[0].shape[0]

        output_lst = []

        if return_attention_matrix:

            attention_mat_lst = []

        for i in range(batch_size):

            node_feature_lst = [embedding[i].view(1,-1) for embedding in embedding_lst]

            node_feature_mat = torch.cat(node_feature_lst, 0) ### 13, 50 

            attention_mat = self.generate_attention_matrx(node_feature_mat)

            output = self.gnn(node_feature_mat, self.adj * attention_mat)

            output = output[-1].view(1,-1)

            output_lst.append(output)

            if return_attention_matrix:

                attention_mat_lst.append(attention_mat)

        output_mat = torch.cat(output_lst, 0)

        if not return_attention_matrix:

            return output_mat 

        else:

            return output_mat, attention_mat_lst

    def interpret(self, complete_dataloader):

        from graph_visualize_interpret import data2graph 

        from HINT.utils import replace_strange_symbol

        for nctid_lst, status_lst, why_stop_lst, label_vec, phase_lst, \

            diseases_lst, icdcode_lst3, drugs_lst, smiles_lst2, criteria_lst in complete_dataloader: 

            output, attention_mat_lst = self.forward(smiles_lst2, icdcode_lst3, criteria_lst, return_attention_matrix=True)

            output = output.view(-1)

            batch_size = len(nctid_lst)

            for i in range(batch_size):

                name = '__'.join([nctid_lst[i], status_lst[i], why_stop_lst[i], \

                                                        str(label_vec[i].item()), str(torch.sigmoid(output[i]).item())[:5], \

                                                        phase_lst[i], diseases_lst[i], drugs_lst[i]])

                if len(name) > 150:

                    name = name[:250]

                name = replace_strange_symbol(name)

                name = name.replace('__', '_')

                name = name.replace('  ', ' ')

                name = 'interpret_result/' + name + '.png'

                print(name)

                data2graph(attention_matrix = attention_mat_lst[i], adj = self.adj, save_name = name)

    def init_pretrain(self, admet_model):

        self.molecule_encoder = admet_model.molecule_encoder

    ### generate attention matrix 

class Only_Molecule(Interaction):

    def __init__(self, molecule_encoder, disease_encoder, protocol_encoder, 

                    global_embed_size, 

                    highway_num_layer,

                    prefix_name, 

                    epoch = 20,

                    lr = 3e-4, 

                    weight_decay = 0):

        super(Only_Molecule, self).__init__(molecule_encoder=molecule_encoder, 

                                            disease_encoder=disease_encoder, 

                                            protocol_encoder=protocol_encoder, 

                                            global_embed_size = global_embed_size, 

                                            highway_num_layer = highway_num_layer,

                                            prefix_name = prefix_name, 

                                            epoch = epoch,

                                            lr = lr, 

                                            weight_decay = weight_decay,)

        self.molecule2out = nn.Linear(self.global_embed_size,1)

    def forward(self, smiles_lst2, icdcode_lst3, criteria_lst):

        molecule_embed = self.molecule_encoder.forward_smiles_lst_lst(smiles_lst2)

        return self.molecule2out(molecule_embed)

class Only_Disease(Only_Molecule):

    def __init__(self, molecule_encoder, disease_encoder, protocol_encoder, 

                    global_embed_size, 

                    highway_num_layer,

                    prefix_name, 

                    epoch = 20,

                    lr = 3e-4, 

                    weight_decay = 0):

        super(Only_Disease, self).__init__(molecule_encoder = molecule_encoder, 

                                            disease_encoder=disease_encoder, 

                                            protocol_encoder=protocol_encoder, 

                                            global_embed_size = global_embed_size, 

                                            highway_num_layer = highway_num_layer,

                                            prefix_name = prefix_name, 

                                            epoch = epoch,

                                            lr = lr, 

                                            weight_decay = weight_decay,)

        self.disease2out = self.molecule2out 

    def forward(self, smiles_lst2, icdcode_lst3, criteria_lst):

        icd_embed = self.disease_encoder.forward_code_lst3(icdcode_lst3)

        return self.disease2out(icd_embed)

def dataloader2Xy(nctid_lst, label_vec, smiles_lst2, icdcode_lst3, criteria_lst, global_icd):

    ## label_vec: (n,)

    y = label_vec 

    num_icd = len(global_icd)

    from HINT.utils import smiles_lst2fp 

    fp_lst = [smiles_lst2fp(smiles_lst).reshape(1,-1) for smiles_lst in smiles_lst2]

    fp_mat = np.concatenate(fp_lst, 0)

    # fp_mat = torch.from_numpy(fp_mat)  ### (n,2048)

    icdcode_lst = []

    for lst2 in icdcode_lst3:

        lst = list(reduce(lambda x,y:x+y, lst2))

        lst = [i.split('.')[0] for i in lst]

        lst = set(lst)  

        icd_feature = np.zeros((1,num_icd), np.int32)

        for ele in lst:

            if ele in global_icd:

                idx = global_icd.index(ele)

                icd_feature[0,idx] = 1 

        icdcode_lst.append(icd_feature)

    icdcode_mat = np.concatenate(icdcode_lst, 0)

    X = np.concatenate([fp_mat, icdcode_mat], 1)

    X = torch.from_numpy(X)

    X = X.float()

    # icdcode_mat = torch.from_numpy(icdcode_mat) 

    # X = torch.cat([fp_mat, icdcode_mat], 1)

    return X, y 

class FFNN(nn.Sequential):

    def __init__(self, molecule_dim, diseasecode_dim, 

                    global_icd, 

                    protocol_dim = 0,

                    prefix_name = 'FFNN', 

                    epoch = 10,

                    lr = 3e-4, 

                    weight_decay = 0, 

                    ):

        super(FFNN, self).__init__()

        self.molecule_dim = molecule_dim 

        self.diseasecode_dim = diseasecode_dim 

        self.protocol_dim = protocol_dim 

        self.prefix_name = prefix_name 

        self.epoch = epoch 

        self.lr = lr 

        self.weight_decay = weight_decay 

        self.global_icd = global_icd 

        self.num_icd = len(global_icd)

        self.fc_dims = [self.molecule_dim + self.diseasecode_dim + self.protocol_dim, 2000, 1000, 200, 50, 1]

        self.fc_layers = nn.ModuleList([nn.Linear(v,self.fc_dims[i+1]) for i,v in enumerate(self.fc_dims[:-1])])

        self.loss = nn.BCEWithLogitsLoss()

        self.save_name = prefix_name 

    def forward(self, X):

        for i in range(len(self.fc_layers) - 1):

            fc_layer = self.fc_layers[i]

            X = fc_layer(X)

        last_layer = self.fc_layers[-1]

        pred = F.sigmoid(last_layer(X))

        return pred 

    def learn(self, train_loader, valid_loader, test_loader):

        opt = torch.optim.Adam(self.parameters(), lr = self.lr, weight_decay = self.weight_decay)

        train_loss_record = [] 

        valid_loss = self.test(valid_loader, return_loss=True)

        valid_loss_record = [valid_loss]

        best_valid_loss = valid_loss

        best_model = deepcopy(self)

        for ep in tqdm(range(self.epoch)):

            for nctid_lst, label_vec, smiles_lst2, icdcode_lst3, criteria_lst in train_loader:

                X, _ = dataloader2Xy(nctid_lst, label_vec, smiles_lst2, icdcode_lst3, criteria_lst, self.global_icd)

                output = self.forward(X).view(-1)  #### 32, 1 -> 32, ||  label_vec 32,

                loss = self.loss(output, label_vec.float())

                train_loss_record.append(loss.item())

                opt.zero_grad() 

                loss.backward() 

                opt.step()

            valid_loss = self.test(valid_loader, return_loss=True)

            valid_loss_record.append(valid_loss)

            if valid_loss < best_valid_loss:

                best_valid_loss = valid_loss 

                best_model = deepcopy(self)

        self.plot_learning_curve(train_loss_record, valid_loss_record)

        self = deepcopy(best_model)

        auc_score, f1score, prauc_score, precision, recall, accuracy, predict_1_ratio, label_1_ratio = self.test(test_loader, return_loss = False, validloader = valid_loader)

    def evaluation(self, predict_all, label_all, threshold = 0.5):

        import pickle, os

        from sklearn.metrics import roc_curve, precision_recall_curve

        with open("predict_label.txt", 'w') as fout:

            for i,j in zip(predict_all, label_all):

                fout.write(str(i)[:4] + '\t' + str(j)[:4]+'\n')

        auc_score = roc_auc_score(label_all, predict_all)

        figure_folder = "figure"

        #### ROC-curve 

        fpr, tpr, thresholds = roc_curve(label_all, predict_all, pos_label=1)

        # roc_curve =plt.figure()

        # plt.plot(fpr,tpr,'-',label=self.save_name + ' ROC Curve ')

        # plt.legend(fontsize = 15)

        #plt.savefig(os.path.join(figure_folder,name+"_roc_curve.png"))

        #### PR-curve

        precision, recall, thresholds = precision_recall_curve(label_all, predict_all)

        # plt.plot(recall,precision, label = self.save_name + ' PR Curve')

        # plt.legend(fontsize = 15)

        # plt.savefig(os.path.join(figure_folder,self.save_name + "_pr_curve.png"))

        label_all = [int(i) for i in label_all]

        float2binary = lambda x:0 if x<threshold else 1

        predict_all = list(map(float2binary, predict_all))

        f1score = f1_score(label_all, predict_all)

        prauc_score = average_precision_score(label_all, predict_all)

        # print(predict_all)

        precision = precision_score(label_all, predict_all)

        recall = recall_score(label_all, predict_all)

        accuracy = accuracy_score(label_all, predict_all)

        predict_1_ratio = sum(predict_all) / len(predict_all)

        label_1_ratio = sum(label_all) / len(label_all)

        return auc_score, f1score, prauc_score, precision, recall, accuracy, predict_1_ratio, label_1_ratio 

    def generate_predict(self, dataloader):

        whole_loss = 0 

        label_all, predict_all = [], []

        for nctid_lst, label_vec, smiles_lst2, icdcode_lst3, criteria_lst in dataloader:

            X, _ = dataloader2Xy(nctid_lst, label_vec, smiles_lst2, icdcode_lst3, criteria_lst, self.global_icd) 

            output = self.forward(X).view(-1)  

            loss = self.loss(output, label_vec.float())

            whole_loss += loss.item()

            predict_all.extend([i.item() for i in torch.sigmoid(output)])

            label_all.extend([i.item() for i in label_vec])

        return whole_loss, predict_all, label_all

    def bootstrap_test(self, dataloader, validloader = None, sample_num = 20):

        best_threshold = 0.5

        # if validloader is not None:

        #   best_threshold = self.select_threshold_for_binary(validloader)

        self.eval()

        whole_loss, predict_all, label_all = self.generate_predict(dataloader)

        from HINT.utils import plot_hist

        plt.clf()

        prefix_name = "./figure/" + self.save_name 

        plot_hist(prefix_name, predict_all, label_all)      

        def bootstrap(length, sample_num):

            idx = [i for i in range(length)]

            from random import choices 

            bootstrap_idx = [choices(idx, k = length) for i in range(sample_num)]

            return bootstrap_idx 

        results_lst = []

        bootstrap_idx_lst = bootstrap(len(predict_all), sample_num = sample_num)

        for bootstrap_idx in bootstrap_idx_lst: 

            bootstrap_label = [label_all[idx] for idx in bootstrap_idx]     

            bootstrap_predict = [predict_all[idx] for idx in bootstrap_idx]

            results = self.evaluation(bootstrap_predict, bootstrap_label, threshold = best_threshold)

            results_lst.append(results)

        self.train() 

        auc = [results[0] for results in results_lst]

        f1score = [results[1] for results in results_lst]

        prauc_score = [results[2] for results in results_lst]

        print("PR-AUC   mean: "+str(np.mean(prauc_score))[:6], "std: "+str(np.std(prauc_score))[:6])

        print("F1       mean: "+str(np.mean(f1score))[:6], "std: "+str(np.std(f1score))[:6])

        print("ROC-AUC  mean: "+ str(np.mean(auc))[:6], "std: " + str(np.std(auc))[:6])

    def test(self, dataloader, return_loss = True, validloader=None):

        # if validloader is not None:

        #   best_threshold = self.select_threshold_for_binary(validloader)

        self.eval()

        best_threshold = 0.5 

        whole_loss, predict_all, label_all = self.generate_predict(dataloader)

        # from HINT.utils import plot_hist

        # plt.clf()

        # prefix_name = "./figure/" + self.save_name 

        # plot_hist(prefix_name, predict_all, label_all)

        self.train()

        if return_loss:

            return whole_loss

        else:

            print_num = 5

            auc_score, f1score, prauc_score, precision, recall, accuracy, \

            predict_1_ratio, label_1_ratio = self.evaluation(predict_all, label_all, threshold = best_threshold)

            print("ROC AUC: " + str(auc_score)[:print_num] + "\nF1: " + str(f1score)[:print_num] \

                 + "\nPR-AUC: " + str(prauc_score)[:print_num] \

                 + "\nPrecision: " + str(precision)[:print_num] \

                 + "\nrecall: "+str(recall)[:print_num] + "\naccuracy: "+str(accuracy)[:print_num] \

                 + "\npredict 1 ratio: " + str(predict_1_ratio)[:print_num] \

                 + "\nlabel 1 ratio: " + str(label_1_ratio)[:print_num])

            return auc_score, f1score, prauc_score, precision, recall, accuracy, predict_1_ratio, label_1_ratio 

    def plot_learning_curve(self, train_loss_record, valid_loss_record):

        plt.plot(train_loss_record)

        plt.savefig("./figure/" + self.save_name + '_train_loss.jpg')

        plt.clf() 

        plt.plot(valid_loss_record)

        plt.savefig("./figure/" + self.save_name + '_valid_loss.jpg')

        plt.clf() 

class ADMET(nn.Sequential):

    def __init__(self, mpnn_model, device):

        super(ADMET, self).__init__()

        self.num = 5 

        self.mpnn_model = mpnn_model 

        self.device = device 

        self.mpnn_dim = mpnn_model.mpnn_hidden_size 

        self.admet_model = []

        self.global_embed_size = self.mpnn_dim 

        self.highway_num_layer = 2 

        for i in range(5):

            admet_fc = nn.Linear(self.mpnn_model.mpnn_hidden_size, self.global_embed_size).to(device)

            admet_highway = Highway(self.global_embed_size, self.highway_num_layer).to(device)

            self.admet_model.append(nn.ModuleList([admet_fc, admet_highway]))

        self.admet_model = nn.ModuleList(self.admet_model)

        self.admet_pred = nn.ModuleList([nn.Linear(self.global_embed_size,1).to(device) for i in range(5)])

        self.f = F.relu 

        self.device = device 

        self = self.to(device)

    def feed_lst_of_module(self, input_feature, lst_of_module):

        x = input_feature

        for single_module in lst_of_module:

            x = self.f(single_module(x))

        return x 

    def forward(self, smiles_lst, idx):

        assert idx in list(range(5))

        '''

            xxxxxxxxxxxx

        '''

        embeds = self.mpnn_model.forward_smiles_lst_lst(smiles_lst)

        embeds = self.feed_lst_of_module(embeds, self.admet_model[idx]) 

        output = self.admet_pred[idx](embeds)

        return output 

    def test(self, valid_loader):

        pass 

    def learn(self, train_loader, valid_loader, idx):

        opt = torch.optim.Adam(self.parameters(), lr = self.lr, weight_decay = self.weight_decay)

        train_loss_record = [] 

        valid_loss = self.test(valid_loader, return_loss=True)

        valid_loss_record = [valid_loss]

        best_valid_loss = valid_loss

        best_model = deepcopy(self)

        for ep in tqdm(range(self.epoch)):

            for smiles_lst in train_loader:

                output = self.forward(smiles_lst).view(-1)  #### 32, 1 -> 32, ||  label_vec 32,

                loss = self.loss(output, label_vec.float())

                train_loss_record.append(loss.item())

                opt.zero_grad() 

                loss.backward() 

                opt.step()

            valid_loss = self.test(valid_loader, return_loss=True)

            valid_loss_record.append(valid_loss)

            if valid_loss < best_valid_loss:

                best_valid_loss = valid_loss 

                best_model = deepcopy(self)

        self = deepcopy(best_model)