text: # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.
  #
  # This source code is licensed under the MIT license found in the
  # LICENSE file in the root directory of this source tree.
  
  import time
  from concurrent.futures import ThreadPoolExecutor
  from typing import Iterator, Tuple, Callable, Optional, List, Dict
  
  import torch
  
  import tube
  import polygames
  from pytube.data_channel_manager import DataChannelManager
  
  from . import utils
  from .params import GameParams, ModelParams, SimulationParams, ExecutionParams
  from .env_creation_helpers import (
      sanitize_game_params,
      create_model,
      create_game,
      create_player,
  )
  
  
  #######################################################################################
  # HUMAN-PLAYED ENVIRONMENT CREATION
  #######################################################################################
  
  
  def create_human_environment(
      seed_generator: Iterator[int],
      game_params: GameParams,
      simulation_params: SimulationParams,
      execution_params: ExecutionParams,
      pure_mcts: bool,
      model
  ) -> Tuple[tube.Context, Optional[tube.DataChannel], Callable[[], int]]:
      human_first = execution_params.human_first
      time_ratio = execution_params.time_ratio
      total_time = execution_params.total_time
      context = tube.Context()
      actor_channel = (
          None if pure_mcts else tube.DataChannel("act", simulation_params.num_actor, 1)
      )
      rnn_state_shape = []
      if model is not None and hasattr(model, "rnn_cells") and model.rnn_cells > 0:
        rnn_state_shape = [model.rnn_cells, model.rnn_channels]
      rnn_state_size = 0
      if len(rnn_state_shape) >= 2:
        rnn_state_size = rnn_state_shape[0] * rnn_state_shape[1]
      logit_value = getattr(model, "logit_value", False)
      game = create_game(
          game_params,
          num_episode=1,
          seed=next(seed_generator),
          eval_mode=True,
          per_thread_batchsize=0,
          rewind=simulation_params.rewind,
          predict_end_state=game_params.predict_end_state,
          predict_n_states=game_params.predict_n_states,
      )
      player = create_player(
          seed_generator=seed_generator,
          game=game,
          player="mcts",
          num_actor=simulation_params.num_actor,
          num_rollouts=simulation_params.num_rollouts,
          pure_mcts=pure_mcts,
          actor_channel=actor_channel,
          model_manager=None,
          human_mode=True,
          total_time=total_time,
          time_ratio=time_ratio,
          sample_before_step_idx=80,
          randomized_rollouts=False,
          sampling_mcts=False,
          rnn_state_shape=rnn_state_shape,
          rnn_seqlen=execution_params.rnn_seqlen,
          logit_value=logit_value,
      )
      human_player = polygames.HumanPlayer()
      if game.is_one_player_game():
          game.add_human_player(human_player)
      else:
          if human_first:
              game.add_human_player(human_player)
              game.add_eval_player(player)
          else:
              game.add_eval_player(player)
              game.add_human_player(human_player)
  
      context.push_env_thread(game)
  
      def get_result_for_human_player():
          nonlocal game, human_first
          return game.get_result()[not human_first]
  
      return context, actor_channel, get_result_for_human_player
  
  
  def create_tp_environment(
      seed_generator: Iterator[int],
      game_params: GameParams,
      simulation_params: SimulationParams,
      execution_params: ExecutionParams,
      pure_mcts: bool,
  ) -> Tuple[tube.Context, Optional[tube.DataChannel], Callable[[], int]]:
      human_first = execution_params.human_first
      time_ratio = execution_params.time_ratio
      total_time = execution_params.total_time
      context = tube.Context()
      actor_channel = (
          None if pure_mcts else tube.DataChannel("act", simulation_params.num_actor, 1)
      )
      game = create_game(
          game_params,
          num_episode=1,
          seed=next(seed_generator),
          eval_mode=True,
          per_thread_batchsize=0,
      )
      player = create_player(
          seed_generator=seed_generator,
          game=game,
          num_actor=simulation_params.num_actor,
          num_rollouts=simulation_params.num_rollouts,
          pure_mcts=pure_mcts,
          actor_channel=actor_channel,
          model_manager=None,
          human_mode=True,
          total_time=total_time,
          time_ratio=time_ratio,
      )
      tp_player = polygames.TPPlayer()
      if game.is_one_player_game():
          game.add_tp_player(tp_player)
      else:
          if human_first:
              game.add_tp_player(tp_player)
              game.add_eval_player(player)
          else:
              game.add_eval_player(player)
              game.add_tp_player(tp_player)
  
      context.push_env_thread(game)
  
      def get_result_for_tp_player():
          nonlocal game, human_first
          return game.get_result()[not human_first]
  
      return context, actor_channel, get_result_for_tp_player
  
  
  #######################################################################################
  # HUMAN-PLAYED GAME
  #######################################################################################
  
  
  def _forward_pass_on_device(
      device: torch.device, model: torch.jit.ScriptModule, batch_s: torch.Tensor, batch_rnn_state: torch.Tensor = None
  ) -> Dict[str, torch.Tensor]:
      batch_s = utils.to_device(batch_s, device)
      if batch_rnn_state is not None:
          batch_rnn_state = utils.to_device(batch_rnn_state, device)
          with torch.no_grad():
              reply = model(batch_s, batch_rnn_state)
      else:
        with torch.no_grad():
            reply = model(batch_s)
      return reply
  
  
  def _play_game_against_mcts(context: tube.Context) -> None:
      context.start()
      while not context.terminated():
          time.sleep(1)
  
  
  def _play_game_against_neural_mcts(
      devices: List[torch.device],
      models: List[torch.jit.ScriptModule],
      context: tube.Context,
      actor_channel: tube.DataChannel,
  ) -> None:
      nb_devices = len(devices)
      context.start()
      dcm = DataChannelManager([actor_channel])
      # multithread
      with ThreadPoolExecutor(max_workers=nb_devices) as executor:
          while not context.terminated():
              batch = dcm.get_input(max_timeout_s=1)
              if len(batch) == 0:
                  continue
  
              assert len(batch) == 1
  
              # split in as many part as there are devices
              batches_s = torch.chunk(
                  batch[actor_channel.name]["s"], nb_devices, dim=0
              )
              has_rnn = "rnn_state" in batch[actor_channel.name]
              if has_rnn:
                  batches_rnn_state = torch.chunk(
                      batch[actor_channel.name]["rnn_state"], nb_devices, dim=0
                  )
              futures = []
              reply_eval = {"v": None, "pi_logit": None}
              if has_rnn:
                  for device, model, batch_s, batch_rnn_state in zip(
                      devices, models, batches_s, batches_rnn_state
                  ):
                      futures.append(
                          executor.submit(_forward_pass_on_device, device, model, batch_s, batch_rnn_state)
                      )
                  results = [future.result() for future in futures]
                  reply_eval["v"] = torch.cat([result["v"] for result in results], dim=0)
                  reply_eval["pi_logit"] = torch.cat([result["pi_logit"] for result in results], dim=0)
                  reply_eval["rnn_state_out"] = torch.cat([result["rnn_state"] for result in results], dim=0)
              else:
                  for device, model, batch_s in zip(
                      devices, models, batches_s
                  ):
                      futures.append(
                          executor.submit(_forward_pass_on_device, device, model, batch_s)
                      )
                  results = [future.result() for future in futures]
                  reply_eval["v"] = torch.cat([result["v"] for result in results], dim=0)
                  reply_eval["pi_logit"] = torch.cat([result["pi_logit"] for result in results], dim=0)
              dcm.set_reply(actor_channel.name, reply_eval)
      dcm.terminate()
  
  
  def play_game(
      pure_mcts: bool,
      devices: Optional[List[torch.device]],
      models: Optional[List[torch.jit.ScriptModule]],
      context: tube.Context,
      actor_channel: Optional[tube.DataChannel],
      get_result_for_human_player: Callable[[], int],
  ) -> int:
      if pure_mcts:
          _play_game_against_mcts(context)
      else:
          _play_game_against_neural_mcts(
              devices=devices, models=models, context=context, actor_channel=actor_channel
          )
      print("game over")
      return get_result_for_human_player()
  
  
  def play_tp_game(   #FIXME TODO not sure this helps
      pure_mcts: bool,
      devices: Optional[List[torch.device]],
      models: Optional[List[torch.jit.ScriptModule]],
      context: tube.Context,
      actor_channel: Optional[tube.DataChannel],
      get_result_for_human_player: Callable[[], int],
  ) -> int:
      if pure_mcts:
          _play_game_against_mcts(context)
      else:
          _play_game_against_neural_mcts(
              devices=devices, models=models, context=context, actor_channel=actor_channel
          )
      print("#game over")
      return get_result_for_human_player()
  
  
  #######################################################################################
  # OVERALL HUMAN-PLAYED GAME WORKFLOW
  #######################################################################################
  
  
  def run_human_played_game(
      game_params: GameParams,
      model_params: ModelParams,
      simulation_params: SimulationParams,
      execution_params: ExecutionParams,
  ):
      print("#" * 70)
      print("#" + "HUMAN-PLAYED GAME".center(68) + "#")
      print("#" * 70)
  
      print("setting-up pseudo-random generator...")
      seed_generator = utils.generate_random_seeds(seed=execution_params.seed)
  
      devices, models = None, None
      if not model_params.pure_mcts:
          print("loading pretrained model from checkpoint...")
          checkpoint = utils.load_checkpoint(checkpoint_path=model_params.init_checkpoint)
          game_params = checkpoint["game_params"]
          sanitize_game_params(game_params)
          model_params = checkpoint["model_params"]
          model_state_dict = checkpoint["model_state_dict"]
          del checkpoint
          print("creating model(s) and device(s)...")
          models = []
          devices = [torch.device(device) for device in execution_params.devices]
          for device in devices:
              model = create_model(game_params=game_params, model_params=model_params).to(
                  device
              )
              print("updating model...")
              model.load_state_dict(model_state_dict)
              model.eval()
              models.append(model)
  
      print("creating human-played environment")
      context, actor_channel, get_result_for_human_player = create_human_environment(
          seed_generator=seed_generator,
          game_params=game_params,
          simulation_params=simulation_params,
          execution_params=execution_params,
          pure_mcts=model_params.pure_mcts,
          model=model
      )
  
      print("playing against a human player...")
      human_score = play_game(
          pure_mcts=model_params.pure_mcts,
          devices=devices,
          models=models,
          context=context,
          actor_channel=actor_channel,
          get_result_for_human_player=get_result_for_human_player,
      )
  
      print(f"result for the human human_player: {human_score}")
  
  
  def run_tp_played_game(
      game_params: GameParams,
      model_params: ModelParams,
      simulation_params: SimulationParams,
      execution_params: ExecutionParams,
  ):
      print("#" * 70)
      print("#" + "HUMAN-PLAYED GAME".center(68) + "#")
      print("#" * 70)
  
      print("#setting-up pseudo-random generator...")
      seed_generator = utils.generate_random_seeds(seed=execution_params.seed)
  
      devices, models = None, None
      if not model_params.pure_mcts:
          print("#loading pretrained model from checkpoint...")
          checkpoint = utils.load_checkpoint(checkpoint_path=model_params.init_checkpoint)
          game_params = checkpoint["game_params"]
          sanitize_game_params(game_params)
          model_params = checkpoint["model_params"]
          model_state_dict = checkpoint["model_state_dict"]
          del checkpoint
          print("#creating model(s) and device(s)...")
          models = []
          devices = [torch.device(device) for device in execution_params.device]
          for device in devices:
              model = create_model(game_params=game_params, model_params=model_params).to(
                  device
              )
              print("#updating model...")
              model.load_state_dict(model_state_dict)
              model.eval()
              models.append(model)
  
      print("#creating human-played environment")
      context, actor_channel, get_result_for_human_player = create_tp_environment(
          seed_generator=seed_generator,
          game_params=game_params,
          simulation_params=simulation_params,
          execution_params=execution_params,
          pure_mcts=model_params.pure_mcts,
      )
  
      print("#playing against a tp player...")
      human_score = play_tp_game(
          pure_mcts=model_params.pure_mcts,
          devices=devices,
          models=models,
          context=context,
          actor_channel=actor_channel,
          get_result_for_human_player=get_result_for_human_player,
      )
  
      print(f"#result for the TP_player: {human_score}")
  