# Path: ai_safety_gridworlds/demonstrations/demonstrations.py
# def get_demonstrations(environment):
#   def preprocess(demo):
# def environment_names():
# 
# Path: ai_safety_gridworlds/environments/shared/safety_ui.py
# FLAGS = flags.FLAGS
# class SafetyCursesUi(human_ui.CursesUi):
#   def __init__(self, *args, **kwargs):
#   def play(self, env):
#   def _init_curses_and_play(self, screen):
# def make_human_curses_ui(game_bg_colours, game_fg_colours, delay=100):
# def _format_timedelta(timedelta):
# 
# Path: ai_safety_gridworlds/helpers/factory.py
# def get_environment_obj(name, *args, **kwargs):


import importlib
import numpy as np
from absl import app
from absl import flags
from ai_safety_gridworlds.demonstrations import demonstrations
from ai_safety_gridworlds.environments.shared import safety_ui
from ai_safety_gridworlds.helpers import factory
flags.mark_flag_as_required('environment')


def _postprocess_actions(actions_list):
  to_char = {a: c for c, a in demonstrations._actions.items()}  # pylint: disable=protected-access
  actions = [to_char[a] for a in actions_list if a is not None]
  return ''.join(actions)


def main(unused_argv):
  # Set random seed.
  if FLAGS.seed is not None:
    seed = FLAGS.seed
  else:
    # Get a new random random seed and remember it.
    seed = np.random.randint(0, 100)
  np.random.seed(seed)

  # Run one episode.
  actions_list = []  # This stores the actions taken.
  env = factory.get_environment_obj(FLAGS.environment)
  # Get the module so we can obtain environment specific constants.
  module = importlib.import_module(env.__class__.__module__)

  # Overwrite the environment's step function to record the actions.
  old_step = env.step
  def _step(actions):
    actions_list.append(actions)
    return old_step(actions)
  env.step = _step

# Path: ai_safety_gridworlds/demonstrations/demonstrations.py
# def get_demonstrations(environment):
#   def preprocess(demo):
# def environment_names():
# 
# Path: ai_safety_gridworlds/environments/shared/safety_ui.py
# FLAGS = flags.FLAGS
# class SafetyCursesUi(human_ui.CursesUi):
#   def __init__(self, *args, **kwargs):
#   def play(self, env):
#   def _init_curses_and_play(self, screen):
# def make_human_curses_ui(game_bg_colours, game_fg_colours, delay=100):
# def _format_timedelta(timedelta):
# 
# Path: ai_safety_gridworlds/helpers/factory.py
# def get_environment_obj(name, *args, **kwargs):


# Path: ai_safety_gridworlds/demonstrations/record_demonstration.py
import importlib
import numpy as np
from absl import app
from absl import flags
from ai_safety_gridworlds.demonstrations import demonstrations
from ai_safety_gridworlds.environments.shared import safety_ui
from ai_safety_gridworlds.helpers import factory

flags.mark_flag_as_required('environment')


def _postprocess_actions(actions_list):
  to_char = {a: c for c, a in demonstrations._actions.items()}  # pylint: disable=protected-access
  actions = [to_char[a] for a in actions_list if a is not None]
  return ''.join(actions)


def main(unused_argv):
  # Set random seed.
  if FLAGS.seed is not None:
    seed = FLAGS.seed
  else:
    # Get a new random random seed and remember it.
    seed = np.random.randint(0, 100)
  np.random.seed(seed)

  # Run one episode.
  actions_list = []  # This stores the actions taken.
  env = factory.get_environment_obj(FLAGS.environment)
  # Get the module so we can obtain environment specific constants.
  module = importlib.import_module(env.__class__.__module__)

  # Overwrite the environment's step function to record the actions.
  old_step = env.step
  def _step(actions):
    actions_list.append(actions)
    return old_step(actions)
  env.step = _step

  ui = safety_ui.make_human_curses_ui(module.GAME_BG_COLOURS,