DiseaseSpecific/attack.py

#%%
import pickle as pkl
from typing import Dict, Tuple, List
import os
import numpy as np
import json
import dill
import logging
import argparse 
import math
from pprint import pprint
import pandas as pd
from collections import defaultdict
import copy
import time
from tqdm import tqdm

import torch
from torch.utils.data import DataLoader
import torch.backends.cudnn as cudnn
import torch.autograd as autograd
from torch.nn import functional as F
from torch.nn.modules.loss import CrossEntropyLoss

from model import Distmult, Complex, Conve
import utils

import sys

import dill

sys.path.append("..")
import Parameters

from transformers import GPT2Tokenizer, GPT2LMHeadModel

logger =  None
def generate_nghbrs_single_entity(x, edge_nghbrs, bound):

    ret_S = set(x)
    ret_L = [x]
    b = 0 
    while(b < len(ret_L)):
        s = ret_L[b]
        if s in edge_nghbrs.keys():
            for v in edge_nghbrs[s]:
                if v not in ret_S:
                    ret_S.add(v)
                    ret_L.append(v)
                    if len(ret_L) == bound:
                        return ret_L
        b += 1
    return ret_L

def generate_nghbrs(target_data, edge_nghbrs, args):
    n_dict = {}
    for i, (s, r, o) in enumerate(target_data):
        L_s = generate_nghbrs_single_entity(s, edge_nghbrs, args.neighbor_num)
        L_o = generate_nghbrs_single_entity(o, edge_nghbrs, args.neighbor_num)
        n_dict[i] = list(set(L_s + L_o))
        n_dict[i].sort()
    return n_dict 
#%%
def check_edge(s, r, o, used_trip = None, args = None):
    """Double check"""
    if args is None:
        return True
    if not args.target_existed:
        assert (s+'_'+o in used_trip) == args.target_existed
    else:
        s = entityid_to_nodetype[s]
        o = entityid_to_nodetype[o]
        r_tp = Parameters.edge_id_to_type[int(r)]
        r_tp = r_tp.split(':')[0]
        r_tp = r_tp.split('-')
        assert s == r_tp[0] and o == r_tp[1]

def get_model_loss(batch, model, device, args = None):
    s,r,o = batch[:,0], batch[:,1], batch[:,2]

    emb_s = model.emb_e(s).squeeze(dim=1)
    emb_r = model.emb_rel(r).squeeze(dim=1)
    emb_o = model.emb_e(o).squeeze(dim=1)

    if args.add_reciprocals:
        r_rev = r + n_rel
        emb_rrev = model.emb_rel(r_rev).squeeze(dim=1)
    else:
        r_rev = r
        emb_rrev = emb_r

    pred_sr = model.forward(emb_s, emb_r, mode='rhs')
    loss_sr = model.loss(pred_sr, o) # Cross entropy loss

    pred_or = model.forward(emb_o, emb_rrev, mode='lhs')
    loss_or = model.loss(pred_or, s)

    train_loss = loss_sr + loss_or
    return train_loss

def get_model_loss_without_softmax(batch, model, device=None):

    with torch.no_grad():
        s,r,o = batch[:,0], batch[:,1], batch[:,2]

        emb_s = model.emb_e(s).squeeze(dim=1)
        emb_r = model.emb_rel(r).squeeze(dim=1)

        pred = model.forward(emb_s, emb_r)
        return -pred[range(o.shape[0]), o]

def lp_regularizer(model, weight, p):
    trainable_params = [model.emb_e.weight, model.emb_rel.weight]
    norm = 0
    for i in range(len(trainable_params)):
        norm += weight * torch.sum( torch.abs(trainable_params[i]) ** p)
    return norm

def n3_regularizer(factors, weight, p):
    norm = 0
    for f in factors:
        norm += weight * torch.sum(torch.abs(f) ** p)
    return norm / factors[0].shape[0] 

def get_train_loss(batch, model, device, args):
    #batch = batch[0].to(device)
    s,r,o = batch[:,0], batch[:,1], batch[:,2]

    emb_s = model.emb_e(s).squeeze(dim=1)
    emb_r = model.emb_rel(r).squeeze(dim=1)
    emb_o = model.emb_e(o).squeeze(dim=1)

    if args.add_reciprocals:
        r_rev = r + n_rel
        emb_rrev = model.emb_rel(r_rev).squeeze(dim=1)
    else:
        r_rev = r
        emb_rrev = emb_r

    pred_sr = model.forward(emb_s, emb_r, mode='rhs')
    loss_sr = model.loss(pred_sr, o) # loss is cross entropy loss

    pred_or = model.forward(emb_o, emb_rrev, mode='lhs')
    loss_or = model.loss(pred_or, s)

    train_loss = loss_sr + loss_or
    
    if (args.reg_weight != 0.0 and args.reg_norm == 3):
        #self.logger.info('Computing regularizer weight')
        if model == 'complex':
            emb_dim = args.embedding_dim #int(self.args.embedding_dim/2)
            lhs = (emb_s[:, :emb_dim], emb_s[:, emb_dim:])
            rel = (emb_r[:, :emb_dim], emb_r[:, emb_dim:])
            rel_rev = (emb_rrev[:, :emb_dim], emb_rrev[:, emb_dim:])
            rhs = (emb_o[:, :emb_dim], emb_o[:, emb_dim:])

            #print(lhs[0].shape, lhs[1].shape)
            factors_sr = (torch.sqrt(lhs[0] ** 2 + lhs[1] ** 2),
                        torch.sqrt(rel[0] ** 2 + rel[1] ** 2),
                        torch.sqrt(rhs[0] ** 2 + rhs[1] ** 2)
                      )
            factors_or = (torch.sqrt(lhs[0] ** 2 + lhs[1] ** 2),
                        torch.sqrt(rel_rev[0] ** 2 + rel_rev[1] ** 2),
                        torch.sqrt(rhs[0] ** 2 + rhs[1] ** 2)
                      )
        else:
            factors_sr = (emb_s, emb_r, emb_o)
            factors_or = (emb_s, emb_rrev, emb_o)

        train_loss  += n3_regularizer(factors_sr, args.reg_weight, p=3)
        train_loss  += n3_regularizer(factors_or, args.reg_weight, p=3)

    if (args.reg_weight != 0.0 and args.reg_norm == 2):
        train_loss += lp_regularizer(model, args.reg_weight, p=2)
    
    return train_loss
def hv(loss, model_params, v):
    grad = autograd.grad(loss, model_params, create_graph=True, retain_graph=True)
    Hv = autograd.grad(grad, model_params, grad_outputs=v)
    return Hv
def gather_flat_grad(grads):
    views = []
    for p in grads:
        if p.data.is_sparse:
            view = p.data.to_dense().view(-1)
        else:
            view = p.data.view(-1)
        views.append(view)
    return torch.cat(views, 0)

def get_inverse_hvp_lissa(v, model, device, param_influence, train_data, args):

    damping = args.damping
    num_samples = args.lissa_repeat
    scale = args.scale 
    train_batch_size = args.lissa_batch_size
    lissa_num_batches = math.ceil(train_data.shape[0]/train_batch_size)
    recursion_depth = int(lissa_num_batches*args.lissa_depth)

    ihvp = None
    # print('inversing hvp...')
    for i in range(num_samples):
        cur_estimate = v
        #lissa_data_iterator = iter(train_loader)
        input_data = torch.from_numpy(train_data.astype('int64'))
        actual_examples = input_data[torch.randperm(input_data.shape[0]), :]
        del input_data
        
        b_begin = 0
        for j in range(recursion_depth):
            model.zero_grad() # same as optimizer.zero_grad()
            if b_begin >= actual_examples.shape[0]:
                b_begin = 0
                input_data = torch.from_numpy(train_data.astype('int64'))
                actual_examples = input_data[torch.randperm(input_data.shape[0]), :]
                del input_data
            
            input_batch = actual_examples[b_begin: b_begin + train_batch_size]
            input_batch = input_batch.to(device)
            
            train_loss = get_train_loss(input_batch, model, device, args)
            
            hvp = hv(train_loss, param_influence, cur_estimate)
            cur_estimate = [_a + (1-damping)*_b - _c / scale for _a, _b, _c in zip(v, cur_estimate, hvp)]
            # if (j%200 == 0) or (j == recursion_depth -1 ):
            #     logger.info("Recursion at depth %s: norm is %f" % (j, np.linalg.norm(gather_flat_grad(cur_estimate).cpu().numpy())))
            
            b_begin += train_batch_size
        
        if ihvp == None:
            ihvp = [_a / scale for _a in cur_estimate]
        else:
            ihvp = [_a + _b / scale for _a, _b in zip(ihvp, cur_estimate)]

    # logger.info("Final ihvp norm is %f" % (np.linalg.norm(gather_flat_grad(ihvp).cpu().numpy())))
    return_ihvp = gather_flat_grad(ihvp)
    return_ihvp /= num_samples
    
    return return_ihvp

#%%
def before_global_attack(device, n_rel, data, target_data, neighbors, model, 
                    filters:Dict[str, Dict[Tuple[str, int], torch.Tensor]],
                    entityid_to_nodetype, batch_size, args, lissa_path, target_disease):

    if os.path.exists(lissa_path) and not args.update_lissa:
        with open(lissa_path, 'rb') as fl:
            ret = dill.load(fl)
        return ret
    ret = {}

    test_data = []
    for i in target_disease:
        tp = entityid_to_nodetype[str(i)]
        # r = torch.LongTensor([[10]]).to(device)
        assert tp == 'disease'
        if tp == 'disease':
            for target in target_data:
                test_data.append([str(target), str(10), str(i)])
    test_data = np.array(test_data)

    for target_trip in tqdm(test_data):

        target_trip_ori = target_trip
        trip_name = '_'.join(list(target_trip_ori))
        target_trip = target_trip[None, :] # add a batch dimension
        target_trip = torch.from_numpy(target_trip.astype('int64')).to(device)
        # target_s, target_r, target_o = target_trip[:,0], target_trip[:,1], target_trip[:,2]
        # target_vec = model.score_triples_vec(target_s, target_r, target_o)

        model.eval()
        model.zero_grad()
        target_loss = get_model_loss(target_trip, model, device)
        target_grads = autograd.grad(target_loss, param_influence)

        model.train()
        inverse_hvp = get_inverse_hvp_lissa(target_grads, model, device, 
                                            param_influence, data, args)
        model.eval()
        inverse_hvp = inverse_hvp.detach().cpu().unsqueeze(0)
        ret[trip_name] = inverse_hvp
    with open(lissa_path, 'wb') as fl:
        dill.dump(ret, fl)
    return ret
    
def global_addtion_attack(device, n_rel, data, target_data, neighbors, model, 
                    filters:Dict[str, Dict[Tuple[str, int], torch.Tensor]],
                    entityid_to_nodetype, batch_size, args, lissa, target_disease):

    logger.info('------  Generating edits per target triple ------')
    start_time = time.time()
    logger.info('Start time: {0}'.format(str(start_time)))

    used_trip = set()
    print("Processing used triples ...")
    for s, r, o in tqdm(data):
        used_trip.add(s+'_'+o)
        # used_trip.add(o+'_'+s)
    print('Size of used triples:', len(used_trip))
    logger.info('Size of used triples: {0}'.format(len(used_trip)))

    ret_trip = []
    score_record = []
    real_add_rank_ratio = 0

    with open(score_path, 'rb') as fl:
        score_record = pkl.load(fl)
    for i, target in enumerate(target_data):

        print('\n\n------  Attacking target tripid:', i, 'tot:', len(target_data), '   ------')
        # lissa_hvp = []
        target_trip = []
        for disease in target_disease:
            target_trip.append([target, str(10), disease])
        #     nm = '{}_{}_{}'.format(target, 10, disease)
        #     lissa_hvp.append(lissa[nm])
        # lissa_hvp = torch.cat(lissa_hvp, dim = 0).to(device)

        target_trip = np.array(target_trip)
        target_trip = torch.from_numpy(target_trip.astype('int64')).to(device)

        model.eval()
        model.zero_grad()
        target_loss = get_model_loss(target_trip, model, device)
        target_grads = autograd.grad(target_loss, param_influence)

        model.train()
        inverse_hvp = get_inverse_hvp_lissa(target_grads, model, device, 
                                            param_influence, data, args)

        model.eval()

        nghbr_trip = []
        s = str(target)
        tp = entityid_to_nodetype[s]
        for nghbr in tqdm(neighbors):
            o = str(nghbr)
            if s!=o and s+'_'+o not in used_trip:
                for r in range(n_rel):
                    if (tp, r) in filters["rhs"].keys() and filters["rhs"][(tp, r)][int(o)] == True:
                            nghbr_trip.append([s, str(r), o])

        nghbr_trip = np.asarray(nghbr_trip)
        influences = []      
        edge_losses = []
        
        # nghbr_cos_log_prob, nghbr_LM_log_prob = score_record[i]
        # assert nghbr_cos_log_prob.shape[0] == nghbr_trip.shape[0]
        
        for train_trip in tqdm(nghbr_trip):
            #model.train() #batch norm cannot be used here
            train_trip = train_trip[None, :] # add batch dim
            train_trip = torch.from_numpy(train_trip.astype('int64')).to(device)
            #### L-train gradient ####
            edge_loss = get_model_loss_without_softmax(train_trip, model, device).squeeze()
            edge_losses.append(edge_loss.unsqueeze(0).detach())
            model.zero_grad()
            train_loss = get_model_loss(train_trip, model, device, args)
            train_grads = autograd.grad(train_loss, param_influence)
            train_grads = gather_flat_grad(train_grads)
            influence = torch.dot(inverse_hvp, train_grads) #default dim=1
            influences.append(influence.unsqueeze(0).detach())

        edge_losses = torch.cat(edge_losses, dim = -1)
        influences = torch.cat(influences, dim = -1)
        edge_losses_log_prob = torch.log(F.softmax(-edge_losses, dim = -1))
        influences_log_prob = torch.log(F.softmax(influences, dim = -1))

        inf_score_sorted, influences_sort = torch.sort(influences_log_prob, -1, descending=True)
        edge_score_sorted, edge_sort = torch.sort(edge_losses_log_prob, -1, descending=True)
        influences_sort = influences_sort.cpu().numpy()
        edge_sort = edge_sort.cpu().numpy()
        inf_score_sorted = inf_score_sorted.cpu().numpy()
        edge_score_sorted = edge_score_sorted.cpu().numpy()

        logger.info('')
        logger.info('Top 8 inf_score: {}'.format(" ".join(map(str, list(inf_score_sorted[:8])))))
        logger.info('Top 8 edge_score: {}'.format(" ".join(map(str, list(edge_score_sorted[:8])))))

        nghbr_cos_log_prob = influences_log_prob.detach().cpu().numpy()
        nghbr_LM_log_prob = edge_losses_log_prob.detach().cpu().numpy()
        max_sim = np.max(nghbr_cos_log_prob)
        min_sim = np.min(nghbr_cos_log_prob)
        max_LM = np.max(nghbr_LM_log_prob)
        min_LM = np.min(nghbr_LM_log_prob)

        # final_score = nghbr_cos_log_prob + nghbr_LM_log_prob
        final_score = nghbr_cos_log_prob

        index = np.argmax(final_score[:-1])
        # p = np.where(index == edge_sort)[0][0]
        # logger.info('Added edge\'s edge rank ratio: {}'.format(p / edge_sort.shape[0]))
        real_add_rank_ratio += p
        add_trip = nghbr_trip[index]
        logger.info('max_inf: {0:.8}, min_inf: {1:.8}, max_edge: {2:.8}, min_edge: {3:.8}'.format(max_sim, min_sim, max_LM, min_LM))
        logger.info('Attack trip: {0}_{1}_{2}.\n Influnce score: {3:.8}. Edge score: {4:.8}.'.format(add_trip[0], add_trip[1], add_trip[2],
                                                                                                                            nghbr_cos_log_prob[index], nghbr_LM_log_prob[index]))
        ret_trip.append(add_trip)   
        score_record.append((nghbr_cos_log_prob, nghbr_LM_log_prob)) 
    real_add_rank_ratio = real_add_rank_ratio  / target_data.shape[0]
    logger.info('Mean real ratio: {}.'.format(real_add_rank_ratio))
    return ret_trip, score_record

def addition_attack(param_influence, device, n_rel, data, target_data, neighbors, model, 
                    filters:Dict[str, Dict[Tuple[str, int], torch.Tensor]],
                    entityid_to_nodetype, batch_size, args, load_Record = False, divide_bound = None, data_mean = None, data_std = None, cache_intermidiate = True):

    if logger:
        logger.info('------  Generating edits per target triple ------')
    start_time = time.time()
    if logger:
        logger.info('Start time: {0}'.format(str(start_time)))

    used_trip = set()
    print("Processing used triples ...")
    for s, r, o in tqdm(data):
        used_trip.add(s+'_'+o)
        # used_trip.add(o+'_'+s)
    print('Size of used triples:', len(used_trip))
    if logger:
        logger.info('Size of used triples: {0}'.format(len(used_trip)))

    nghbr_trip_len = []
    ret_trip = []
    score_record = []
    direct_add_rank_ratio = 0
    real_add_rank_ratio = 0
    bad_ratio = 0

    RRcord = []
    print('****'*10)
    if load_Record:
        print('Load intermidiate file')
        with open(intermidiate_path, 'rb') as fl:
            RRcord = dill.load(fl)
    else:
        print('Donnot load intermidiate file')

    for i, target_trip in enumerate(target_data):

        print('\n\n------  Attacking target tripid:', i, '   ------')
        target_nghbrs = neighbors[i]
        for a in target_nghbrs:
            if str(a) == '-1':
                raise Exception('pppp')

        target_trip_ori = target_trip
        check_edge(target_trip[0], target_trip[1], target_trip[2], used_trip)
        target_trip = target_trip[None, :] # add a batch dimension
        target_trip = torch.from_numpy(target_trip.astype('int64')).to(device)
        # target_s, target_r, target_o = target_trip[:,0], target_trip[:,1], target_trip[:,2]
        # target_vec = model.score_triples_vec(target_s, target_r, target_o)

        model.eval()

        if load_Record:
            o_target_trip, nghbr_trip, edge_losses, influences, edge_losses_log_prob, influences_log_prob = RRcord[i]
            assert (o_target_trip.cpu() == target_trip.cpu()).sum().item() == 3
        else:
            model.zero_grad()
            target_loss = get_model_loss(target_trip, model, device, args)
            target_grads = autograd.grad(target_loss, param_influence)

            model.train()
            inverse_hvp = get_inverse_hvp_lissa(target_grads, model, device, 
                                                param_influence, data, args)

            model.eval()
            nghbr_trip = []
            valid_trip = 0
            if args.candidate_mode == 'quadratic':
                s_o_list = [(i, j) for i in target_nghbrs for j in target_nghbrs]
            elif args.candidate_mode == 'linear':
                s_o_list = [(j, i) for i in target_nghbrs for j in [target_trip_ori[0], target_trip_ori[2]]] \
                        +  [(i, j) for i in target_nghbrs for j in [target_trip_ori[0], target_trip_ori[2]]]
            else:
                raise Exception('Wrong candidate_mode: '+args.candidate_mode)
            for s, o in tqdm(s_o_list): 
                tp = entityid_to_nodetype[s]
                if s!=o and s+'_'+o not in used_trip:
                    for r in range(n_rel):
                        if (tp, r) in filters["rhs"].keys() and filters["rhs"][(tp, r)][int(o)] == True:
                            # check_edge(s, r, o)
                            valid_trip += 1
                            nghbr_trip.append([s, str(r), o])
                            # logger.info('{0}_{1}_{2}'.format(s, str(r), o))
            nghbr_trip_len.append(len(nghbr_trip))
            print('Valid trip:', valid_trip)

            if target_trip_ori[0]+'_'+target_trip_ori[2] not in used_trip:
                nghbr_trip.append(target_trip_ori)
            nghbr_trip = np.asarray(nghbr_trip)
            print("Edge scoring ...")

            influences = []
            edge_losses = []

            for train_trip in tqdm(nghbr_trip):
                #model.train() #batch norm cannot be used here
                train_trip = train_trip[None, :] # add batch dim
                train_trip = torch.from_numpy(train_trip.astype('int64')).to(device)
                #### L-train gradient ####
                edge_loss = get_model_loss_without_softmax(train_trip, model, device).squeeze()
                edge_losses.append(edge_loss.unsqueeze(0).detach())
                model.zero_grad()
                train_loss = get_model_loss(train_trip, model, device, args)
                train_grads = autograd.grad(train_loss, param_influence)
                train_grads = gather_flat_grad(train_grads)
                influence = torch.dot(inverse_hvp, train_grads) #default dim=1
                influences.append(influence.unsqueeze(0).detach())  

            edge_losses = torch.cat(edge_losses, dim = -1)
            influences = torch.cat(influences, dim = -1)
            edge_losses_log_prob = torch.log(F.softmax(-edge_losses, dim = -1))
            influences_log_prob = torch.log(F.softmax(influences, dim = -1))
            std_scale = torch.std(edge_losses_log_prob) / torch.std(influences_log_prob)
            influences_log_prob = (influences_log_prob - influences_log_prob.mean()) * std_scale + edge_losses_log_prob.mean()
            
            RRcord.append([target_trip.detach(), nghbr_trip, edge_losses, influences, edge_losses_log_prob, influences_log_prob])

        inf_score_sorted, influences_sort = torch.sort(influences_log_prob, -1, descending=True)
        edge_score_sorted, edge_sort = torch.sort(edge_losses_log_prob, -1, descending=True)

        influences_sort = influences_sort.cpu().numpy()
        edge_sort = edge_sort.cpu().numpy()
        inf_score_sorted = inf_score_sorted.cpu().numpy()
        edge_score_sorted = edge_score_sorted.cpu().numpy()
        edge_losses = edge_losses.cpu().numpy()

        p = np.where(influences_sort[0] == edge_sort)[0][0]
        direct_add_rank_ratio += p / edge_sort.shape[0]
        if logger:
            logger.info('Top 8 inf_score: {}'.format(" ".join(map(str, list(inf_score_sorted[:8])))))
            logger.info('Top 8 edge_score: {}'.format(" ".join(map(str, list(edge_score_sorted[:8])))))

        nghbr_cos_log_prob = influences_log_prob.detach().cpu().numpy()
        nghbr_LM_log_prob = edge_losses_log_prob.detach().cpu().numpy()
        max_sim = nghbr_cos_log_prob[influences_sort[0]]
        min_sim = nghbr_cos_log_prob[influences_sort[-1]]
        max_LM = nghbr_LM_log_prob[edge_sort[0]]
        min_LM = nghbr_LM_log_prob[edge_sort[-1]]
        direct_score_0 = 0
        direct_score_1 = 0
        if target_trip_ori[0]+'_'+target_trip_ori[2] not in used_trip:
            direct_score_0 = nghbr_cos_log_prob[-1]
            direct_score_1 = nghbr_LM_log_prob[-1]
    
        # bound = math.log(1 / nghbr_LM_log_prob.shape[0])
        bound = 1 - args.reasonable_rate
        edge_losses = (edge_losses - data_mean) / data_std
        edge_losses_prob =  1 / ( 1 + np.exp(edge_losses - divide_bound) )
        nghbr_LM_log_prob[edge_losses_prob < bound] = -(1e20)

        final_score = nghbr_cos_log_prob + nghbr_LM_log_prob

        index = np.argmax(final_score[:-1])
        sort_index = [(i, final_score[i])for i in range(len(final_score) - 1)]
        sort_index = sorted(sort_index, key=lambda x: x[1], reverse=True)
        assert sort_index[0][0] == index

        p = np.where(index == edge_sort)[0][0]
        if logger:
            logger.info('Bad edge ratio: {}'.format((edge_losses_prob < bound).mean()))
            logger.info('Bounded edge\'s edge rank ratio: {}'.format(p / edge_sort.shape[0]))
        real_add_rank_ratio += p / edge_sort.shape[0]
        bad_ratio += (edge_losses_prob < bound).mean()

        add_trip = nghbr_trip[index]

        if (int(add_trip[0]) == int(-1)):
            add_trip[0], add_trip[1], add_trip[2] = -1, -1, -1
            print(final_score.shape, index, edge_losses_prob[index], bound)
            raise Exception('??')

        if logger:
            logger.info('max_inf: {0:.8}, min_inf: {1:.8}, max_edge: {2:.8}, min_edge: {3:.8}'.format(max_sim, min_sim, max_LM, min_LM))
            logger.info('Target trip: {0}_{1}_{2}. Attack trip: {3}_{4}_{5}.\n Influnce score: {6:.8}. Edge score: {7:.8}. Direct score: {8:.8} + {9:.8}'.format(target_trip_ori[0],target_trip_ori[1], target_trip_ori[2], 
                                                                                                                            add_trip[0], add_trip[1], add_trip[2],
                                                                                                                            nghbr_cos_log_prob[index], nghbr_LM_log_prob[index],
                                                                                                                            direct_score_0, direct_score_1))
        if (args.added_edge_num == '' or int(args.added_edge_num) == 1):
            ret_trip.append(add_trip)
        else:
            edge_num = int(args.added_edge_num)
            for i in range(edge_num):
                ret_trip.append(nghbr_trip[sort_index[i][0]])
        score_record.append((nghbr_cos_log_prob, nghbr_LM_log_prob)) 
    
    if not load_Record and cache_intermidiate:
        with open(intermidiate_path, 'wb') as fl:
            dill.dump(RRcord, fl)
    direct_add_rank_ratio = direct_add_rank_ratio / target_data.shape[0]
    real_add_rank_ratio = real_add_rank_ratio  / target_data.shape[0]
    bad_ratio = bad_ratio / target_data.shape[0]
    if logger:
        logger.info('Mean direct ratio: {}. Mean real ratio: {}. Mean bad ratio: {}'.format(direct_add_rank_ratio, real_add_rank_ratio, bad_ratio))
    return ret_trip, score_record

def calculate_edge_bound(data, model, device, n_ent):

    tmp = np.random.choice(a = data.shape[0], size = data.shape[0] // 10, replace=False)
    existed_data= data[tmp, :]

    print('calculating edge bound ...')
    print(existed_data.shape)

    existed_edge = set()
    for src_trip in existed_data:
        existed_edge.add('_'.join(list(src_trip)))
    
    not_existed = []
    for s, r, o in  existed_data:

        if np.random.randint(0, n_ent) % 2 == 0:
            while True:
                oo = np.random.randint(0, n_ent)
                if '_'.join([s, r, str(oo)]) not in existed_edge:
                    not_existed.append([s, r, str(oo)])
                    break
        else:
            while True:
                ss = np.random.randint(0, n_ent)
                if '_'.join([str(ss), r, o]) not in existed_edge:
                    not_existed.append([str(ss), r, o])
                    break   
    existed_data = np.array(existed_data)
    not_existed = np.array(not_existed)
    existed_data = torch.from_numpy(existed_data.astype('int64')).to(device)
    not_existed = torch.from_numpy(not_existed.astype('int64')).to(device)
    loss_existed = get_model_loss_without_softmax(existed_data, model).cpu().numpy()
    loss_not_existed = get_model_loss_without_softmax(not_existed, model).cpu().numpy()
    tot_loss = np.hstack((loss_existed, loss_not_existed))
    tot_mean, tot_std = np.mean(tot_loss), np.std(tot_loss)

    loss_existed = (loss_existed - tot_mean) / tot_std
    loss_not_existed = (loss_not_existed - tot_mean) / tot_std

    print('Tot mean: {}, Tot std: {}'.format(tot_mean, tot_std))

    # print(np.mean(loss_existed), np.std(loss_existed), np.max(loss_existed))
    # print(np.mean(loss_not_existed), np.std(loss_not_existed), np.min(loss_not_existed))
    l_mean, l_std = np.mean(loss_existed), np.std(loss_existed)
    r_mean, r_std = np.mean(loss_not_existed), np.std(loss_not_existed)

    A = -1/(l_std**2) + 1/(r_std**2)
    B = 2 * (-r_mean/(r_std**2) + l_mean/(l_std**2))
    C = (r_mean**2)/(r_std**2)-(l_mean**2)/(l_std**2) + np.log((r_std**2)/(l_std**2))

    delta = B**2 - 4*A*C

    x_1 = ( -B + math.sqrt(delta) ) / (2*A)
    x_2 = ( -B - math.sqrt(delta) ) / (2*A)

    x = None
    if (x_1 > l_mean and x_1 < r_mean):
        x = x_1
    if (x_2 > l_mean and x_2 < r_mean):
        x = x_2
    if not x:
        raise Exception('Bad model!!!!')
    TP = (loss_existed < x).mean()
    TN = (loss_not_existed > x).mean()
    FP = (loss_not_existed < x).mean()
    FN = (loss_existed > x).mean()
    print('X:{}, TP:{}, TN:{}, FP:{}, FN{}'.format(x, TP, TN, FP, FN))

    sig_existed = 1 / ( 1 + np.exp(loss_existed- x) ) # negtive important
    sig_not_existed = 1 / ( 1 + np.exp(loss_not_existed - x) )

    print('Positive mean score:', sig_existed.mean(),'Negetive mean score:', sig_not_existed.mean())

    return x, tot_mean, tot_std


#%%
if __name__ == '__main__':
    parser = utils.get_argument_parser()
    parser = utils.add_attack_parameters(parser)
    args = parser.parse_args()
    args = utils.set_hyperparams(args)

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    args.device = device
    args.device1 = device
    if torch.cuda.device_count() >= 2:
        args.device = "cuda:0"
        args.device1 = "cuda:1"

    utils.seed_all(args.seed)
    np.set_printoptions(precision=5)
    cudnn.benchmark = False

    model_name = '{0}_{1}_{2}_{3}_{4}'.format(args.model, args.embedding_dim, args.input_drop, args.hidden_drop, args.feat_drop)
    model_path = 'saved_models/{0}_{1}.model'.format(args.data, model_name)
    data_path = os.path.join('processed_data', args.data)
    target_path = os.path.join(data_path, 'DD_target_{0}_{1}_{2}_{3}_{4}_{5}.txt'.format(args.model, args.data, args.target_split, args.target_size, 'exists:'+str(args.target_existed), args.attack_goal))
    lissa_path = 'lissa/{0}_{1}_{2}'.format(args.model, 
                                                args.data, 
                                                args.target_size)
    intermidiate_path = 'intermidiate/{0}_{1}_{2}_{3}_{4}_{5}_{6}'.format(args.model, 
                                                                args.target_split, 
                                                                args.target_size, 
                                                                'exists:'+str(args.target_existed),
                                                                args.neighbor_num,
                                                                args.candidate_mode,
                                                                args.attack_goal)
    log_path = 'logs/attack_logs/cos_{0}_{1}_{2}_{3}_{4}_{5}_{6}_{7}'.format(args.model, 
                                                                args.target_split, 
                                                                args.target_size, 
                                                                'exists:'+str(args.target_existed),
                                                                args.neighbor_num,
                                                                args.candidate_mode,
                                                                args.attack_goal,
                                                                str(args.reasonable_rate))
    print(log_path)
    attack_path = os.path.join('attack_results', args.data, 'cos_{0}_{1}_{2}_{3}_{4}_{5}_{6}_{7}{8}.txt'.format(args.model, 
                                                                                                        args.target_split, 
                                                                                                        args.target_size, 
                                                                                                        'exists:'+str(args.target_existed),
                                                                                                        args.neighbor_num,
                                                                                                        args.candidate_mode,
                                                                                                        args.attack_goal,
                                                                                                        str(args.reasonable_rate),
                                                                                                        str(args.added_edge_num)))

    logging.basicConfig(format = '%(asctime)s - %(levelname)s - %(name)s -   %(message)s',
                                datefmt = '%m/%d/%Y %H:%M:%S',
                                level = logging.INFO,
                                filename = log_path
                            )
    logger = logging.getLogger(__name__)
    logger.info(vars(args))
    #%%
    n_ent, n_rel, ent_to_id, rel_to_id = utils.generate_dicts(data_path)
    data  = utils.load_data(os.path.join(data_path, 'all.txt'))
    with open(os.path.join(data_path, 'filter.pickle'), 'rb') as fl:
        filters = pkl.load(fl)
    with open(os.path.join(data_path, 'entityid_to_nodetype.json'), 'r') as fl:
        entityid_to_nodetype = json.load(fl)
    with open(os.path.join(data_path, 'edge_nghbrs.pickle'), 'rb') as fl:
        edge_nghbrs = pkl.load(fl)
    with open(os.path.join(data_path, 'disease_meshid.pickle'), 'rb') as fl:
        disease_meshid = pkl.load(fl)
    with open(os.path.join(data_path, 'entities_dict.json'), 'r') as fl:
        entity_to_id = json.load(fl)
    with open(Parameters.GNBRfile+'entity_raw_name', 'rb') as fl:
        entity_raw_name = pkl.load(fl)
    #%%
    init_mask = np.asarray([0] * n_ent).astype('int64')
    init_mask = (init_mask == 1)
    for k, v in filters.items():
        for kk, vv in v.items():
            tmp = init_mask.copy()
            tmp[np.asarray(vv)] = True
            t = torch.ByteTensor(tmp).to(args.device)
            filters[k][kk] = t
    #%%
    model = utils.load_model(model_path, args, n_ent, n_rel, args.device)
    divide_bound, data_mean, data_std = calculate_edge_bound(data, model, args.device, n_ent)
    # index = torch.LongTensor([0, 1]).to(device)
    # print(model.emb_rel(index)[:, :32])
    # print(model.emb_e(index)[:, :32])
    # raise Exception
    #%%
    target_data = utils.load_data(target_path)
    if args.attack_goal == 'single':
        neighbors = generate_nghbrs(target_data, edge_nghbrs, args)
    elif args.attack_goal == 'global':
        s_set = set()
        for s, r, o in target_data:
            s_set.add(s)
        target_data = list(s_set)
        target_data.sort()
        target_data = np.array(target_data, dtype=str)
        neighbors = []
        for i in list(range(n_ent)):
            tp = entityid_to_nodetype[str(i)]
            # r = torch.LongTensor([[10]]).to(device)
            if tp == 'gene':
                neighbors.append(str(i))
        target_disease = []
        tid = 1
        bound = 50
        while True:
            meshid = disease_meshid[tid][0]
            fre = disease_meshid[tid][1]
            if len(entity_raw_name[meshid]) > 4:
                target_disease.append(entity_to_id[meshid])
                bound -= 1
                if bound == 0:
                    break
            tid += 1
    else:
        raise Exception('Wrong attack_goal: '+args.attack_goal)

    param_optimizer = list(model.named_parameters())
    param_influence = []
    for n,p in param_optimizer:
        param_influence.append(p)
    if args.attack_goal == 'single':
        len_list = []
        for v in neighbors.values():
            len_list.append(len(v))
        mean_len = np.mean(len_list)
    else:
        mean_len = len(neighbors)
    print('Mean length of neighbors:', mean_len)
    logger.info("Mean length of neighbors: {0}".format(mean_len))

    # GPT_LM = LMscore_calculator(data_path, args)
    lissa_num_batches = math.ceil(data.shape[0]/args.lissa_batch_size)
    logger.info('-------- Lissa Params for IHVP --------')
    logger.info('Damping: {0}'.format(args.damping))
    logger.info('Lissa_repeat: {0}'.format(args.lissa_repeat))
    logger.info('Lissa_depth: {0}'.format(args.lissa_depth))
    logger.info('Scale: {0}'.format(args.scale))
    logger.info('Lissa batch size: {0}'.format(args.lissa_batch_size))
    logger.info('Lissa num bacthes: {0}'.format(lissa_num_batches))

    score_path = os.path.join('attack_results', args.data, 'score_cos_{0}_{1}_{2}_{3}_{4}_{5}_{6}_{7}{8}.txt'.format(args.model, 
                                                                                                        args.target_split, 
                                                                                                        args.target_size, 
                                                                                                        'exists:'+str(args.target_existed),
                                                                                                        args.neighbor_num,
                                                                                                        args.candidate_mode,
                                                                                                        args.attack_goal,
                                                                                                        str(args.reasonable_rate),
                                                                                                        str(args.added_edge_num)))

    if args.attack_goal == 'single':
        attack_trip, score_record = addition_attack(param_influence, args.device, n_rel, data, target_data, neighbors, model, filters, entityid_to_nodetype, args.attack_batch_size, args, load_Record = args.load_existed, divide_bound = divide_bound, data_mean = data_mean, data_std = data_std)
    else:
        # lissa = before_global_attack(args.device, n_rel, data, target_data, neighbors, model, filters, entityid_to_nodetype, args.attack_batch_size, args, lissa_path, target_disease)

        attack_trip, score_record = global_addtion_attack(args.device, n_rel, data, target_data, neighbors, model, filters, entityid_to_nodetype, args.attack_batch_size, args, None, target_disease)

    utils.save_data(attack_path, attack_trip)

    logger.info("Attack triples are saved in " + attack_path)
    with open(score_path, 'wb') as fl:
        pkl.dump(score_record, fl)