# coding=utf-8
# coding=utf-8
# Copyright 2019 The RecSim Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
r"""An executable class to run agents in the simulator."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import os
import time

from absl import flags
from dopamine.discrete_domains import checkpointer
import gin.tf
from gym import spaces
import numpy as np
from recsim.simulator import environment
import tensorflow.compat.v1 as tf


flags.DEFINE_bool(
    'debug_mode', False,
    'If set to true, the agent will output in-episode statistics '
    'to Tensorboard. Disabled by default as this results in '
    'slower training.')
flags.DEFINE_string('agent_name', None, 'Name of the agent.')
flags.DEFINE_string('base_dir', None,
                    'Base directory to host all required sub-directories.')
flags.DEFINE_string(
    'environment_name', 'interest_evolution',
    'The environment with which to run the experiment. Supported choices are '
    '{interest_evolution, interest_exploration}.')
flags.DEFINE_string(
    'episode_log_file', '',
    'Filename under base_dir to output simulated episodes in SequenceExample.')
flags.DEFINE_multi_string(
    'gin_files', [], 'List of paths to gin configuration files (e.g.'
    '"third_party/py/dopamine/agents/dqn/dqn.gin").')
flags.DEFINE_multi_string(
    'gin_bindings', [],
    'Gin bindings to override the values set in the config files '
    '(e.g. "runner_lib.Runner.max_steps_per_episode=100')


FLAGS = flags.FLAGS


def load_gin_configs(gin_files, gin_bindings):
  """Loads gin configuration files.

  Args:
    gin_files: list, of paths to the gin configuration files for this
      experiment.
    gin_bindings: list, of gin parameter bindings to override the values in the
      config files.
  """
  gin.parse_config_files_and_bindings(
      gin_files, bindings=gin_bindings, skip_unknown=False)


@gin.configurable
class Runner(object):
  """Object that handles running experiments.

  Here we use the term 'experiment' to mean simulating interactions between the
  agent and the environment and reporting some statistics pertaining to these
  interactions.
  """

  _output_dir = None
  _checkpoint_dir = None
  _agent = None
  _checkpointer = None

  def __init__(self,
               base_dir,
               create_agent_fn,
               env,
               episode_log_file='',
               checkpoint_file_prefix='ckpt',
               max_steps_per_episode=27000):
    """Initializes the Runner object in charge of running a full experiment.

    Args:
      base_dir: str, the base directory to host all required sub-directories.
      create_agent_fn: A function that takes as args a Tensorflow session and an
        environment, and returns an agent.
      env: A Gym environment for running the experiments.
      episode_log_file: Path to output simulated episodes in tf.SequenceExample.
        Disable logging if episode_log_file is an empty string.
      checkpoint_file_prefix: str, the prefix to use for checkpoint files.
      max_steps_per_episode: int, maximum number of steps after which an episode
        terminates.
    """
    tf.logging.info('max_steps_per_episode = %s', max_steps_per_episode)

    if base_dir is None:
      raise ValueError('Missing base_dir.')

    self._base_dir = base_dir
    self._create_agent_fn = create_agent_fn
    self._env = env
    self._checkpoint_file_prefix = checkpoint_file_prefix
    self._max_steps_per_episode = max_steps_per_episode
    self._episode_log_file = episode_log_file
    self._episode_writer = None

  def _set_up(self, eval_mode):
    """Sets up the runner by creating and initializing the agent."""
    # Reset the tf default graph to avoid name collisions from previous runs
    # before doing anything else.
    tf.reset_default_graph()
    self._summary_writer = tf.summary.FileWriter(self._output_dir)
    if self._episode_log_file:
      self._episode_writer = tf.io.TFRecordWriter(
          os.path.join(self._output_dir, self._episode_log_file))
    # Set up a session and initialize variables.
    self._sess = tf.Session(config=tf.ConfigProto(allow_soft_placement=True))
    self._agent = self._create_agent_fn(
        self._sess,
        self._env,
        summary_writer=self._summary_writer,
        eval_mode=eval_mode)
    # type check: env/agent must both be multi- or single-user
    if self._agent.multi_user and not isinstance(
        self._env.environment, environment.MultiUserEnvironment):
      raise ValueError('Multi-user agent requires multi-user environment.')
    if not self._agent.multi_user and isinstance(
        self._env.environment, environment.MultiUserEnvironment):
      raise ValueError('Single-user agent requires single-user environment.')
    self._summary_writer.add_graph(graph=tf.get_default_graph())
    self._sess.run(tf.global_variables_initializer())
    self._sess.run(tf.local_variables_initializer())

  def _initialize_checkpointer_and_maybe_resume(self, checkpoint_file_prefix):
    """Reloads the latest checkpoint if it exists.

    This method will first create a `Checkpointer` object and then call
    `checkpointer.get_latest_checkpoint_number` to determine if there is a valid
    checkpoint in self._checkpoint_dir, and what the largest file number is.
    If a valid checkpoint file is found, it will load the bundled data from this
    file and will pass it to the agent for it to reload its data.
    If the agent is able to successfully unbundle, this method will increase and
    return the iteration number keyed by 'current_iteration' and the step number
    keyed by 'total_steps' as the return values.

    Args:
      checkpoint_file_prefix: str, the checkpoint file prefix.
    Returns:
      start_iteration: The iteration number to be continued after the latest
        checkpoint.
      start_step: The step number to be continued after the latest checkpoint.
    """
    self._checkpointer = checkpointer.Checkpointer(self._checkpoint_dir,
                                                   checkpoint_file_prefix)
    start_iteration = 0
    start_step = 0
    # Check if checkpoint exists.
    # Note that the existence of checkpoint 0 means that we have finished
    # iteration 0 (so we will start from iteration 1).
    latest_checkpoint_version = checkpointer.get_latest_checkpoint_number(
        self._checkpoint_dir)
    if latest_checkpoint_version >= 0:
      assert not self._episode_writer, 'Can only log episodes from scratch.'
      experiment_data = self._checkpointer.load_checkpoint(
          latest_checkpoint_version)
      start_iteration = experiment_data['current_iteration'] + 1
      del experiment_data['current_iteration']
      start_step = experiment_data['total_steps'] + 1
      del experiment_data['total_steps']
      if self._agent.unbundle(self._checkpoint_dir, latest_checkpoint_version,
                              experiment_data):
        tf.logging.info(
            'Reloaded checkpoint and will start from '
            'iteration %d', start_iteration)
    return start_iteration, start_step

  def _log_one_step(self, user_obs, doc_obs, slate, responses, reward,
                    is_terminal, sequence_example):
    """Adds one step of agent-environment interaction into SequenceExample.

    Args:
      user_obs: An array of floats representing user state observations
      doc_obs: A list of observations of the documents
      slate: An array of indices to doc_obs
      responses: A list of observations of responses for items in the slate
      reward: A float for the reward returned after this step
      is_terminal: A boolean for whether a terminal state has been reached
      sequence_example: A SequenceExample proto for logging current episode
    """

    def _add_float_feature(feature, values):
      feature.feature.add(float_list=tf.train.FloatList(value=values))

    def _add_int64_feature(feature, values):
      feature.feature.add(int64_list=tf.train.Int64List(value=values))

    if self._episode_writer is None:
      return
    fl = sequence_example.feature_lists.feature_list

    if isinstance(self._env.environment, environment.MultiUserEnvironment):
      for i, (single_user,
              single_slate,
              single_user_responses,
              single_reward) in enumerate(zip(user_obs,
                                              slate,
                                              responses,
                                              reward)):
        user_space = list(self._env.observation_space.spaces['user'].spaces)[i]
        _add_float_feature(fl['user_%d' % i], spaces.flatten(
            user_space, single_user))
        _add_int64_feature(fl['slate_%d' % i], single_slate)
        _add_float_feature(fl['reward_%d' % i], [single_reward])
        for j, response in enumerate(single_user_responses):
          resp_space = self._env.observation_space.spaces['response'][i][0]
          for k in response:
            _add_float_feature(fl['response_%d_%d_%s' % (i, j, k)],
                               spaces.flatten(resp_space, response))
    else:  # single-user environment
      _add_float_feature(
          fl['user'],
          spaces.flatten(self._env.observation_space.spaces['user'], user_obs))
      _add_int64_feature(fl['slate'], slate)
      for i, response in enumerate(responses):