text: """RDF datasets
  
  Datasets from "A Collection of Benchmark Datasets for
  Systematic Evaluations of Machine Learning on
  the Semantic Web"
  """
  import os
  from collections import OrderedDict
  import itertools
  import rdflib as rdf
  import abc
  import re
  
  import networkx as nx
  import numpy as np
  
  import dgl
  import dgl.backend as F
  from .utils import download, extract_archive, get_download_dir, _get_dgl_url
  
  __all__ = ['AIFB', 'MUTAG', 'BGS', 'AM']
  
  class Entity:
      """Class for entities
  
      Parameters
      ----------
      id : str
          ID of this entity
      cls : str
          Type of this entity
      """
      def __init__(self, id, cls):
          self.id = id
          self.cls = cls
  
      def __str__(self):
          return '{}/{}'.format(self.cls, self.id)
  
  class Relation:
      """Class for relations
  
      Parameters
      ----------
      cls : str
          Type of this relation
      """
      def __init__(self, cls):
          self.cls = cls
  
      def __str__(self):
          return str(self.cls)
  
  class RDFGraphDataset:
      """Base graph dataset class from RDF tuples.
  
      To derive from this, implement the following abstract methods:
      * ``parse_entity``
      * ``parse_relation``
      * ``process_tuple``
      * ``process_idx_file_line``
      * ``predict_category``
  
      Preprocessed graph and other data will be cached in the download folder
      to speedup data loading.
  
      The dataset should contain a "trainingSet.tsv" and a "testSet.tsv" file
      for training and testing samples.
  
      Attributes
      ----------
      graph : dgl.DGLHeteroGraph
          Graph structure
      num_classes : int
          Number of classes to predict
      predict_category : str
          The entity category (node type) that has labels for prediction
      train_idx : Tensor
          Entity IDs for training. All IDs are local IDs w.r.t. to ``predict_category``.
      test_idx : Tensor
          Entity IDs for testing. All IDs are local IDs w.r.t. to ``predict_category``.
      labels : Tensor
          All the labels of the entities in ``predict_category``
  
      Parameters
      ----------
      url : str or path
          URL to download the raw dataset.
      name : str
          Name of the dataset
      force_reload : bool, optional
          If true, force load and process from raw data. Ignore cached pre-processed data.
      print_every : int, optional
          Log for every X tuples.
      insert_reverse : bool, optional
          If true, add reverse edge and reverse relations to the final graph.
      """
      def __init__(self, url, name,
                   force_reload=False,
                   print_every=10000,
                   insert_reverse=True):
          download_dir = get_download_dir()
          zip_file_path = os.path.join(download_dir, '{}.zip'.format(name))
          download(url, path=zip_file_path)
          self._dir = os.path.join(download_dir, name)
          extract_archive(zip_file_path, self._dir)
          self._print_every = print_every
          self._insert_reverse = insert_reverse
          if not force_reload and self.has_cache():
              print('Found cached graph. Load cache ...')
              self.load_cache()
          else:
              raw_tuples = self.load_raw_tuples()
              self.process_raw_tuples(raw_tuples)
          print('#Training samples:', len(self.train_idx))
          print('#Testing samples:', len(self.test_idx))
          print('#Classes:', self.num_classes)
          print('Predict category:', self.predict_category)
  
      def load_raw_tuples(self):
          raw_rdf_graphs = []
          for i, filename in enumerate(os.listdir(self._dir)):
              fmt = None
              if filename.endswith('nt'):
                  fmt = 'nt'
              elif filename.endswith('n3'):
                  fmt = 'n3'
              if fmt is None:
                  continue
              g = rdf.Graph()
              print('Parsing file %s ...' % filename)
              g.parse(os.path.join(self._dir, filename), format=fmt)
              raw_rdf_graphs.append(g)
          return itertools.chain(*raw_rdf_graphs)
  
      def process_raw_tuples(self, raw_tuples):
          mg = nx.MultiDiGraph()
          ent_classes = OrderedDict()
          rel_classes = OrderedDict()
          entities = OrderedDict()
          src = []
          dst = []
          ntid = []
          etid = []
          for i, (sbj, pred, obj) in enumerate(raw_tuples):
              if i % self._print_every == 0:
                  print('Processed %d tuples, found %d valid tuples.' % (i, len(src)))
              sbjent = self.parse_entity(sbj)
              rel = self.parse_relation(pred)
              objent = self.parse_entity(obj)
              processed = self.process_tuple((sbj, pred, obj), sbjent, rel, objent)
              if processed is None:
                  # ignored
                  continue
              # meta graph
              sbjclsid = _get_id(ent_classes, sbjent.cls)
              objclsid = _get_id(ent_classes, objent.cls)
              relclsid = _get_id(rel_classes, rel.cls)
              mg.add_edge(sbjent.cls, objent.cls, key=rel.cls)
              if self._insert_reverse:
                  mg.add_edge(objent.cls, sbjent.cls, key='rev-%s' % rel.cls)
              # instance graph
              src_id = _get_id(entities, str(sbjent))
              if len(entities) > len(ntid):  # found new entity
                  ntid.append(sbjclsid)
              dst_id = _get_id(entities, str(objent))
              if len(entities) > len(ntid):  # found new entity
                  ntid.append(objclsid)
              src.append(src_id)
              dst.append(dst_id)
              etid.append(relclsid)
  
          src = np.array(src)
          dst = np.array(dst)
          ntid = np.array(ntid)
          etid = np.array(etid)
          ntypes = list(ent_classes.keys())
          etypes = list(rel_classes.keys())
  
          # add reverse edge with reverse relation
          if self._insert_reverse:
              print('Adding reverse edges ...')
              newsrc = np.hstack([src, dst])
              newdst = np.hstack([dst, src])
              src = newsrc
              dst = newdst
              etid = np.hstack([etid, etid + len(etypes)])
              etypes.extend(['rev-%s' % t for t in etypes])
  
          self.build_graph(mg, src, dst, ntid, etid, ntypes, etypes)
  
          print('Load training/validation/testing split ...')
          idmap = F.asnumpy(self.graph.nodes[self.predict_category].data[dgl.NID])
          glb2lcl = {glbid : lclid for lclid, glbid in enumerate(idmap)}
          def findidfn(ent):
              if ent not in entities:
                  return None
              else:
                  return glb2lcl[entities[ent]]
          self.load_data_split(findidfn)
  
          self.save_cache(mg, src, dst, ntid, etid, ntypes, etypes)
  
      def build_graph(self, mg, src, dst, ntid, etid, ntypes, etypes):
          # create homo graph
          print('Creating one whole graph ...')
          g = dgl.graph((src, dst))
          g.ndata[dgl.NTYPE] = F.tensor(ntid)
          g.edata[dgl.ETYPE] = F.tensor(etid)
          print('Total #nodes:', g.number_of_nodes())
          print('Total #edges:', g.number_of_edges())
  
          # convert to heterograph
          print('Convert to heterograph ...')
          hg = dgl.to_hetero(g,
                             ntypes,
                             etypes,
                             metagraph=mg)
          print('#Node types:', len(hg.ntypes))
          print('#Canonical edge types:', len(hg.etypes))
          print('#Unique edge type names:', len(set(hg.etypes)))
          self.graph = hg
  
      def save_cache(self, mg, src, dst, ntid, etid, ntypes, etypes):
          nx.write_gpickle(mg, os.path.join(self._dir, 'cached_mg.gpickle'))
          np.save(os.path.join(self._dir, 'cached_src.npy'), src)
          np.save(os.path.join(self._dir, 'cached_dst.npy'), dst)
          np.save(os.path.join(self._dir, 'cached_ntid.npy'), ntid)
          np.save(os.path.join(self._dir, 'cached_etid.npy'), etid)
          save_strlist(os.path.join(self._dir, 'cached_ntypes.txt'), ntypes)
          save_strlist(os.path.join(self._dir, 'cached_etypes.txt'), etypes)
          np.save(os.path.join(self._dir, 'cached_train_idx.npy'), F.asnumpy(self.train_idx))
          np.save(os.path.join(self._dir, 'cached_test_idx.npy'), F.asnumpy(self.test_idx))
          np.save(os.path.join(self._dir, 'cached_labels.npy'), F.asnumpy(self.labels))
  
      def has_cache(self):
          return (os.path.exists(os.path.join(self._dir, 'cached_mg.gpickle'))
              and os.path.exists(os.path.join(self._dir, 'cached_src.npy'))
              and os.path.exists(os.path.join(self._dir, 'cached_dst.npy'))
              and os.path.exists(os.path.join(self._dir, 'cached_ntid.npy'))
              and os.path.exists(os.path.join(self._dir, 'cached_etid.npy'))
              and os.path.exists(os.path.join(self._dir, 'cached_ntypes.txt'))
              and os.path.exists(os.path.join(self._dir, 'cached_etypes.txt'))
              and os.path.exists(os.path.join(self._dir, 'cached_train_idx.npy'))
              and os.path.exists(os.path.join(self._dir, 'cached_test_idx.npy'))
              and os.path.exists(os.path.join(self._dir, 'cached_labels.npy')))
  
      def load_cache(self):
          mg = nx.read_gpickle(os.path.join(self._dir, 'cached_mg.gpickle'))
          src = np.load(os.path.join(self._dir, 'cached_src.npy'))
          dst = np.load(os.path.join(self._dir, 'cached_dst.npy'))
          ntid = np.load(os.path.join(self._dir, 'cached_ntid.npy'))
          etid = np.load(os.path.join(self._dir, 'cached_etid.npy'))
          ntypes = load_strlist(os.path.join(self._dir, 'cached_ntypes.txt'))
          etypes = load_strlist(os.path.join(self._dir, 'cached_etypes.txt'))
          self.train_idx = F.tensor(np.load(os.path.join(self._dir, 'cached_train_idx.npy')))
          self.test_idx = F.tensor(np.load(os.path.join(self._dir, 'cached_test_idx.npy')))
          labels = np.load(os.path.join(self._dir, 'cached_labels.npy'))
          self.num_classes = labels.max() + 1
          self.labels = F.tensor(labels)
  
          self.build_graph(mg, src, dst, ntid, etid, ntypes, etypes)
  
      def load_data_split(self, ent2id):
          label_dict = {}
          labels = np.zeros((self.graph.number_of_nodes(self.predict_category),)) - 1
          train_idx = self.parse_idx_file(
              os.path.join(self._dir, 'trainingSet.tsv'),
              ent2id, label_dict, labels)
          test_idx = self.parse_idx_file(
              os.path.join(self._dir, 'testSet.tsv'),
              ent2id, label_dict, labels)
          self.train_idx = F.tensor(train_idx)
          self.test_idx = F.tensor(test_idx)
          self.labels = F.tensor(labels).long()
          self.num_classes = len(label_dict)
  
      def parse_idx_file(self, filename, ent2id, label_dict, labels):
          idx = []
          with open(filename, 'r') as f:
              for i, line in enumerate(f):
                  if i == 0:
                      continue  # first line is the header
                  sample, label = self.process_idx_file_line(line)
                  #person, _, label = line.strip().split('\t')
                  ent = self.parse_entity(sample)
                  entid = ent2id(str(ent))
                  if entid is None:
                      print('Warning: entity "%s" does not have any valid links associated. Ignored.' % str(ent))
                  else:
                      idx.append(entid)
                      lblid = _get_id(label_dict, label)
                      labels[entid] = lblid
          return idx
  
      @abc.abstractmethod
      def parse_entity(self, term):
          """Parse one entity from an RDF term.
  
          Return None if the term does not represent a valid entity and the
          whole tuple should be ignored.
  
          Parameters
          ----------
          term : rdflib.term.Identifier
              RDF term
  
          Returns
          -------
          Entity or None
              An entity.
          """
          pass
  
      @abc.abstractmethod
      def parse_relation(self, term):
          """Parse one relation from an RDF term.
  
          Return None if the term does not represent a valid relation and the
          whole tuple should be ignored.
  
          Parameters
          ----------
          term : rdflib.term.Identifier
              RDF term
  
          Returns
          -------
          Relation or None
              A relation
          """
          pass
  
      @abc.abstractmethod
      def process_tuple(self, raw_tuple, sbj, rel, obj):
          """Process the tuple.
  
          Return (Entity, Relation, Entity) tuple for as the final tuple.
          Return None if the tuple should be ignored.
          
          Parameters
          ----------
          raw_tuple : tuple of rdflib.term.Identifier
              (subject, predicate, object) tuple
          sbj : Entity
              Subject entity
          rel : Relation
              Relation
          obj : Entity
              Object entity
  
          Returns
          -------
          (Entity, Relation, Entity)
              The final tuple or None if should be ignored
          """
          pass
  
      @abc.abstractmethod
      def process_idx_file_line(self, line):
          """Process one line of ``trainingSet.tsv`` or ``testSet.tsv``.
  
          Parameters
          ----------
          line : str
              One line of the file
  
          Returns
          -------
          (str, str)
              One sample and its label
          """
          pass
  
      @property
      @abc.abstractmethod
      def predict_category(self):
          """Return the category name that has labels."""
          pass
  
  def _get_id(dict, key):
      id = dict.get(key, None)
      if id is None:
          id = len(dict)
          dict[key] = id
      return id
  
  def save_strlist(filename, strlist):
      with open(filename, 'w') as f:
          for s in strlist:
              f.write(s + '\n')
  
  def load_strlist(filename):
      with open(filename, 'r') as f:
          ret = []
          for line in f:
              ret.append(line.strip())
          return ret
  
  class AIFB(RDFGraphDataset):
      """AIFB dataset.
      
      Examples
      --------
      >>> dataset = dgl.data.rdf.AIFB()
      >>> print(dataset.graph)
      """
  
      employs = rdf.term.URIRef("http://swrc.ontoware.org/ontology#employs")
      affiliation = rdf.term.URIRef("http://swrc.ontoware.org/ontology#affiliation")
      entity_prefix = 'http://www.aifb.uni-karlsruhe.de/'
      relation_prefix = 'http://swrc.ontoware.org/'
  
      def __init__(self,
                   force_reload=False,
                   print_every=10000,
                   insert_reverse=True):
          url = _get_dgl_url('dataset/rdf/aifb-hetero.zip')
          name = 'aifb-hetero'
          super(AIFB, self).__init__(url, name,
                                     force_reload=force_reload,
                                     print_every=print_every,
                                     insert_reverse=insert_reverse)
  
      def parse_entity(self, term):
          if isinstance(term, rdf.Literal):
              return Entity(id=str(term), cls="_Literal")
          if isinstance(term, rdf.BNode):
              return None
          entstr = str(term)
          if entstr.startswith(self.entity_prefix):
              sp = entstr.split('/')
              return Entity(id=sp[5], cls=sp[3])
          else:
              return None
  
      def parse_relation(self, term):
          if term == self.employs or term == self.affiliation:
              return None
          relstr = str(term)
          if relstr.startswith(self.relation_prefix):
              return Relation(cls=relstr.split('/')[3])
          else:
              relstr = relstr.split('/')[-1]
              return Relation(cls=relstr)
  
      def process_tuple(self, raw_tuple, sbj, rel, obj):
          if sbj is None or rel is None or obj is None:
              return None
          return (sbj, rel, obj)
  
      def process_idx_file_line(self, line):
          person, _, label = line.strip().split('\t')
          return person, label
  
      @property
      def predict_category(self):
          return 'Personen'
  
  class MUTAG(RDFGraphDataset):
      """MUTAG dataset.
      
      Examples
      --------
      >>> dataset = dgl.data.rdf.MUTAG()
      >>> print(dataset.graph)
      """
  
      d_entity = re.compile("d[0-9]")
      bond_entity = re.compile("bond[0-9]")
  
      is_mutagenic = rdf.term.URIRef("http://dl-learner.org/carcinogenesis#isMutagenic")
      rdf_type = rdf.term.URIRef("http://www.w3.org/1999/02/22-rdf-syntax-ns#type")
      rdf_subclassof = rdf.term.URIRef("http://www.w3.org/2000/01/rdf-schema#subClassOf")
      rdf_domain = rdf.term.URIRef("http://www.w3.org/2000/01/rdf-schema#domain")
  
      entity_prefix = 'http://dl-learner.org/carcinogenesis#'
      relation_prefix = entity_prefix
  
      def __init__(self,
                   force_reload=False,
                   print_every=10000,
                   insert_reverse=True):
          url = _get_dgl_url('dataset/rdf/mutag-hetero.zip')
          name = 'mutag-hetero'
          super(MUTAG, self).__init__(url, name,
                                     force_reload=force_reload,
                                     print_every=print_every,
                                     insert_reverse=insert_reverse)
  
      def parse_entity(self, term):
          if isinstance(term, rdf.Literal):
              return Entity(id=str(term), cls="_Literal")
          elif isinstance(term, rdf.BNode):
              return None
          entstr = str(term)
          if entstr.startswith(self.entity_prefix):
              inst = entstr[len(self.entity_prefix):]
              if self.d_entity.match(inst):
                  cls = 'd'
              elif self.bond_entity.match(inst):
                  cls = 'bond'
              else:
                  cls = None
              return Entity(id=inst, cls=cls)
          else:
              return None
  
      def parse_relation(self, term):
          if term == self.is_mutagenic:
              return None
          relstr = str(term)
          if relstr.startswith(self.relation_prefix):
              cls = relstr[len(self.relation_prefix):]
              return Relation(cls=cls)
          else:
              relstr = relstr.split('/')[-1]
              return Relation(cls=relstr)
  
      def process_tuple(self, raw_tuple, sbj, rel, obj):
          if sbj is None or rel is None or obj is None:
              return None
  
          if not raw_tuple[1].startswith('http://dl-learner.org/carcinogenesis#'):
              obj.cls = 'SCHEMA'
              if sbj.cls is None:
                  sbj.cls = 'SCHEMA'
          if obj.cls is None:
              obj.cls = rel.cls
  
          assert sbj.cls is not None and obj.cls is not None
          
          return (sbj, rel, obj)
  
      def process_idx_file_line(self, line):
          bond, _, label = line.strip().split('\t')
          return bond, label
  
      @property
      def predict_category(self):
          return 'd'
  
  class BGS(RDFGraphDataset):
      """BGS dataset.
  
      BGS namespace convention:
      http://data.bgs.ac.uk/(ref|id)/<Major Concept>/<Sub Concept>/INSTANCE
  
      We ignored all literal nodes and the relations connecting them in the
      output graph. We also ignored the relation used to mark whether a
      term is CURRENT or DEPRECATED.
  
      Examples
      --------
      >>> dataset = dgl.data.rdf.BGS()
      >>> print(dataset.graph)
      """
  
      lith = rdf.term.URIRef("http://data.bgs.ac.uk/ref/Lexicon/hasLithogenesis")
      entity_prefix = 'http://data.bgs.ac.uk/'
      status_prefix = 'http://data.bgs.ac.uk/ref/CurrentStatus'
      relation_prefix = 'http://data.bgs.ac.uk/ref'
  
      def __init__(self,
                   force_reload=False,
                   print_every=10000,
                   insert_reverse=True):
          url = _get_dgl_url('dataset/rdf/bgs-hetero.zip')
          name = 'bgs-hetero'
          super(BGS, self).__init__(url, name,
                                     force_reload=force_reload,
                                     print_every=print_every,
                                     insert_reverse=insert_reverse)
  
      def parse_entity(self, term):
          if isinstance(term, rdf.Literal):
              return None
          elif isinstance(term, rdf.BNode):
              return None
          entstr = str(term)
          if entstr.startswith(self.status_prefix):
              return None
          if entstr.startswith(self.entity_prefix):
              sp = entstr.split('/')
              if len(sp) != 7:
                  return None
              # instance
              cls = '%s/%s' % (sp[4], sp[5])
              inst = sp[6]
              return Entity(id=inst, cls=cls)
          else:
              return None
  
      def parse_relation(self, term):
          if term == self.lith:
              return None
          relstr = str(term)
          if relstr.startswith(self.relation_prefix):
              sp = relstr.split('/')
              if len(sp) < 6:
                  return None
              assert len(sp) == 6, relstr
              cls = '%s/%s' % (sp[4], sp[5])
              return Relation(cls=cls)
          else:
              relstr = relstr.replace('.', '_')
              return Relation(cls=relstr)
  
      def process_tuple(self, raw_tuple, sbj, rel, obj):
          if sbj is None or rel is None or obj is None:
              return None
          return (sbj, rel, obj)
  
      def process_idx_file_line(self, line):
          _, rock, label = line.strip().split('\t')
          return rock, label
  
      @property
      def predict_category(self):
          return 'Lexicon/NamedRockUnit'
  
  class AM(RDFGraphDataset):
      """AM dataset.
  
      Namespace convention:
      Instance: http://purl.org/collections/nl/am/<type>-<id>
      Relation: http://purl.org/collections/nl/am/<name>
  
      We ignored all literal nodes and the relations connecting them in the
      output graph.
  
      Examples
      --------
      >>> dataset = dgl.data.rdf.AM()
      >>> print(dataset.graph)
      """
  
      objectCategory = rdf.term.URIRef("http://purl.org/collections/nl/am/objectCategory")
      material = rdf.term.URIRef("http://purl.org/collections/nl/am/material")
      entity_prefix = 'http://purl.org/collections/nl/am/'
      relation_prefix = entity_prefix
  
      def __init__(self,
                   force_reload=False,
                   print_every=10000,
                   insert_reverse=True):
          url = _get_dgl_url('dataset/rdf/am-hetero.zip')
          name = 'am-hetero'
          super(AM, self).__init__(url, name,
                                     force_reload=force_reload,
                                     print_every=print_every,
                                     insert_reverse=insert_reverse)
  
      def parse_entity(self, term):
          if isinstance(term, rdf.Literal):
              return None
          elif isinstance(term, rdf.BNode):
              return Entity(id=str(term), cls='_BNode')
          entstr = str(term)
          if entstr.startswith(self.entity_prefix):
              sp = entstr.split('/')
              assert len(sp) == 7, entstr
              spp = sp[6].split('-')
              if len(spp) == 2:
                  # instance
                  cls, inst = spp
              else:
                  cls = 'TYPE'
                  inst = spp
              return Entity(id=inst, cls=cls)
          else:
              return None
  
      def parse_relation(self, term):
          if term == self.objectCategory or term == self.material:
              return None
          relstr = str(term)
          if relstr.startswith(self.relation_prefix):
              sp = relstr.split('/')
              assert len(sp) == 7, relstr
              cls = sp[6]
              return Relation(cls=cls)
          else:
              relstr = relstr.replace('.', '_')
              return Relation(cls=relstr)
  
      def process_tuple(self, raw_tuple, sbj, rel, obj):
          if sbj is None or rel is None or obj is None:
              return None
          return (sbj, rel, obj)
  
      def process_idx_file_line(self, line):
          proxy, _, label = line.strip().split('\t')
          return proxy, label
  
      @property
      def predict_category(self):
          return 'proxy'
  
  if __name__ == '__main__':
      AIFB()
  