text: # cython: language_level=3
  import numpy as np
  from .foundation import AlgorithmConfig
  import copy
  from UTILS.colorful import *
  from UTILS.tensor_ops import __hash__, my_view
  def _flatten_helper(T, N, _tensor):
      return _tensor.view(T * N, *_tensor.size()[2:])
  
  '''
      轨迹
  '''
  
  class TRAJ_BASE():
      key_data_type = {}
      key_data_shape = {}
      max_mem_length = -1
      
      def __init__(self, traj_limit, env_id):
          self.traj_limit = traj_limit
          self.env_id = env_id
          self.readonly_lock = False
          self.key_dict = []
          self.time_pointer = 0
          self.need_reward_bootstrap = False
          self.deprecated_flag = False
  
      def check_type_shape(self, key, first_content=None):
          if first_content is not None:
              content_type = first_content.dtype
              content_shape = first_content.shape
              if key in trajectory.key_data_type: 
                  assert trajectory.key_data_type[key] == content_type
              else: 
                  trajectory.key_data_type[key] = content_type
                  trajectory.key_data_shape[key] = content_shape
              return content_type, content_shape
          assert key in trajectory.key_data_type
          return trajectory.key_data_type[key], trajectory.key_data_shape[key]
  
      def init_track(self, key, first_content):
          content = first_content
          self.check_type_shape(key, first_content)
          assert isinstance(content, np.ndarray) or isinstance(content, float), (key, content.__class__)
          # assert self.time_pointer == 0, key  # 问题
          tensor_size = ((self.traj_limit,) + tuple(content.shape))
          set_item = np.zeros(shape=tensor_size, dtype=content.dtype)
          set_item[:] = np.nan  if np.issubdtype(content.dtype, np.floating) else 0
          setattr(self, key, set_item)
          self.key_dict.append(key)
  
      def init_track_none(self, key):
          content_dtype, content_shape = self.check_type_shape(key)
          tensor_size = ((self.traj_limit,) + tuple(content_shape))
          set_item = np.zeros(shape=tensor_size, dtype=content_dtype)
          set_item[:] = np.nan  if np.issubdtype(content_dtype, np.floating) else 0
          setattr(self, key, set_item)
          self.key_dict.append(key)
  
      def remember(self, key, content):
          assert not self.readonly_lock
          if not (key in self.key_dict) and (content is not None):
              self.init_track(key=key, first_content=content)
              getattr(self, key)[self.time_pointer] = content
          elif not (key in self.key_dict) and (content is None):
              self.init_track_none(key=key)
          elif (key in self.key_dict) and (content is not None):
              getattr(self, key)[self.time_pointer] = content
          else:
              pass
  
      def time_shift(self):
          assert self.time_pointer < self.traj_limit
          self.time_pointer += 1
  
      def cut_tail(self): 
          TJ = lambda key: getattr(self, key)
          self.readonly_lock = True
          n_frame = self.time_pointer
          # check is buffer size too big
          if n_frame > trajectory.max_mem_length: 
              trajectory.max_mem_length = n_frame
              print('max_mem_length：%d, traj_limit:%d'%(trajectory.max_mem_length, self.traj_limit))
          # clip tail
          for key in self.key_dict: setattr(self, key, TJ(key)[:n_frame])
  
  class trajectory(TRAJ_BASE):
  
      def __init__(self, traj_limit, env_id):
          super().__init__(traj_limit, env_id)
          self.reference_track_name = 'value'
  
      def early_finalize(self):
          assert not self.readonly_lock   # unfinished traj
          self.need_reward_bootstrap = True
  
      def set_terminal_obs(self, tobs):
          self.tobs = copy.deepcopy(tobs)
  
      def cut_tail(self): # get_hyper_reward(self):
          # 删去多余的预留空间
          super().cut_tail()
          TJ = lambda key: getattr(self, key)
          # 进一步地， 根据这个轨迹上的NaN，删除所有无效时间点
          reference_track = getattr(self, self.reference_track_name)
          if self.need_reward_bootstrap:
              # print('need_reward_bootstrap')
              # 找到最后一个不是nan的位置
              T = np.where(~np.isnan(reference_track.squeeze()))[0][-1]
              self.boot_strap_value = {
                  'bootstrap_value':TJ('value').squeeze()[T].copy(), 
              }
              assert not hasattr(self,'tobs')
              self.set_terminal_obs(TJ('g_obs')[T].copy())
              reference_track[T] = np.nan
  
          # deprecated if nothing in it
          p_invalid = np.isnan(my_view(reference_track, [0, -1])).any(axis=-1)
          p_valid = ~p_invalid
          if p_invalid.all(): #invalid traj
              self.deprecated_flag = True
              return
  
          # adjust reward position
          reward = TJ('reward')
          for i in reversed(range(self.time_pointer)):
              if p_invalid[i] and i != 0: # invalid, push reward forward
                  reward[i-1] += reward[i]; reward[i] = np.nan
          setattr(self, 'reward', reward)
          # clip NaN
          for key in self.key_dict: setattr(self, key, TJ(key)[p_valid])
          # all done
          return
  
      # new finalize
      def finalize(self, hyper_reward=None):
          if hyper_reward is not None: 
              assert self.finalize
          self.readonly_lock = True
          assert hyper_reward is not None
          assert not self.deprecated_flag
  
          assert not np.isnan(getattr(self,'reward')).any()
          self.gae_finalize_return(reward_key='reward', value_key='value', new_return_name='return')
          return
  
      def copy_track(self, origin_key, new_key):
          if hasattr(self, origin_key):
              origin_handle = getattr(self, origin_key)
              setattr(self, new_key, origin_handle.copy())
              new_handle = getattr(self, new_key)
              self.key_dict.append(new_key)
              #return origin_handle, new_handle
          else:
              real_key_list = [real_key for real_key in self.__dict__ if (origin_key+'>' in real_key)]
              assert len(real_key_list)>0
              for real_key in real_key_list:
                  mainkey, subkey = real_key.split('>')
                  self.copy_track(real_key, (new_key+'>'+subkey))
              #return
  
      def gae_finalize_return(self, reward_key, value_key, new_return_name):
  
          gamma = AlgorithmConfig.gamma # ------- gae parameters -------
          tau = AlgorithmConfig.tau
          # ------- -------------- -------
          rewards = getattr(self, reward_key)
          value = getattr(self, value_key)
          length = rewards.shape[0]
          assert rewards.shape[0]==value.shape[0]
          # if dimension not aligned
          if rewards.ndim == value.ndim-1: rewards = np.expand_dims(rewards, -1)
          # initalize two more tracks
          setattr(self, new_return_name, np.zeros_like(value))
          self.key_dict.append(new_return_name)
  
          returns = getattr(self, new_return_name)
          boot_strap = 0 if not self.need_reward_bootstrap else self.boot_strap_value['bootstrap_'+value_key]
  
          for step in reversed(range(length)):
              if step==(length-1): # 最后一帧
                  value_preds_delta = rewards[step] + gamma * boot_strap      - value[step]
                  gae = value_preds_delta
              else:
                  value_preds_delta = rewards[step] + gamma * value[step + 1] - value[step]
                  gae = value_preds_delta + gamma * tau * gae
              returns[step] = gae + value[step]
  
  
  
  def calculate_sample_entropy(samples):
      key = []
      freq = []
      n_sample = len(samples)
      for s in samples:
          if s not in key:
              key.append(s)
              freq.append(1)
          else:
              i = key.index(s)
              freq[i] += 1
  
      entropy = 0.0
      for j,f in enumerate(freq):
          freq[j] /= n_sample
          entropy += -freq[j] * np.log(freq[j])
      # print亮红(key)
      # print亮红(freq)
      return entropy
  
  
  
  
  
  
  
  
  
  
  
  
  class TrajPoolManager(object):
      def __init__(self, n_pool):
          self.n_pool =  n_pool
          # self.traj_pool_history = []
          self.hyper_reward = []
          self.traj_pool_index = []
          self.cnt = 0
  
          self.clip_entropy_max = 4
          self.entropy_coef = 0
  
      def absorb_finalize_pool(self, pool):
          # self.traj_pool_history.append(pool)   # OOM
          pattern = []
          for traj_handle in pool:
              traj_handle.cut_tail()
          # h_reward = np.array([]).mean()
          # pattern_entropy = calculate_sample_entropy(pattern)
          # print亮绿('entropy:%.3f'%pattern_entropy)
          # h_reward = min(self.clip_entropy_max, pattern_entropy)*self.entropy_coef
          h_reward = 0
          # print亮绿('h_reward:%.3f'%h_reward)
          pool = list(filter(lambda traj: not traj.deprecated_flag, pool))
          for traj_handle in pool: 
              traj_handle.finalize(hyper_reward=h_reward)
          # delete deprecated
  
          self.cnt += 1
          task = ['train']
          return task, pool
  
  
  
  
  
  
  
  
  
  
  
  
  '''
      轨迹池管理
  '''
  class BatchTrajManager():
      templete = {
          # exam that trajectory have at least following things 
          "on-policy":  ['_SKIP_', 'obs', 'actions', 'reward', '_DONE_', 'value', 'actionLogProbs'],
          "off-policy": ['_SKIP_', 'obs', 'actions', 'reward', '_DONE_'],
  
      }
      def __init__(self, n_env, traj_limit, templete, trainer_hook, purge_outdated):
          self.trainer_hook = trainer_hook
          self.n_env = n_env
          self.traj_limit = traj_limit
          self.train_traj_needed = AlgorithmConfig.train_traj_needed
          self.upper_training_epoch = AlgorithmConfig.upper_training_epoch
          self.live_trajs = [trajectory(self.traj_limit, env_id=i) for i in range(n_env)]
          self.live_traj_frame = [0 for _ in range(self.n_env)]
          self.traj_pool = []
          self.registered_keys = []
          self._traj_lock_buf = None
          self.pool_manager = TrajPoolManager(n_pool=self.upper_training_epoch)
          self.patience = 1e3
          self.update_cnt = 0
          self.purge_outdated = purge_outdated
  
      def update(self, traj_frag, index):
          assert traj_frag is not None
          for j, env_i in enumerate(index):
              traj_handle = self.live_trajs[env_i]
              for key in traj_frag:
                  if traj_frag[key] is None:
                      assert False, key
                  if isinstance(traj_frag[key], dict):  # 如果是二重字典，特殊处理
                      for sub_key in traj_frag[key]:
                          content = traj_frag[key][sub_key][j]
                          traj_handle.remember(key + ">" + sub_key, content)
                  else:
                      content = traj_frag[key][j]
                      traj_handle.remember(key, content)
              self.live_traj_frame[env_i] += 1
              traj_handle.time_shift()
          return
  
      # 函数入口
      def feed_traj(self, traj_frag, require_hook=False):
          assert self._traj_lock_buf is None
          # an unlock hook must be exected before new trajectory feed in
          if require_hook: # the traj_frag is not intact, lock up traj_frag, wait for more
              assert '_DONE_' not in traj_frag
              assert 'reward' not in traj_frag
              self._traj_lock_buf = traj_frag
              return self._unlock_hook
          else:
              assert '_DONE_' in traj_frag
              assert '_SKIP_' in traj_frag
              self.__batch_update(traj_frag=traj_frag)
  
      def get_traj_frame(self):
          return self.live_traj_frame
          
      def train_and_clear_traj_pool(self):
          print('do update %d'%self.update_cnt)
          
          if self.purge_outdated:
              for env_index in range(self.n_env): 
                  if self.live_traj_frame[env_index] == 0: continue # 正常结束，不用分割
                  assert False
                  self.live_trajs[env_index].early_finalize() # 提前结束，进行切割
  
                  self.traj_pool.append(self.live_trajs[env_index])
                  self.live_trajs[env_index] = trajectory(self.traj_limit, env_id=env_index)
                  self.live_traj_frame[env_index] = 0
  
                  
          current_task_l, self.traj_pool = self.pool_manager.absorb_finalize_pool(pool=self.traj_pool)
          for current_task in current_task_l:
              ppo_update_cnt = self.trainer_hook(self.traj_pool, current_task)
  
          self.traj_pool = []
          self.update_cnt += 1
          # assert ppo_update_cnt == self.update_cnt
          return self.update_cnt
  
      def can_exec_training(self):
          if len(self.traj_pool) >= self.train_traj_needed:
              return True
          else:
              return False
              
      def _unlock_hook(self, traj_frag):
          assert self._traj_lock_buf is not None
          traj_frag.update(self._traj_lock_buf)
          self._traj_lock_buf = None
          assert '_DONE_' in traj_frag
          assert '_SKIP_' in traj_frag
          self.__batch_update(traj_frag=traj_frag)
  
      def ___check_integraty(self, traj_frag):
          # can not alway waste time checking this
          if self.patience < 0: return
          self.patience -= 1
          for key in traj_frag:
              if key not in self.registered_keys and (not key.startswith('_')):
                  self.registered_keys.append(key)
          for key in self.registered_keys:
              assert key in traj_frag, ('this key sometimes disappears from the traj_frag:', key)
  
      def __batch_update(self, traj_frag):
          self.___check_integraty(traj_frag)
          done = traj_frag['_DONE_']; traj_frag.pop('_DONE_') # done flag
          skip = traj_frag['_SKIP_']; traj_frag.pop('_SKIP_') # skip/frozen flag
          tobs = traj_frag['_TOBS_']; traj_frag.pop('_TOBS_') # terminal obs
          # single bool to list bool
          if isinstance(done, bool): done = [done for i in range(self.n_env)]
          if isinstance(skip, bool): skip = [skip for i in range(self.n_env)]
          n_active = sum(~skip)
          # feed
          cnt = 0
          for env_i in range(self.n_env):
              if skip[env_i]: continue
              # otherwise
              frag_index = cnt; cnt += 1
              env_index = env_i
              traj_handle = self.live_trajs[env_index]
              for key in traj_frag:
                  self.make_traj_remember(traj=traj_handle, key=key, content=traj_frag[key],frag_index=frag_index, n_active=n_active)
              self.live_traj_frame[env_index] += 1
              traj_handle.time_shift()
              if done[env_i]:
                  assert tobs[env_i] is not None # get the final obs
                  traj_handle.set_terminal_obs(tobs[env_i])
                  self.traj_pool.append(traj_handle)
                  self.live_trajs[env_index] = trajectory(self.traj_limit, env_id=env_index)
                  self.live_traj_frame[env_index] = 0
          # print红(self.live_traj_frame)
      def make_traj_remember(self, traj, key, content, frag_index, n_active):
          if content is None:
              traj.remember(key, None)
          elif isinstance(content, dict):
              for sub_key in content: 
                  self.make_traj_remember(traj, "".join((key , ">" , sub_key)), content=content[sub_key], frag_index=frag_index, n_active=n_active)
          else:
              assert n_active == len(content), ('length error')
              traj.remember(key, content[frag_index])
  