text: # -*- coding: utf-8 -*-
  #
  # Copyright 2017-2019 Data61, CSIRO
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  # http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  
  __all__ = ["train_val_test_split", "NodeSplitter"]
  
  import numpy as np
  import pandas as pd
  from stellargraph.core.graph import StellarGraphBase
  from stellargraph import globalvar
  
  
  # Easier functional interface for the splitter:
  def train_val_test_split(
      G,
      node_type=None,
      test_size=0.4,
      train_size=0.2,
      targets=None,
      split_equally=False,
      seed=None,
  ):
      """
          Splits node data into train, test, validation, and unlabeled sets.
  
          Any nodes that have a target value equal to globals.UNKNOWN_TARGET_ATTRIBUTE are added to the unlabeled set.
  
          The validation set includes all nodes that remain after the train, test and unlabeled sets have been
          created. As a result, it is possible the the validation set is empty.
  
      Args:
          G : StellarGraph containing the nodes to be split.
  
          test_size: float, int
              If float, should be between 0.0 and 1.0 and represent the proportion of the dataset to
              include in the test split. If int, represents the absolute number of test samples.
  
          train_size: float, int
              If float, should be between 0.0 and 1.0 and represent the proportion of the dataset to
              include in the train split. If int, represents the absolute number of train samples.
              If None, the value is automatically set to the complement of the test size.
  
          seed: int or None, optional (default=None)
              If this is an int the seed will be used to initialize a random number generator,
              otherwith the numpy default will be used.
  
          shuffle : boolean, optional (default=True)
              Whether or not to shuffle the data before splitting. If shuffle=False then stratify must be None.
  
          targets: None or DataFrame or dict, optional (default=None)
              If False the nodes are randomly assigned to each partition. If this is a Pandas DataFrame
              with node_ids as the index, or a dictionary with node_ids as keys and target as values,
              then these values will be used to find unlabelled nodes and if `split_equally` is True
              these target values will be used to sample equal numbers of each class.
  
          split_equally: bool (default=False)
              if `split_equally` is True the values passed into the targets argument will be used to
              sample equal numbers for the train split by class label.
  
      Returns:
              y_train, y_val, y_test, y_unlabeled
      """
      node_splitter = NodeSplitter()
  
      # Get list of nodes to split
      if node_type is None:
          nodes = list(G)
  
      elif isinstance(G, StellarGraphBase):
          nodes = G.nodes_of_type(node_type)
  
      else:
          raise TypeError("G must be a StellarGraph is node_type is not None")
  
      # Number of nodes and number without a label
      n_nodes = len(nodes)
  
      # Extract the target information
      if targets is not None:
          # TODO: The equal sampling option will fail if these values are not hashable.
          # Check that split_equally_by_target is the correct type
          if isinstance(targets, pd.DataFrame):
              target_values = [
                  targets.loc[n]
                  if n in targets.index
                  else globalvar.UNKNOWN_TARGET_ATTRIBUTE
                  for n in nodes
              ]
  
          elif isinstance(targets, dict):
              target_values = [
                  targets.get(n, globalvar.UNKNOWN_TARGET_ATTRIBUTE) for n in nodes
              ]
  
          else:
              raise TypeError(
                  "The targets are expected to be either a Pandas DataFrame or a dict."
              )
  
          n_known = sum(t != globalvar.UNKNOWN_TARGET_ATTRIBUTE for t in targets)
  
      else:
          n_known = n_nodes
          target_values = [0] * n_nodes
  
      if n_known == 0:
          raise RuntimeError("No nodes with target attribute to split.")
  
      # Find the number of nodes to use in the training set
      if isinstance(train_size, float) and (0 < train_size <= 1):
          train_size_n = int(n_known * train_size)
  
      elif isinstance(train_size, int):
          train_size_n = train_size
  
      else:
          raise ValueError("Splitter: train_size should be specified as a float or int")
  
      # Find the number of nodes to use in the test set
      if isinstance(test_size, float) and (0 < test_size <= 1):
          test_size_n = int(n_known * test_size)
  
      elif isinstance(test_size, int):
          test_size_n = test_size
  
      else:
          raise ValueError("Splitter: train_size should be specified as a float or int")
  
      # Find the number of nodes to use in the validation set
      val_size = None
      if isinstance(val_size, float) and (0 < val_size <= 1):
          val_size_n = int(n_known * val_size)
  
      elif isinstance(val_size, int):
          val_size_n = val_size
  
      else:
          val_size_n = max(0, n_known - (train_size_n + test_size_n))
  
      # Check that these sizes make sense
      if (train_size_n + test_size_n + val_size_n) > n_known:
          raise ValueError(
              "Number of train, test and val nodes "
              "is greater than the total number of labelled nodes."
          )
  
      # Now the splitter needs the node IDs and labels zipped together
      # TODO: This is a hack as the splitter only works when this array is sting type
      nodeid_and_label = np.array([nl for nl in enumerate(target_values)], dtype="U")
  
      # If stratified sampling, we need the target labels.
      if split_equally and targets is not None:
          class_set = set(target_values)
  
          # Remove the unknown target type
          class_set.discard(globalvar.UNKNOWN_TARGET_ATTRIBUTE)
  
          # The number of classes we have
          n_classes = len(class_set)
  
          if n_classes == 0:
              raise RuntimeError(
                  "Found no usable target classes in split_equally_by_targets."
              )
  
          if train_size_n < n_classes:
              raise RuntimeError(
                  "The number of classes must be smaller than the training size."
              )
  
          # The number of nodes we want per class
          p = int(train_size_n / n_classes)
  
          splits = node_splitter.train_test_split(
              y=nodeid_and_label, p=p, method="count", test_size=test_size_n, seed=seed
          )
  
      else:
          splits = node_splitter.train_test_split(
              y=nodeid_and_label,
              method="absolute",
              train_size=train_size_n,
              test_size=test_size_n,
              seed=seed,
          )
  
      # Get the node_ids out of the splitter
      node_ids_out = [[nodes[int(ind)] for ind in split[:, 0]] for split in splits]
  
      return node_ids_out
  
  
  class NodeSplitter(object):
      def __init__(self):
          self.format_epgm = False
          self.g_epgm = None
          self.g_id = None
          self._random = None
  
      def _get_nodes(self, graph_nodes, node_type, target_attribute):
          """
          Returns a list of node IDs for the subset of graph_nodes that have the given node type.
  
          Args:
              graph_nodes: <list> List of OrderedDict with vertex data for graph in EPGM format
              node_type: <str> The node type of interest
              target_attribute: <str> The target attribute key
  
          Returns:
              <list> List of node IDs that have given node type
  
          """
          # This code will fail if a node of node_type is missing the target_attribute.
          # We can fix this by using node['data'].get(target_attribute, None) so that at least all nodes of the
          # given type are returned. However, we must check for None in target_attribute later to exclude these nodes
          # from being added to train, test, and validation datasets.
          y = [
              (
                  node["id"],
                  node["data"].get(target_attribute, globalvar.UNKNOWN_TARGET_ATTRIBUTE),
              )
              for node in graph_nodes
              if node["meta"][globalvar.TYPE_ATTR_NAME] == node_type
          ]
  
          return y
  
      def _check_parameters(self, y, p, method, test_size, train_size, seed):
          """
          Checks that the parameters have valid values. It not, then it raises a ValueError exception with a
          message corresponding to the invalid parameter.
  
          Args:
              y: <numpy array> Array of size Nx2 containing node id, label columns.
              p: <int or float> Percent or count of the number of points for each class to sample.
              method: <str> One of 'count', 'percent', or 'absolute'.
              test_size: <int> number of points in the test set. For method 'count', it should be less than or equal to
              N - (np.unique(labels) * nc) where N is the number of labeled points in y.
              train_size: <int> The number of points in the train set only used by method 'absolute'.
              seed: <int> seed for random number generator, positive int or 0
  
          """
          if not isinstance(y, np.ndarray):
              raise ValueError("({}) y should be numpy array".format(type(self).__name__))
  
          if method != "count" and method != "percent" and method != "absolute":
              raise ValueError(
                  "({}) Valid methods are 'count', 'percent', and 'absolute' not {}".format(
                      type(self).__name__, method
                  )
              )
  
          if seed is not None:
              if seed < 0:
                  raise ValueError(
                      "({}) The random number generator seed value, seed, should be positive integer or None.".format(
                          type(self).__name__
                      )
                  )
              if not isinstance(seed, int):
                  raise ValueError(
                      "({}) The random number generator seed value, seed, should be integer type or None.".format(
                          type(self).__name__
                      )
                  )
  
          if method == "count":
              if not isinstance(p, int) or p <= 0:
                  raise ValueError(
                      "({}) p should be positive integer".format(type(self).__name__)
                  )
              if test_size is None or not isinstance(test_size, int) or test_size <= 0:
                  raise ValueError(
                      "({}) test_size must be positive integer".format(
                          type(self).__name__
                      )
                  )
  
          elif method == "percent":
              if not isinstance(p, float) or p < 0.0 or p > 1.0:
                  raise ValueError(
                      "({}) p should be float in the range [0,1].".format(
                          type(self).__name__
                      )
                  )
  
          elif method == "absolute":
              if test_size is None or not isinstance(test_size, int) or test_size <= 0:
                  raise ValueError(
                      "({}) test_size should be positive integer".format(
                          type(self).__name__
                      )
                  )
              if train_size is None or not isinstance(train_size, int) or train_size <= 0:
                  raise ValueError(
                      "({}) train_size should be positive integer".format(
                          type(self).__name__
                      )
                  )
  
      def train_test_split(
          self, y=None, p=10, method="count", test_size=None, train_size=None, seed=None
      ):
          """
          Splits node data into train, test, validation, and unlabeled sets.
  
          Any points in y that have value globals.UNKNOWN_TARGET_ATTRIBUTE are added to the unlabeled set.
  
          The validation set includes all the point that remain after the train, test and unlabeled sets have been
          created. As a result, it is possible the the validation set is empty, e.g., when method is set to 'percent'.
  
          The train, and test sets are build based on the specified method, 'count', 'percent', or 'absolute'.
  
          method='count': The value of parameter p specifies the number of points in the train set for each class. The
          test set size must be specified using the test_size parameter.
  
          method='percent': The value of parameter p specifies the train set size (and 1-p the test set size) as a
          percentage of the total number of points in y (including the unlabeled points.) The split is performed uniformly
          at random and the point labels (as specified in y) are not taken into account.
  
          method='absolute': The values of the parameters train_size and test_size specify the size of the train and test
          sets respectively. Points are selected uniformly at random and the label (as specified in y) are not taken into
          account.
  
          Args:
              y: <numpy array> Array of size Nx2 containing node id, label columns.
              p: <int or float> Percent or count of the number of points for each class to sample.
              method: <str> One of 'count', 'percent', or 'absolute'.
              test_size: <int> number of points in the test set. For method 'count', it should be less than or equal to
              N - (np.unique(labels) * nc) where N is the number of labeled points in y.
              train_size: <int> The number of points in the train set only used by method 'absolute'.
              seed: <int> seed for random number generator, positive int or 0
  
          Returns:
              y_train, y_val, y_test, y_unlabeled
          """
          self._check_parameters(
              y=y,
              p=p,
              method=method,
              test_size=test_size,
              train_size=train_size,
              seed=seed,
          )
  
          if self._random is None:
              self._random = np.random.RandomState(seed=seed)
  
          if method == "count":
              return self._split_data(y, p, test_size)
          elif method == "percent":
              n_unlabelled_points = np.sum(y[:, 1] == globalvar.UNKNOWN_TARGET_ATTRIBUTE)
              train_size = int((y.shape[0] - n_unlabelled_points) * p)
              test_size = y.shape[0] - n_unlabelled_points - train_size
              return self._split_data_absolute(
                  y=y, test_size=test_size, train_size=train_size
              )
          elif method == "absolute":
              return self._split_data_absolute(
                  y=y, test_size=test_size, train_size=train_size
              )
  
      def _split_data_absolute(self, y, test_size, train_size):
          """
          Splits given data such that the sizes of the test and train sets are fixed to the values given.
  
          Args:
              y: <numpy.ndarray> Array of size N x 2 containing node id + labels.
              test_size: <int> number of points in test set.
              train_size: <int> The number of points in the train set.
  
          Returns:
              y_train, y_val, y_test, y_unlabeled
  
          """
          # The label column in y could include None type, that is point with no ground truth label. These, if any,will
          # be returned separately in y_unlabeled dataset
          y_used = np.zeros(y.shape[0])  # initialize all the points are available
  
          # indexes of points with no class label:
          ind = np.nonzero(y[:, 1] == globalvar.UNKNOWN_TARGET_ATTRIBUTE)
          y_unlabeled = y[ind]
          y_used[ind] = 1
  
          ind = np.nonzero(y_used == 0)  # unused points
          ind_sampled = self._random.choice(ind[0], train_size, replace=False)
          y_train = y[ind_sampled]
          # mark these as used to make sure that they are not sampled for the test set
          y_used[ind_sampled] = 1
  
          # now sample test_size points for the test set
          ind = np.nonzero(y_used == 0)  # indexes of points that are not in training set
          if len(ind[0]) < test_size:
              raise Exception(
                  "Not enough nodes available for the test set: available {} nodes, needed {}. Aborting".format(
                      len(ind[0]), test_size
                  )
              )
          ind_sampled = self._random.choice(ind[0], test_size, replace=False)
          y_test = y[ind_sampled]
          y_used[ind_sampled] = 1
          # print("y_test shape: ", y_test.shape)
  
          # Validation set
          # the remaining labeled points (if any) go into the validation set
          ind = np.nonzero(y_used == 0)
          y_val = y[ind[0]]
          # print("y_val shape:", y_val.shape)
  
          return y_train, y_val, y_test, y_unlabeled
  
      def _split_data(self, y, nc, test_size):
          """
          Splits the data according to the scheme in Yang et al, ICML 2016, Revisiting semi-supervised learning with graph
          embeddings.
  
          Args:
              y: <numpy.ndarray> Array of size N x 2 containing node id + labels.
              nc: <int> number of points from each class in train set.
              test_size: <int> number of points in test set; it should be less than or equal to
              N - (np.unique(labels) * nc).
  
          Returns:
              y_train, y_val, y_test, y_unlabeled
          """
          y_used = np.zeros(y.shape[0])  # initialize all the points are available
  
          # indexes of points with no class label:
          ind = np.nonzero(y[:, 1] == globalvar.UNKNOWN_TARGET_ATTRIBUTE)
          y_unlabeled = y[ind]
          y_used[ind] = 1
  
          y_train = None
  
          class_labels = np.unique(y[:, 1])
          ind = class_labels == globalvar.UNKNOWN_TARGET_ATTRIBUTE
          class_labels = class_labels[np.logical_not(ind)]
  
          if test_size > y.shape[0] - class_labels.size * nc:
              # re-adjust so that none of the training samples end up in the test set
              test_size = y.shape[0] - class_labels.size * nc
  
          for clabel in class_labels:
              # indexes of points with class label clabel:
              ind = np.nonzero(y[:, 1] == clabel)
              # select nc of these at random for the training set
              if ind[0].size <= nc:
                  # too few labeled examples for class so use half for training and half for testing
                  ind_selected = self._random.choice(
                      ind[0], ind[0].size // 2, replace=False
                  )
              else:
                  ind_selected = self._random.choice(ind[0], nc, replace=False)
              # mark these as used to make sure that they are not sampled for the test set:
              y_used[ind_selected] = 1
              if y_train is None:
                  y_train = y[ind_selected]
              else:
                  # print("y_train shape:", y_train.shape)
                  # print("y[ind_selected] shape:", y[ind_selected].shape)
                  y_train = np.vstack((y_train, y[ind_selected]))
  
          # now sample test_size points for the test set
          ind = np.nonzero(y_used == 0)  # indexes of points that are not in training set
          if len(ind[0]) < test_size:
              raise Exception(
                  "Not enough nodes available for the test set: available {} nodes, needed {}. Aborting".format(
                      len(ind[0]), test_size
                  )
              )
          ind_selected = self._random.choice(ind[0], test_size, replace=False)
          y_test = y[ind_selected]
          y_used[ind_selected] = 1
          # print("y_test shape: ", y_test.shape)
  
          # the remaining points (if any) go into the validation set
          ind = np.nonzero(y_used == 0)
          y_val = y[ind[0]]
          # print("y_val shape:", y_val.shape)
  
          return y_train, y_val, y_test, y_unlabeled
  