text: from __future__ import absolute_import
  from __future__ import division
  import numpy as np
  from rl_code.data.feed_data import RelationEntityBatcher
  from rl_code.data.grapher import RelationEntityGrapher
  import logging
  
  logger = logging.getLogger()
  
  
  class Episode(object):
  
      def __init__(self, graph, data, params):
          """
  
          :param graph:
          :param data:
                  e1: shape = (batch_size, ).
                  r: shape = (batch_size, ).
                  e2: shape = (batch_size, ).
                  all_e2s: len = batch_size. It stores all corresponding e2s according to (e1, r).
          :param params:
          :return:
                  self.start_entities/current_entities/end_entities/query_relation: shape = (batch_size*num_rollouts, )
                  self.state = {'next_relations'/'next_entities':(batch_size*num_rollouts, max_num_actions),
                                'current_entities': (batch_size*num_rollouts, )}
          """
          self.grapher = graph
          self.batch_size, self.path_len, num_rollouts, test_rollouts, positive_reward, negative_reward, mode, batcher = params
          self.mode = mode
          if self.mode == 'train':
              self.num_rollouts = num_rollouts
          else:
              self.num_rollouts = test_rollouts
          self.current_hop = 0
          start_entities, query_relation,  end_entities, all_answers = data
          self.no_examples = start_entities.shape[0]      # = batch_size
          self.positive_reward = positive_reward          # a parameter, normally =1
          self.negative_reward = negative_reward          # a parameter, normally =0
          start_entities = np.repeat(start_entities, self.num_rollouts)       # shape = (batch_size*num_rollouts, )
          batch_query_relation = np.repeat(query_relation, self.num_rollouts)
          end_entities = np.repeat(end_entities, self.num_rollouts)
          self.start_entities = start_entities
          self.end_entities = end_entities
          self.current_entities = np.array(start_entities)
          self.query_relation = batch_query_relation
          self.all_answers = all_answers
  
          next_actions = self.grapher.return_next_actions(self.current_entities, self.start_entities, self.query_relation,
                                                          self.end_entities, self.all_answers, self.current_hop == self.path_len - 1,
                                                          self.num_rollouts)
          self.state = {}
          self.state['next_relations'] = next_actions[:, :, 1]
          self.state['next_entities'] = next_actions[:, :, 0]
          self.state['current_entities'] = self.current_entities
  
      def get_state(self):
          return self.state
  
      def get_query_relation(self):
          return self.query_relation
  
      def get_reward(self):
          reward = (self.current_entities == self.end_entities)
  
          # set the True and False values to the values of positive and negative rewards.
          condlist = [reward == True, reward == False]
          choicelist = [self.positive_reward, self.negative_reward]
          reward = np.select(condlist, choicelist)  # [B,]
          return reward
  
      def __call__(self, action):
          self.current_hop += 1
          self.current_entities = self.state['next_entities'][np.arange(self.no_examples*self.num_rollouts), action]
  
          next_actions = self.grapher.return_next_actions(self.current_entities, self.start_entities, self.query_relation,
                                                          self.end_entities, self.all_answers, self.current_hop == self.path_len - 1,
                                                          self.num_rollouts )
  
          self.state['next_relations'] = next_actions[:, :, 1]
          self.state['next_entities'] = next_actions[:, :, 0]
          self.state['current_entities'] = self.current_entities
          return self.state
  
  
  class env(object):
      def __init__(self, params, mode='train'):
  
          self.batch_size = params['batch_size']
          self.num_rollouts = params['num_rollouts']
          self.positive_reward = params['positive_reward']
          self.negative_reward = params['negative_reward']
          self.mode = mode
          self.path_len = params['path_length']
          self.test_rollouts = params['test_rollouts']
          input_dir = params['data_input_dir']
          if mode == 'train':
              self.batcher = RelationEntityBatcher(input_dir=input_dir,
                                                   batch_size=params['batch_size'],
                                                   entity_vocab=params['entity_vocab'],
                                                   relation_vocab=params['relation_vocab']
                                                   )
          else:
              self.batcher = RelationEntityBatcher(input_dir=input_dir,
                                                   mode =mode,
                                                   batch_size=params['batch_size']//4,
                                                   entity_vocab=params['entity_vocab'],
                                                   relation_vocab=params['relation_vocab'])
  
              self.total_no_examples = self.batcher.store.shape[0]
          self.grapher = RelationEntityGrapher(triple_store=params['data_input_dir'] + '/' + 'train.txt',   # 'graph.txt'
                                               max_num_actions=params['max_num_actions'],
                                               entity_vocab=params['entity_vocab'],
                                               relation_vocab=params['relation_vocab'])
  
      def get_episodes(self):
          params = self.batch_size, self.path_len, self.num_rollouts, self.test_rollouts, self.positive_reward, self.negative_reward, self.mode, self.batcher
          if self.mode == 'train':
              for data in self.batcher.yield_next_batch_train():
                  yield Episode(self.grapher, data, params)
          else:
              for data in self.batcher.yield_next_batch_test():
                  if data == None:
                      return
                  yield Episode(self.grapher, data, params)
  