text: """
  Utility functions for link prediction
  Most code is adapted from authors' implementation of RGCN link prediction:
  https://github.com/MichSchli/RelationPrediction
  
  """
  
  import os
  import numpy as np
  import torch
  import dgl
  import logging
  from tqdm import tqdm
  
  logging.basicConfig(level=logging.DEBUG)
  logger = logging.getLogger('Framework')
  logger.setLevel(logging.DEBUG)
  
  
  #######################################################################
  #
  # Utility function for building training and testing graphs
  #
  #######################################################################
  
  def get_adj_and_degrees(num_nodes, triplets):
      """ Get adjacency list and degrees of the graph
      """
      adj_list = [[] for _ in range(num_nodes)]
      for i, triplet in enumerate(triplets):
          adj_list[triplet[0]].append([i, triplet[2]])
          adj_list[triplet[2]].append([i, triplet[0]])
  
      degrees = np.array([len(a) for a in adj_list])
      adj_list = [np.array(a) for a in adj_list]
      return adj_list, degrees
  
  
  def sample_edge_neighborhood(adj_list, degrees, n_triplets, sample_size):
      """Sample edges by neighborhool expansion.
  
      This guarantees that the sampled edges form a connected graph, which
      may help deeper GNNs that require information from more than one hop.
      """
      edges = np.zeros((sample_size,), dtype=np.int32)
  
      # initialize
      sample_counts = np.array([d for d in degrees])
      picked = np.array([False for _ in range(n_triplets)])
      seen = np.array([False for _ in degrees])
  
      for i in range(0, sample_size):
          weights = sample_counts * seen
  
          if np.sum(weights) == 0:
              weights = np.ones_like(weights)
              weights[np.where(sample_counts == 0)] = 0
  
          probabilities = weights / np.sum(weights)
          chosen_vertex = np.random.choice(np.arange(degrees.shape[0]),
                                           p=probabilities)
          chosen_adj_list = adj_list[chosen_vertex]
          seen[chosen_vertex] = True
  
          chosen_edge = np.random.choice(np.arange(chosen_adj_list.shape[0]))
          chosen_edge = chosen_adj_list[chosen_edge]
          edge_number = chosen_edge[0]
  
          while picked[edge_number]:
              chosen_edge = np.random.choice(np.arange(chosen_adj_list.shape[0]))
              chosen_edge = chosen_adj_list[chosen_edge]
              edge_number = chosen_edge[0]
  
          edges[i] = edge_number
          other_vertex = chosen_edge[1]
          picked[edge_number] = True
          sample_counts[chosen_vertex] -= 1
          sample_counts[other_vertex] -= 1
          seen[other_vertex] = True
  
      return edges
  
  
  def sample_edge_uniform(_, __, n_triplets, sample_size):
      """Sample edges uniformly from all the edges."""
      all_edges = np.arange(n_triplets)
      return np.random.choice(all_edges, sample_size, replace=False)
  
  
  def generate_sampled_graph_and_labels(triplets, sample_size, split_size,
                                        num_rels, adj_list, degrees,
                                        negative_rate, sampler="uniform"):
      """Get training graph and signals
      First perform edge neighborhood sampling on graph, then perform negative
      sampling to generate negative samples
      """
      # perform edge neighbor sampling
      if sampler == "uniform":
          edges = sample_edge_uniform(
              adj_list, degrees, len(triplets), sample_size)
      elif sampler == "neighbor":
          edges = sample_edge_neighborhood(
              adj_list, degrees, len(triplets), sample_size)
      else:
          raise ValueError(
              "Sampler type must be either 'uniform' or 'neighbor'.")
  
      # relabel nodes to have consecutive node ids
      edges = triplets[edges]
      src, rel, dst = edges.transpose()
      uniq_v, edges = np.unique((src, dst), return_inverse=True)
      src, dst = np.reshape(edges, (2, -1))
      relabeled_edges = np.stack((src, rel, dst)).transpose()
  
      # negative sampling
      samples, labels = negative_sampling(relabeled_edges, len(uniq_v),
                                          negative_rate)
  
      # further split graph, only half of the edges will be used as graph
      # structure, while the rest half is used as unseen positive samples
      split_size = int(sample_size * split_size)
      graph_split_ids = np.random.choice(np.arange(sample_size),
                                         size=split_size, replace=False)
      src = src[graph_split_ids]
      dst = dst[graph_split_ids]
      rel = rel[graph_split_ids]
  
      # build DGL graph
      logger.info("# sampled nodes: {}".format(len(uniq_v)))
      logger.info("# sampled edges: {}".format(len(src) * 2))
      g, rel, norm = build_graph_from_triplets(len(uniq_v), num_rels,
                                               (src, rel, dst))
      return g, uniq_v, rel, norm, samples, labels
  
  
  def comp_deg_norm(g):
      g = g.local_var()
      in_deg = g.in_degrees(range(g.number_of_nodes())).float().numpy()
      norm = 1.0 / in_deg
      norm[np.isinf(norm)] = 0
      return norm
  
  
  def node_norm_to_edge_norm(g, node_norm):
      g = g.local_var()
      # convert to edge norm
      g.ndata['norm'] = node_norm
      g.apply_edges(lambda edges: {'norm': edges.dst['norm']})
      return g.edata['norm']
  
  
  def build_graph_from_triplets(num_nodes, num_rels, triplets):
      """ Create a DGL graph. The graph is bidirectional because RGCN authors
          use reversed relations.
          This function also generates edge type and normalization factor
          (reciprocal of node incoming degree)
      """
      g = dgl.DGLGraph()
      g.add_nodes(num_nodes)
      src, rel, dst = triplets
      src, dst = np.concatenate((src, dst)), np.concatenate((dst, src))
      rel = np.concatenate((rel, rel + num_rels))
      edges = sorted(zip(dst, src, rel))
      dst, src, rel = np.array(edges).transpose()
      g.add_edges(src, dst)
      norm = comp_deg_norm(g)
      logger.info("# nodes: {}, # edges: {}".format(num_nodes, len(src)))
      return g, rel.astype('int64'), norm  # .astype('int64') # This is buggy
  
  
  def build_test_graph(num_nodes, num_rels, edges):
      src, rel, dst = edges.transpose()
      logger.info("Test graph:")
      return build_graph_from_triplets(num_nodes, num_rels, (src, rel, dst))
  
  
  def negative_sampling(pos_samples, num_entity, negative_rate):
      size_of_batch = len(pos_samples)
      num_to_generate = size_of_batch * negative_rate
      neg_samples = np.tile(pos_samples, (negative_rate, 1))
      labels = np.zeros(size_of_batch * (negative_rate + 1), dtype=np.float32)
      labels[: size_of_batch] = 1
      values = np.random.randint(num_entity, size=num_to_generate)
      choices = np.random.uniform(size=num_to_generate)
      subj = choices > 0.5
      obj = choices <= 0.5
      neg_samples[subj, 0] = values[subj]
      neg_samples[obj, 2] = values[obj]
  
      return np.concatenate((pos_samples, neg_samples)), labels
  
  
  #######################################################################
  #
  # Utility functions for evaluations (filtered)
  #
  #######################################################################
  
  def filter_o(triplets_to_filter, target_s, target_r, target_o, num_entities,
               entity_filter=None):
      target_s, target_r, target_o = int(target_s), int(target_r), int(target_o)
      filtered_o = []
      # Do not filter out the test triplet, since we want to predict on it
      if (target_s, target_r, target_o) in triplets_to_filter:
          triplets_to_filter.remove((target_s, target_r, target_o))
      # Do not consider an object if it is part of a triplet to filter
      for o in range(num_entities):
          if (target_s, target_r, o) not in triplets_to_filter:
              filtered_o.append(o)
      # Do further filtering
      if entity_filter:
          filtered_o = list(filter(lambda idx: idx in entity_filter, filtered_o))
      return torch.LongTensor(filtered_o)
  
  
  def filter_s(triplets_to_filter, target_s, target_r, target_o, num_entities,
               entity_filter=None):
      target_s, target_r, target_o = int(target_s), int(target_r), int(target_o)
      filtered_s = []
      # Do not filter out the test triplet, since we want to predict on it
      if (target_s, target_r, target_o) in triplets_to_filter:
          triplets_to_filter.remove((target_s, target_r, target_o))
      # Do not consider a subject if it is part of a triplet to filter
      for s in range(num_entities):
          if (s, target_r, target_o) not in triplets_to_filter:
              filtered_s.append(s)
      # Do further filtering
      if entity_filter:
          filtered_s = list(filter(lambda idx: idx in entity_filter, filtered_s))
      return torch.LongTensor(filtered_s)
  
  
  def perturb_o_and_get_filtered_rank(embedding, decoder, target_triples,
                                      triplets_to_filter, entity_filter=None):
      """ Perturb object in the triplets
      """
      num_entities = embedding.shape[0]
      ranks = []
      s, r, o = target_triples[:, 0], target_triples[:, 1], target_triples[:, 2]
      for idx in tqdm(range(len(target_triples)), desc='perturbing object...'):
          target_s = s[idx]
          target_r = r[idx]
          target_o = o[idx]
          filtered_o = filter_o(triplets_to_filter, target_s,
                                target_r, target_o, num_entities,
                                entity_filter)
          target_o_idx = int((filtered_o == target_o).nonzero())
          scores = decoder.inference(embedding, target_s, target_r, filtered_o)
          _, indices = torch.sort(scores, descending=True)
          rank = int((indices == target_o_idx).nonzero())
          ranks.append(rank)
      ranks = torch.LongTensor(ranks) + 1
      return ranks
  
  
  def perturb_s_and_get_filtered_rank(embedding, decoder, target_triples,
                                      triplets_to_filter, entity_filter=None):
      """ Perturb subject in the triplets
      """
      num_entities = embedding.shape[0]
      ranks = []
      s, r, o = target_triples[:, 0], target_triples[:, 1], target_triples[:, 2]
      for idx in tqdm(range(len(target_triples)), desc='perturbing subject...'):
          target_s = s[idx]
          target_r = r[idx]
          target_o = o[idx]
          filtered_s = filter_s(triplets_to_filter, target_s,
                                target_r, target_o, num_entities,
                                entity_filter)
          target_s_idx = int((filtered_s == target_s).nonzero())
          scores = decoder.inference(embedding, filtered_s, target_r, target_o)
          _, indices = torch.sort(scores, descending=True)
          rank = int((indices == target_s_idx).nonzero())
          ranks.append(rank)
      ranks = torch.LongTensor(ranks) + 1
      return ranks
  
  
  def eval_filtered(embedding, decoder, train_triplets, valid_triplets, test_triplets,
                    hits=[], eval_type="valid", entity_filters=(None, None)):
      target_triples = None
      if eval_type == 'test':
          target_triples = test_triplets
      elif eval_type == 'valid':
          target_triples = valid_triplets
  
      triplets_to_filter = torch.cat(
          [train_triplets, valid_triplets, test_triplets]).tolist()
      triplets_to_filter = {tuple(triplet) for triplet in triplets_to_filter}
  
      logger.info('Perturbing subject...')
      ranks_s = perturb_s_and_get_filtered_rank(embedding, decoder, target_triples,
                                                triplets_to_filter,
                                                entity_filters[0])
      logger.info('Perturbing object...')
      ranks_o = perturb_o_and_get_filtered_rank(embedding, decoder, target_triples,
                                                triplets_to_filter,
                                                entity_filters[1])
  
      ranks = torch.cat([ranks_s, ranks_o])
      result = {'head': {}, 'tail': {}, 'all': {}}
  
      for rank_name, rank_data in {'head': ranks_s, 'tail': ranks_o, 'all': ranks}.items():
          logger.info('Eval results for ' + rank_name)
          mr = torch.mean(rank_data.float()).item()
          mrr = torch.mean(1.0 / rank_data.float()).item()
          logger.info("MR: {:.6f}".format(mr))
          logger.info("MRR: {:.6f}".format(mrr))
  
          result[rank_name] = {'mr': mr, 'mrr': mrr}
  
          for hit in hits:
              avg_count = torch.mean((rank_data <= hit).float()).item()
              logger.info(
                  "Hits @ {}: {:.6f}".format(hit, avg_count))
              result[rank_name]['hit_{}'.format(hit)] = avg_count
  
      return result
  
  
  #######################################################################
  #
  # Model util function
  #
  #######################################################################
  
  def load_model(model_state_file, graph_encoder, relation_decoder):
      if os.path.exists(model_state_file):
          try:
              model_dict = torch.load(model_state_file)
              graph_encoder.load_state_dict(model_dict['graph'])
              relation_decoder.load_state_dict(model_dict['relation'])
              epoch = model_dict['epoch']
              return epoch
          except Exception as e:
              logger.warning('Fails to load model because of ' + str(e))
              return 0
      else:
          return 0
  
  
  def save_model(model_state_file, graph_encoder, relation_decoder, epoch):
      model_folder = model_state_file[:model_state_file.rfind('/')]
      if not os.path.exists(model_folder):
          os.makedirs(model_folder)
      torch.save({
          'graph': graph_encoder.state_dict(),
          'relation': relation_decoder.state_dict(),
          'epoch': epoch
      }, model_state_file)
  