text: """Abstract classes and utility operations for building graph representations and
  data loaders (known as Sequence objects in Keras).
  
  Most users will not need to interact with this module."""
  from abc import ABCMeta, abstractmethod
  from operator import itemgetter
  import numpy as np
  from amdnet.utils.general import expand_1st, to_list
  from monty.json import MSONable
  from amdnet.data import local_env
  from inspect import signature
  from pymatgen.analysis.local_env import NearNeighbors
  from tensorflow.keras.utils import Sequence
  
  
  class StructureGraph(MSONable):
      """
      This is a base class for converting converting structure into graphs or model inputs
      Methods to be implemented are follows:
          1. convert(self, structure)
              This is to convert a structure into a graph dictionary
          2. get_input(self, structure)
              This method convert a structure directly to a model input
          3. get_flat_data(self, graphs, targets)
              This method process graphs and targets pairs and output model input list.
  
      """
  
      # TODO (wardlt): Consider making "num_*_features" funcs to simplify making a MEGNet model
  
      def __init__(self,
                   nn_strategy=None,
                   atom_converter=None,
                   bond_converter=None,
                   **kwargs):
  
          if isinstance(nn_strategy, str):
              strategy = local_env.get(nn_strategy)
              parameters = signature(strategy).parameters
              param_dict = {i: j.default for i, j in parameters.items()}
              for i, j in kwargs.items():
                  if i in param_dict:
                      setattr(self, i, j)
                      param_dict.update({i: j})
              self.nn_strategy = strategy(**param_dict)
          elif isinstance(nn_strategy, NearNeighbors):
              self.nn_strategy = nn_strategy
          elif nn_strategy is None:
              self.nn_strategy = None
          else:
              raise RuntimeError("Strategy not valid")
  
          self.atom_converter = atom_converter
          self.bond_converter = bond_converter
          if self.atom_converter is None:
              self.atom_converter = self._get_dummy_converter()
          if self.bond_converter is None:
              self.bond_converter = self._get_dummy_converter()
  
      def convert(self, structure, state_attributes=None):
          """
          Take a pymatgen structure and convert it to a index-type graph representation
          The graph will have node, distance, index1, index2, where node is a vector of Z number
          of atoms in the structure, index1 and index2 mark the atom indices forming the bond and separated by
          distance.
          For state attributes, you can set structure.state = [[xx, xx]] beforehand or the algorithm would
          take default [[0, 0]]
  
          Args:
              state_attributes: (list) state attributes
              structure: (pymatgen structure)
              (dictionary)
          """
          state_attributes = state_attributes or [[0, 0]]
          index1 = []
          index2 = []
          bonds = []
          if self.nn_strategy is None:
              raise RuntimeError("NearNeighbor strategy is not provided!")
          for n, neighbors in enumerate(self.nn_strategy.get_all_nn_info(structure)):
              index1.extend([n] * len(neighbors))
              for neighbor in neighbors:
                  index2.append(neighbor['site_index'])
                  bonds.append(neighbor['weight'])
  
          atoms = [i.specie.Z for i in structure]
  
          if np.size(np.unique(index1)) < len(atoms):
              raise RuntimeError("Isolated atoms found in the structure")
          else:
              return {'atom': np.array(atoms, dtype='int32').tolist(),
                      'bond': bonds,
                      'state': state_attributes,
                      'index1': index1,
                      'index2': index2
                      }
  
      def __call__(self, structure):
          return self.convert(structure)
  
      def get_input(self, structure):
          """
          Turns a structure into model input
          """
          graph = self.convert(structure)
          return self.graph_to_input(graph)
  
      def graph_to_input(self, graph):
          """
          Turns a graph into model input
  
          Args:
              (dict): Dictionary description of the graph
          Return:
              ([np.ndarray]): Inputs in the form needed by MEGNet
          """
          gnode = [0] * len(graph['atom'])
          gbond = [0] * len(graph['index1'])
  
          ret = [expand_1st(self.atom_converter.convert(graph['atom'])),
                  expand_1st(self.bond_converter.convert(graph['bond'])),
                  expand_1st(np.array(graph['state'])),
                  expand_1st(np.array(graph['index1'])),
                  expand_1st(np.array(graph['index2'])),
                  expand_1st(np.array(gnode)),
                  expand_1st(np.array(gbond))]
          ret = [np.array(i_, dtype=np.float) for i_ in ret]
          
          return ret
  
      def get_flat_data(self, graphs, targets=None):
          """
          Expand the graph dictionary to form a list of features and targets tensors.
          This is useful when the model is trained on assembled graphs on the fly.
  
          Args:
              graphs: (list of dictionary) list of graph dictionary for each structure
              targets: (list of float or list) Optional: corresponding target
                  values for each structure
  
          Returns:
              tuple(node_features, edges_features, global_values, index1, index2, targets)
          """
  
          output = []  # Will be a list of arrays
  
          # Convert the graphs to matrices
          for feature in ['atom', 'bond', 'state', 'index1', 'index2']:
              output.append([np.array(x[feature], dtype=np.float) for x in graphs])
  
          # If needed, add the targets
          if targets is not None:
              output.append([to_list(t) for t in targets])
          
          return tuple(output)
  
      @staticmethod
      def _get_dummy_converter():
          return DummyConverter()
  
      def as_dict(self):
          all_dict = super().as_dict()
          if 'nn_strategy' in all_dict:
              nn_strategy = all_dict.pop('nn_strategy')
              all_dict.update({'nn_strategy': local_env.serialize(nn_strategy)})
          return all_dict
  
      @classmethod
      def from_dict(cls, d):
          if 'nn_strategy' in d:
              nn_strategy = d.pop('nn_strategy')
              nn_strategy_obj = local_env.deserialize(nn_strategy)
              d.update({'nn_strategy': nn_strategy_obj})
              return super().from_dict(d)
          return super().from_dict(d)
  
  
  class DistanceConverter(MSONable):
      """
      Base class for distance conversion. The class needs to have a convert method.
      """
  
      def convert(self, d):
          raise NotImplementedError
  
  
  class DummyConverter(DistanceConverter):
      """
      Dummy converter as a placeholder
      """
  
      def convert(self, d):
          return d
  
  
  class GaussianDistance(DistanceConverter):
      """
      Expand distance with Gaussian basis sit at centers and with width 0.5.
  
      Args:
          centers: (np.array)
          width: (float)
      """
  
      def __init__(self, centers=np.linspace(0, 5, 100), width=0.5):
          self.centers = centers
          self.width = width
  
      def convert(self, d):
          """
          expand distance vector d with given parameters
  
          Args:
              d: (1d array) distance array
  
          Returns
              (matrix) N*M matrix with N the length of d and M the length of centers
          """
          d = np.array(d)
          return np.exp(-(d[:, None] - self.centers[None, :]) ** 2 / self.width ** 2)
  
  
  class MoorseLongRange(DistanceConverter):
      """
      This is an attempt to implement a Moorse/long range interactomic potential like
      distance expansion. The distance is expanded with this basis at different equilibrium
      bond distance, r_eq. It is still a work in progress. Do not use if you do not know
      much about the parameters
      ref: https://en.wikipedia.org/wiki/Morse/Long-range_potential#Function
  
      Args:
          d_e: (float) dissociate energy
          r_ref: (float) reference bond length
          r_eq: (list) equilibrium bond length
          p: (int) exponential term in the original equation, see ref
          q: (int) exponential term in the original equaiton, see ref
          cm: (list) long range coefficients in u_LR = \Sigma_i_N (cm_i / r^i)
          betas: (list) parameters determining the transition between long range and short range
      """
  
      def __init__(self, d_e=1, r_ref=2, r_eq=[1, 2, 3],
                   p=2, q=2, cm=[1, 2, 3, 4],
                   betas=[0.1, 0.2, 0.3, 0.4]):
          self.d_e = d_e
          self.r_ref = r_ref
          self.r_eq = np.array(r_eq)
          self.p = p
          self.q = q
          self.cm = np.array(cm).ravel()
          self.n_cm = len(self.cm)
          self.betas = np.array(betas).ravel()
  
      def convert(self, d):
          return self.d_e * (1 - self.u(d)[:, None] / self.u(self.r_eq)[None, :] *
                             np.exp(-self.beta(d) * self.y(d[:, None], self.r_eq[None, :], self.p))) ** 2
  
      def u(self, r):
          m_i = np.arange(1, self.n_cm + 1)
          if np.size(r) == 1:
              return np.sum(self.cm / r ** m_i)
          return np.sum(self.cm[None, :] / r[:, None] ** m_i[None, :], axis=1).ravel()
  
      @staticmethod
      def y(r, r_ref, p):
          return (r ** p - r_ref ** p) / (r ** p + r_ref ** p)
  
      def beta(self, r):
          y_p_ref = self.y(r, self.r_ref, self.p)
          y_q_ref = self.y(r, self.r_ref, self.q)
          return (self.beta_inf[None, :] * y_p_ref[:, None] + (1 - y_p_ref[:, None]) *
                  np.sum(self.betas[None, :] * y_q_ref[:, None]
                         ** np.arange(0, len(self.betas))[None, :], axis=1).ravel()[:, None])
  
      @property
      def beta_inf(self):
          return np.log(2 * self.d_e / self.u(self.r_eq))
  
  
  class BaseGraphBatchGenerator(Sequence):
      """Base class for classes that generate batches of training data for MEGNet.
      Based on the Sequence class, which is the data loader equivalent for Keras.
  
      Implementations of this base class must implement the :meth:`_generate_inputs`,
      which generates the lists of graph descriptions for a batch.
  
      The :meth:`process_atom_features` function and related functions are used to modify
      the features for each atom, bond, and global features when creating a batch.
      """
  
      def __init__(self, dataset_size, targets, batch_size=128, shuffle=True):
          """
          Args:
              dataset_size (int): Number of entries in dataset
              targets (ndarray): Feature to be predicted for each network
              batch_size (int): Maximum batch size
              shuffle (bool): Whether to shuffle the data after each step
          """
          if targets is not None:
              self.targets = np.array(targets)
          else:
              self.targets = None
          self.batch_size = batch_size
          self.total_n = dataset_size
          self.is_shuffle = shuffle
          self.max_step = int(np.ceil(self.total_n / batch_size))
          self.mol_index = np.arange(self.total_n)
          if self.is_shuffle:
              self.mol_index = np.random.permutation(self.mol_index)
  
      def __len__(self):
          return self.max_step
  
      def _combine_graph_data(self, feature_list_temp, connection_list_temp, global_list_temp,
                              index1_temp, index2_temp):
          """Compile the matrices describing each graph into single matrices for the entire graph
  
          Beyond concatenating the graph descriptions, this operation updates the indices of each
          node to be sequential across all graphs so they are not duplicated between graphs
  
          Args:
              feature_list_temp ([ndarray]): List of features for each node
              connection_list_temp ([ndarray]): List of features for each connection
              global_list_temp ([ndarray]): List of global state for each graph
              index1_temp ([ndarray]): List of indices for the start of each bond
              index2_temp ([ndarray]): List of indices for the end of each bond
          Returns:
              (tuple): Input arrays describing the entire batch of networks:
                  - ndarray: Features for each node
                  - ndarray: Features for each connection
                  - ndarray: Global state for each graph
                  - ndarray: Indices for the start of each bond
                  - ndarray: Indices for the end of each bond
                  - ndarray: Index of graph associated with each node
                  - ndarray: Index of graph associated with each connection
          """
          # get atom's structure id
          gnode = []
          for i, j in enumerate(feature_list_temp):
              gnode += [i] * len(j)
          # get bond features from a batch of structures
          # get bond's structure id
          gbond = []
          for i, j in enumerate(connection_list_temp):
              gbond += [i] * len(j)
  
          # assemble atom features together
          feature_list_temp = np.concatenate(feature_list_temp, axis=0)
          feature_list_temp = self.process_atom_feature(feature_list_temp)
  
          # assemble bond feature together
          connection_list_temp = np.concatenate(connection_list_temp, axis=0)
          connection_list_temp = self.process_bond_feature(connection_list_temp)
  
          # assemble state feature together
          global_list_temp = np.concatenate(global_list_temp, axis=0)
          global_list_temp = self.process_state_feature(global_list_temp)
  
          # assemble bond indices
          index1 = []
          index2 = []
          offset_ind = 0
          for ind1, ind2 in zip(index1_temp, index2_temp):
              index1 += [i + offset_ind for i in ind1]
              index2 += [i + offset_ind for i in ind2]
              offset_ind += (max(ind1) + 1)
  
          # Compile the inputs in needed order
          inputs = [expand_1st(feature_list_temp),
                    expand_1st(connection_list_temp),
                    expand_1st(global_list_temp),
                    expand_1st(index1),
                    expand_1st(index2),
                    expand_1st(gnode),
                    expand_1st(gbond)]
          return inputs
  
      def on_epoch_end(self):
          if self.is_shuffle:
              self.mol_index = np.random.permutation(self.mol_index)
  
      def process_atom_feature(self, x):
          return x
  
      def process_bond_feature(self, x):
          return x
  
      def process_state_feature(self, x):
          return x
  
      def __getitem__(self, index):
          # Get the indices for this batch
          batch_index = self.mol_index[index * self.batch_size:(index + 1) * self.batch_size]
  
          # Get the inputs for each batch
          inputs = self._generate_inputs(batch_index)
  
          # Make the graph data
          inputs = self._combine_graph_data(*inputs)
          
          # Return the batch
          if self.targets is None:
              return inputs
          else:
              # get targets
              it = itemgetter(*batch_index)
              target_temp = itemgetter_list(self.targets, batch_index)
              target_temp = np.atleast_2d(target_temp)
  
              return inputs, expand_1st(target_temp)
  
      @abstractmethod
      def _generate_inputs(self, batch_index):
          """Get the graph descriptions for each batch
  
          Args:
               batch_index ([int]): List of indices for training batch
          Returns:
              (tuple): Input arrays describing each network:
                  - [ndarray]: List of features for each node
                  - [ndarray]: List of features for each connection
                  - [ndarray]: List of global state for each graph
                  - [ndarray]: List of indices for the start of each bond
                  - [ndarray]: List of indices for the end of each bond
          """
          pass
  
  
  class GraphBatchGenerator(BaseGraphBatchGenerator):
      """
      A generator class that assembles several structures (indicated by
      batch_size) and form (x, y) pairs for model training.
  
      Args:
          atom_features: (list of np.array) list of atom feature matrix,
          bond_features: (list of np.array) list of bond features matrix
          state_features: (list of np.array) list of [1, G] state features,
              where G is the global state feature dimension
          index1_list: (list of integer) list of (M, ) one side atomic index of the bond,
          M is different for different structures
          index2_list: (list of integer) list of (M, ) the other side atomic
              index of the bond, M is different for different structures,
              but it has to be the same as the corresponding index1.
          targets: (numpy array), N*1, where N is the number of structures
          batch_size: (int) number of samples in a batch
      """
  
      def __init__(self,
                   atom_features,
                   bond_features,
                   state_features,
                   index1_list,
                   index2_list,
                   targets=None,
                   batch_size=128,
                   is_shuffle=True):
          super().__init__(len(atom_features), targets, batch_size, is_shuffle)
          self.atom_features = atom_features
          self.bond_features = bond_features
          self.state_features = state_features
          self.index1_list = index1_list
          self.index2_list = index2_list
  
      def _generate_inputs(self, batch_index):
          """Get the graph descriptions for each batch
  
          Args:
               batch_index ([int]): List of indices for training batch
          Returns:
              (tuple): Input arrays describe each network:
                  - [ndarray]: List of features for each nodes
                  - [ndarray]: List of features for each connection
                  - [ndarray]: List of global state for each graph
                  - [ndarray]: List of indices for the start of each bond
                  - [ndarray]: List of indices for the end of each bond
          """
  
          # Get the features and connectivity lists for this batch
          it = itemgetter(*batch_index)
          feature_list_temp = itemgetter_list(self.atom_features, batch_index)
          connection_list_temp = itemgetter_list(self.bond_features, batch_index)
          global_list_temp = itemgetter_list(self.state_features, batch_index)
          index1_temp = itemgetter_list(self.index1_list, batch_index)
          index2_temp = itemgetter_list(self.index2_list, batch_index)
  
          return feature_list_temp, connection_list_temp, global_list_temp, index1_temp, index2_temp
  
  
  class GraphBatchDistanceConvert(GraphBatchGenerator):
      """
      Generate batch of structures with bond distance being expanded using a Expansor
  
      Args:
          atom_features: (list of np.array) list of atom feature matrix,
          bond_features: (list of np.array) list of bond features matrix
          state_features: (list of np.array) list of [1, G] state features, where G is the global state feature dimension
          index1_list: (list of integer) list of (M, ) one side atomic index of the bond, M is different for different
              structures
          index2_list: (list of integer) list of (M, ) the other side atomic index of the bond, M is different for
              different structures, but it has to be the same as the correponding index1.
          targets: (numpy array), N*1, where N is the number of structures
          batch_size: (int) number of samples in a batch
          is_shuffle: (bool) whether to shuffle the structure, default to True
          distance_converter: (bool) converter for processing the distances
  
      """
  
      def __init__(self,
                   atom_features,
                   bond_features,
                   state_features,
                   index1_list,
                   index2_list,
                   targets=None,
                   batch_size=128,
                   is_shuffle=True,
                   distance_converter=None):
          super().__init__(atom_features=atom_features,
                           bond_features=bond_features,
                           state_features=state_features,
                           index1_list=index1_list,
                           index2_list=index2_list,
                           targets=targets,
                           batch_size=batch_size,
                           is_shuffle=is_shuffle)
          self.distance_converter = distance_converter
  
      def process_bond_feature(self, x):
          return self.distance_converter.convert(x)
  
  
  def itemgetter_list(l, indices):
      """
      Get indices of l and return a tuple
  
      Args:
          l:  (list)
          indices: (list) indices
  
      Returns:
          (tuple)
      """
      it = itemgetter(*indices)
      if np.size(indices) == 1:
          return it(l),
      else:
          return it(l)
  
  
  
  class HierarchGraphBatchDistanceConvert(BaseGraphBatchGenerator):
      """
      Generate batch of structures with bond distance being expanded using a Expansor
  
      Args:
          atom_features: (list of np.array) list of atom feature matrix,
          bond_features: (list of np.array) list of bond features matrix
          state_features: (list of np.array) list of [1, G] state features, where G is the global state feature dimension
          index1_list: (list of integer) list of (M, ) one side atomic index of the bond, M is different for different
              structures
          index2_list: (list of integer) list of (M, ) the other side atomic index of the bond, M is different for
              different structures, but it has to be the same as the correponding index1.
          targets: (numpy array), N*1, where N is the number of structures
          batch_size: (int) number of samples in a batch
          is_shuffle: (bool) whether to shuffle the structure, default to True
          distance_converter: (bool) converter for processing the distances
  
      """
  
      def __init__(self,
                   a_node_features,
                   a_bond_features,
                   a_state_features,
                   a_index1_list,
                   a_index2_list,
                   m_node_features,
                   m_bond_features,
                   m_state_features,
                   m_index1_list,
                   m_index2_list,
                   targets=None,
                   batch_size=128,
                   is_shuffle=True,
                   atom_distance_converter=None,
                   motif_distance_converter=None):
          super().__init__(len(a_node_features), targets, batch_size, is_shuffle)
          
          self.atom_distance_converter = atom_distance_converter
          self.motif_distance_converter = motif_distance_converter
          
          self.a_node_features = a_node_features
          self.a_bond_features = a_bond_features
          self.a_state_features = a_state_features
          self.a_index1_list = a_index1_list
          self.a_index2_list = a_index2_list
          
          self.m_node_features = m_node_features
          self.m_bond_features = m_bond_features
          self.m_state_features = m_state_features
          self.m_index1_list = m_index1_list
          self.m_index2_list = m_index2_list
  
      def process_bond_feature(self, x, is_atom):
          if is_atom:
              return self.atom_distance_converter.convert(x)
          else:
              return self.motif_distance_converter.convert(x)
  
      def _generate_inputs(self, batch_index):
          """Get the graph descriptions for each batch
  
          Args:
               batch_index ([int]): List of indices for training batch
          Returns:
              (tuple): Input arrays describe each network:
                  - [ndarray]: List of features for each nodes
                  - [ndarray]: List of features for each connection
                  - [ndarray]: List of global state for each graph
                  - [ndarray]: List of indices for the start of each bond
                  - [ndarray]: List of indices for the end of each bond
          """
  
          # Get the features and connectivity lists for this batch
          it = itemgetter(*batch_index)
          
          a_node_list = itemgetter_list(self.a_node_features, batch_index)
          a_bond_list = itemgetter_list(self.a_bond_features, batch_index)
          a_state_list = itemgetter_list(self.a_state_features, batch_index)
          a_index1_list = itemgetter_list(self.a_index1_list, batch_index)
          a_index2_list = itemgetter_list(self.a_index2_list, batch_index)
          
          m_node_list = itemgetter_list(self.m_node_features, batch_index)
          m_bond_list = itemgetter_list(self.m_bond_features, batch_index)
          m_state_list = itemgetter_list(self.m_state_features, batch_index)
          m_index1_list = itemgetter_list(self.m_index1_list, batch_index)
          m_index2_list = itemgetter_list(self.m_index2_list, batch_index)
  
          return a_node_list, a_bond_list, a_state_list, a_index1_list, a_index2_list, \
                 m_node_list, m_bond_list, m_state_list, m_index1_list, m_index2_list
  
  
      def _combine_graph_data(self, a_node_list, a_bond_list, a_state_list, a_index1_list, a_index2_list, \
                              m_node_list, m_bond_list, m_state_list, m_index1_list, m_index2_list):
          
          # ATOM part
          # get atom's structure id
          a_gnode = []
          for i, j in enumerate(a_node_list):
              a_gnode += [i] * len(j)
          # get bond features from a batch of structures
          # get bond's structure id
          a_gbond = []
          for i, j in enumerate(a_bond_list):
              a_gbond += [i] * len(j)
  
          # assemble atom features together
          a_node_list = np.concatenate(a_node_list, axis=0)
          a_node_list = self.process_atom_feature(a_node_list)
  
          # assemble bond feature together
          a_bond_list = np.concatenate(a_bond_list, axis=0)
          a_bond_list = self.process_bond_feature(a_bond_list, is_atom=True)
  
          # assemble state feature together
          a_state_list = np.concatenate(a_state_list, axis=0)
          a_state_list = self.process_state_feature(a_state_list)
  
          # assemble bond indices
          a_index1 = []
          a_index2 = []
          offset_ind = 0
          for ind1, ind2 in zip(a_index1_list, a_index2_list):
              a_index1 += [i + offset_ind for i in ind1]
              a_index2 += [i + offset_ind for i in ind2]
              offset_ind += (max(ind1) + 1)
          
          # MOTIF part
          # get atom's structure id
          m_gnode = []
          for i, j in enumerate(m_node_list):
              m_gnode += [i] * len(j)
          # get bond features from a batch of structures
          # get bond's structure id
          m_gbond = []
          for i, j in enumerate(m_bond_list):
              m_gbond += [i] * len(j)
  
          # assemble atom features together
          m_node_list = np.concatenate(m_node_list, axis=0)
          m_node_list = self.process_atom_feature(m_node_list)
  
          # assemble bond feature together
          m_bond_list = np.concatenate(m_bond_list, axis=0)
          m_bond_list = self.process_bond_feature(m_bond_list, is_atom=False)
  
          # assemble state feature together
          m_state_list = np.concatenate(m_state_list, axis=0)
          m_state_list = self.process_state_feature(m_state_list)
  
          # assemble bond indices
          m_index1 = []
          m_index2 = []
          offset_ind = 0
          for ind1, ind2 in zip(m_index1_list, m_index2_list):
              m_index1 += [i + offset_ind for i in ind1]
              m_index2 += [i + offset_ind for i in ind2]
              offset_ind += (max(ind1) + 1)
  
          # Compile the inputs in needed order
          inputs = [expand_1st(a_node_list),
                    expand_1st(a_bond_list),
                    expand_1st(a_state_list),
                    expand_1st(a_index1),
                    expand_1st(a_index2),
                    expand_1st(a_gnode),
                    expand_1st(a_gbond),
                    expand_1st(m_node_list),
                    expand_1st(m_bond_list),
                    expand_1st(m_state_list),
                    expand_1st(m_index1),
                    expand_1st(m_index2),
                    expand_1st(m_gnode),
                    expand_1st(m_gbond)]
          return inputs
  
  