text: #
  # Copyright (c) Facebook, Inc. and its affiliates.
  #
  # This source code is licensed under the MIT license found in the
  # LICENSE file in the root directory of this source tree.
  #
  
  from rlstructures import Agent, DictTensor
  import torch
  import os
  import sys
  import gym
  from gym.wrappers import TimeLimit
  from gym.spaces import Discrete
  from rlstructures.env_wrappers import GymEnv
  from rlstructures.batchers import EpisodeBatcher, Batcher
  import gym
  from gym.utils import seeding
  
  # Advanced Topics - Trajectory Batcher
  
  # A trajectory batcher will just acquire N timesteps (and not full episodes)
  class MyEnv(gym.Env):
      def __init__(self):
          super().__init__()
          self.action_space = Discrete(2)
  
      def seed(self, seed=None):
          self.np_random, seed = seeding.np_random(seed)
  
      def reset(self, env_info={"env_id": 0}):
          assert "env_id" in env_info
          self.env_id = env_info["env_id"]
          self.x = self.np_random.rand() * 2.0 - 1.0
          self.identifier = self.np_random.rand()
          obs = {"x": self.x, "identifier": self.identifier, "env_id": self.env_id}
          return obs
  
      def step(self, action):
          if action == 0:
              self.x -= 0.3
          else:
              self.x += 0.3
          done = self.x < -1 or self.x > 1
  
          obs = (
              {"x": self.x, "identifier": self.identifier, "env_id": self.env_id},
              self.x,
              done,
              {},
          )
          return obs
  
  
  # As you can see, the env_info can be used as an input parameter for the environment allowing to model multiple environments through a single class
  
  # We can do the same with agents, and implement an *Agent* that is parametrized by an *agent_info*. In our case, the agent is just an agent outputting its agent_id as an action.
  # Advanced examples are shown in the *rlaglos* directory (e.g stochastic/deterministic polices, epsilon-greedy policies, ...)
  
  
  class UniformAgent(Agent):
      def __init__(self, n_actions):
          super().__init__()
          self.n_actions = n_actions
  
      def __call__(self, state, observation, agent_info=None, history=None):
          B = observation.n_elems()
          agent_state = None
  
          # Initialize agent_info is not specified
          if agent_info is None:
              agent_info = DictTensor({"agent_id": torch.tensor([0]).repeat(B)})
  
          # initialize the state of the agent if not specified
          if state is None:
              agent_state = DictTensor({"timestep": torch.zeros(B).long()})
          else:
              agent_state = state
  
          scores = torch.randn(B, self.n_actions)
          probabilities = torch.softmax(scores, dim=1)
          actions = torch.distributions.Categorical(probabilities).sample()
          new_state = DictTensor({"timestep": agent_state["timestep"] + 1})
          # We also decide to output the action probabilities
          return (
              agent_state,
              DictTensor(
                  {
                      "action": actions,
                      "action_probabilities": probabilities,
                      "agent_id": agent_info["agent_id"],
                  }
              ),
              new_state,
          )
  
  
  # By specifying a particular value of *env_info* and *agent_infoo* when calling the *batcher.execute* method, the user may control which agent interacts with which environment
  # Let us illustrate this using **Multi-processes batchers**
  # * **NOTE:** For using multi-processes batchers, the user has to add the *slot* and *position_in_slot* arguments in the *Agent.__call__* methods (this will be explained later, but can be ignored)
  
  
  def create_env(seed=0, max_episode_steps=100):
      envs = []
      for k in range(4):
          e = MyEnv()
          e = TimeLimit(e, max_episode_steps=max_episode_steps)
          envs.append(e)
      return GymEnv(envs, seed=seed)
  
  
  def create_agent(n_actions=None):
      # Here, the buffer argument must be specified
      return UniformAgent(n_actions)
  
  
  if __name__ == "__main__":
      import torch.multiprocessing as mp
  
      mp.set_start_method("spawn")
  
      # The **EpisodeBatcher** will sample full episodes (until the environment returns done==True)
      # If one consider a rlstructures.VecEnv env, and n_threads (or processes), then the batcher will sample n_episodes = N * env.n_envs()*n_threads episodes at each execution (where N is chosen by the user)
      # *seeds* is a list of environment seeds, one seed per process
      # The batcher has to be configured 'at the right size' since all the processes are sharing a common *Buffer* to store trajectories
      # The simplest case is:
      # *n_slots = env.n_envs() x n_threads *
      # *n_timeteps* is the number of timesteps that will be acquired at each call
  
      batcher = Batcher(
          n_timesteps=10,
          n_slots=16,
          n_threads=4,
          seeds=[1, 2, 3, 4],
          create_agent=create_agent,
          agent_args={"n_actions": 2},
          create_env=create_env,
          env_args={"max_episode_steps": 100},
      )
  
      # A traajectory batcher has to be *reset*
      # Then calling *execute* will acquire the next T steps
      # The *execute* method will return *None* if all environments have stopped
      batcher.reset(
          agent_info=DictTensor({"agent_id": torch.arange(16)}),
          env_info=DictTensor({"env_id": torch.arange(16)}),
      )
      import time
  
      batcher.execute()
      t = batcher.get()
      while not t is None:
          print(t.lengths)
          batcher.execute()
          t = batcher.get()
  