# coding=utf-8
# Copyright 2018 The Dopamine 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.
"""Compact implementation of a DQN agent."""

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

import math
import os
import random

from absl import logging

from dopamine.discrete_domains import atari_lib
from dopamine.replay_memory import circular_replay_buffer
import numpy as np
import tensorflow as tf
import collections
import tensorflow.contrib.layers as layers

gfile = tf.gfile


# These are aliases which are used by other classes.
NATURE_DQN_OBSERVATION_SHAPE = atari_lib.NATURE_DQN_OBSERVATION_SHAPE
NATURE_DQN_DTYPE = atari_lib.NATURE_DQN_DTYPE
NATURE_DQN_STACK_SIZE = atari_lib.NATURE_DQN_STACK_SIZE
nature_dqn_network = atari_lib.NatureDQNNetwork


def linearly_decaying_epsilon(decay_period, step, warmup_steps, epsilon):
  """Returns the current epsilon for the agent's epsilon-greedy policy.

  This follows the Nature DQN schedule of a linearly decaying epsilon (Mnih et
  al., 2015). The schedule is as follows:
    Begin at 1. until warmup_steps steps have been taken; then
    Linearly decay epsilon from 1. to epsilon in decay_period steps; and then
    Use epsilon from there on.

  Args:
    decay_period: float, the period over which epsilon is decayed.
    step, the number of training steps completed so far.
    warmup_steps, the number of steps taken before epsilon is decayed.
    epsilon: float, the final value to which to decay the epsilon parameter.

  Returns:
    A float, the current epsilon value computed according to the schedule.
  """
  steps_left = decay_period + warmup_steps - step
  bonus = (1.0 - epsilon) * steps_left / decay_period
  bonus = np.clip(bonus, 0., 1. - epsilon)
  return epsilon + bonus


def identity_epsilon(unused_decay_period, unused_step, unused_warmup_steps,
                     epsilon):
  return epsilon


class DQNAgent(object):
  """An implementation of the DQN agent."""

  def __init__(self,num_actions,
               observation_shape=atari_lib.NATURE_DQN_OBSERVATION_SHAPE,
               observation_dtype=atari_lib.NATURE_DQN_DTYPE,
               stack_size=atari_lib.NATURE_DQN_STACK_SIZE,
               network=None,
               gamma=0.99,
               update_horizon=1,
               min_replay_history=20000,
               update_period=4,
               target_update_period=8000,
               epsilon_fn=linearly_decaying_epsilon,
               epsilon_train=0.01,
               epsilon_eval=0.001,
               epsilon_decay_period=250000,
               tf_device='/cpu:*',
               eval_mode=False,
               use_staging=False,
               max_tf_checkpoints_to_keep=4,
               optimizer=tf.train.AdamOptimizer(
                   learning_rate=0.0003),
               summary_writer=None,
               summary_writing_frequency=500,
               allow_partial_reload=False):
    """Initializes the agent and constructs the components of its graph.

    Args:
      sess: `tf.Session`, for executing ops.
      num_actions, number of actions the agent can take at any state.
      observation_shape: tuple of ints describing the observation shape.
      observation_dtype: tf.DType, specifies the type of the observations. Note
        that if your inputs are continuous, you should set this to tf.float32.
      stack_size, number of frames to use in state stack.
      network: tf.Keras.Model, expecting 2 parameters: num_actions,
        network_type. A call to this object will return an instantiation of the
        network provided. The network returned can be run with different inputs
        to create different outputs. See
        dopamine.discrete_domains.atari_lib.NatureDQNNetwork as an example.
      gamma: float, discount factor with the usual RL meaning.
      update_horizon, horizon at which updates are performed, the 'n' in
        n-step update.
      min_replay_history, number of transitions that should be experienced
        before the agent begins training its value function.
      update_period, period between DQN updates.
      target_update_period, update period for the target network.
      epsilon_fn: function expecting 4 parameters:
        (decay_period, step, warmup_steps, epsilon). This function should return
        the epsilon value used for exploration during training.
      epsilon_train: float, the value to which the agent's epsilon is eventually
        decayed during training.
      epsilon_eval: float, epsilon used when evaluating the agent.
      epsilon_decay_period, length of the epsilon decay schedule.
      tf_device: str, Tensorflow device on which the agent's graph is executed.
      eval_mode: bool, True for evaluation and False for training.
      use_staging: bool, when True use a staging area to prefetch the next
        training batch, speeding training up by about 30%.
      max_tf_checkpoints_to_keep, the number of TensorFlow checkpoints to
        keep.
      optimizer: `tf.train.Optimizer`, for training the value
        function.
      summary_writer: SummaryWriter object for outputting training statistics.
        Summary writing disabled if set to None.
      summary_writing_frequency, frequency with which summaries will be
        written. Lower values will result in slower training.
      allow_partial_reload: bool, whether we allow reloading a partial agent
        (for instance, only the network parameters).
    """
    assert isinstance(observation_shape, tuple)
    logging.info('Creating %s agent with the following parameters:',
                 self.__class__.__name__)
    logging.info('\t gamma: %f', gamma)
    logging.info('\t update_horizon: %f', update_horizon)
    logging.info('\t min_replay_history: %d', min_replay_history)
    logging.info('\t update_period: %d', update_period)
    logging.info('\t target_update_period: %d', target_update_period)
    logging.info('\t epsilon_train: %f', epsilon_train)
    logging.info('\t epsilon_eval: %f', epsilon_eval)
    logging.info('\t epsilon_decay_period: %d', epsilon_decay_period)
    logging.info('\t tf_device: %s', tf_device)
    logging.info('\t use_staging: %s', use_staging)
    logging.info('\t optimizer: %s', optimizer)
    logging.info('\t max_tf_checkpoints_to_keep: %d',
                 max_tf_checkpoints_to_keep)
    self.state_dim = 44
    self.num_actions = num_actions
    self.observation_shape = tuple(observation_shape)
    self.observation_dtype = observation_dtype
    self.stack_size = stack_size
    self.network = network
    self.gamma = gamma
    self.update_horizon = update_horizon
    self.cumulative_gamma = math.pow(gamma, update_horizon)
    self.min_replay_history = min_replay_history
    self.target_update_period = target_update_period
    self.epsilon_fn = epsilon_fn
    self.epsilon_train = epsilon_train
    self.epsilon_eval = epsilon_eval
    self.epsilon_decay_period = epsilon_decay_period
    self.update_period = update_period
    self.eval_mode = eval_mode
    self.training_steps = 0
    self.optimizer = optimizer
    self.summary_writer = summary_writer
    self.summary_writing_frequency = summary_writing_frequency
    self.allow_partial_reload = allow_partial_reload
    self.tf_device = tf_device
    # self.sess = sess
    # self.sess.run(tf.global_variables_initializer())

    with tf.device(self.tf_device):
      # Create a placeholder for the state input to the DQN network.
      # The last axis indicates the number of consecutive frames stacked.
      state_shape = (1,) + self.observation_shape + (stack_size,)
      self.state = np.zeros(state_shape)
      self.state_ph = tf.placeholder(
          tf.float32, [None, self.state_dim], name='state')
      self.Lambda_ph = tf.placeholder(
          tf.float32, [None, 1], name='lambda')
      self.next_state_ph = tf.placeholder(
        tf.float32, [None, self.state_dim], name='next_state')
      self.action_ph = tf.placeholder(
        tf.int32, [None, 1], name='current_action')
      self.reward_ph = tf.placeholder(
        tf.float32, [None, 1], name='reward_ph')
      self.done_ph = tf.placeholder(
        tf.float32, [None, 1], name='done_ph')
      self.next_action_ph = tf.placeholder(
        tf.int32, [None, 1])
      self.batch_size_state_ph = tf.placeholder(
        self.observation_dtype, (None, ) + state_shape[1:], name='batch_state_ph'
      )
      #self._replay = self._build_replay_buffer(use_staging)

      #self._build_networks()

      #self._train_op = self._build_train_op()
      self._sync_qt_ops = self._build_sync_op()

    if self.summary_writer is not None:
      # All tf.summaries should have been defined prior to running this.
      self._merged_summaries = tf.summary.merge_all()
    #self._sess = sess

    # var_map = atari_lib.maybe_transform_variable_names(
    #     tf.global_variables())
    # self._saver = tf.train.Saver(
    #     var_list=var_map, max_to_keep=max_tf_checkpoints_to_keep)

    # Variables to be initialized by the agent once it interacts with the
    # environment.
    self._observation = None
    self._last_observation = None


  def _create_network(self, name):
    """Builds the convolutional network used to compute the agent's Q-values.

    Args:
      name: str, this name is passed to the tf.keras.Model and used to create
        variable scope under the hood by the tf.keras.Model.
    Returns:
      network: tf.keras.Model, the network instantiated by the Keras model.
    """
    network = self.network(self.num_actions, name=name)
    return network

  def _build_networks(self):
    """Builds the Q-value network computations needed for acting and training.

    These are:
      self.online_convnet: For computing the current state's Q-values.
      self.target_convnet: For computing the next state's target Q-values.
      self._net_outputs: The actual Q-values.
      self._q_argmax: The action maximizing the current state's Q-values.
      self._replay_net_outputs: The replayed states' Q-values.
      self._replay_next_target_net_outputs: The replayed next states' target
        Q-values (see Mnih et al., 2015 for details).
    """

    # _network_template instantiates the model and returns the network object.
    # The network object can be used to generate different outputs in the graph.
    # At each call to the network, the parameters will be reused.
    self.online_convnet = tf.make_template('Online', self._network_template)
    self.target_convnet = tf.make_template('Target', self._network_template)
    # self.online_convnet = self._create_network(name='Online')
    # self.target_convnet = self._create_network(name='Target')
    self._net_outputs = self.online_convnet(self.state_ph, self.Lambda_ph)
    # TODO(bellemare): Ties should be broken. They are unlikely to happen when
    # using a deep network, but may affect performance with a linear
    # approximation scheme.
    self._q_argmax = tf.expand_dims(tf.argmax(self._net_outputs.q_values, axis=1), axis= 1, name = 'next_action')
    self._replay_net_outputs = self.online_convnet(self.state_ph, self.Lambda_ph)
    # self._replay_net_outputs = self.online_convnet(tf.cast(self.replay_buffer.new_batch_data["state"], dtype = tf.float32),
    #                                                self.replay_buffer.new_batch_data["Lambda"])
    self._replay_next_target_net_outputs = self.target_convnet(self.next_state_ph, self.Lambda_ph)

    # self._replay_next_target_net_outputs = self.target_convnet(
    #     tf.cast(self.replay_buffer.new_batch_data["next_state"], dtype = tf.float32), self.replay_buffer.new_batch_data["Lambda"])

  def _network_template(self, state, Lambda):
    """Builds the convolutional network used to compute the agent's Q-values.

    Args:
      state: `tf.Tensor`, contains the agent's current state.

    Returns:
      net: _network_type object containing the tensors output by the network.
    """
    activation_fn = tf.nn.relu
    r10 = tf.layers.batch_normalization(tf.concat([state, Lambda], 1), axis=-1, momentum=0.99,
                                        epsilon=0.001,
                                        center=True,
                                        scale=True,
                                        beta_initializer=tf.zeros_initializer(),
                                        gamma_initializer=tf.ones_initializer(),
                                        moving_mean_initializer=tf.zeros_initializer(),
                                        moving_variance_initializer=tf.ones_initializer(),
                                        beta_regularizer=None,
                                        gamma_regularizer=None,
                                        beta_constraint=None,
                                        gamma_constraint=None,
                                        training=False,
                                        trainable=True,
                                        name=None,
                                        reuse=None,
                                        renorm=False,
                                        renorm_clipping=None,
                                        renorm_momentum=0.99,
                                        fused=None,
                                        virtual_batch_size=None,
                                        adjustment=None)
    r11 = layers.fully_connected(r10, 256, activation_fn=tf.nn.relu)
    net = layers.fully_connected(r11, 256, activation_fn=tf.nn.relu)
    net = tf.layers.dense(
      net,
      units=512,
      activation=activation_fn,
    )
    q_values = tf.layers.dense(
      net,
      units=self.num_actions,
    )
    return self._get_network_type()(q_values)

  def _get_network_type(self):
    """Returns the type of the outputs of a Q value network.

    Returns:
      net_type: _network_type object defining the outputs of the network.
    """
    return collections.namedtuple('DQN_network', ['q_values'])
  #
  # def _build_replay_buffer(self, replay_buffer):
  #
  #     self._replay.states = self.replay_buffer.new_batch_data["state"]
  #     self._replay.next_states = self.replay_buffer.new_batch_data["next_state"]
  #     self._replay.rewards = self.replay_buffer.new_batch_data["rewards"]
  #     self._replay.actions = self.replay_buffer.new_batch_data["actions"]
  #     self._replay.terminals = self.replay_buffer.new_batch_data["done"]
  #     self._replay.Lambda = self.replay_buffer.new_batch_data["Lambda"]
  #     self._replay.next_action = self.replay_buffer.new_batch_data["next_action"],


  def _build_target_q_op(self):
    """Build an op used as a target for the Q-value.

    Returns:
      target_q_op: An op calculating the Q-value.
    """
    # Get the maximum Q-value across the actions dimension.
    replay_next_qt_max = tf.reduce_max(
        self._replay_next_target_net_outputs.q_values, 1)
    # Calculate the Bellman target value.
    #   Q_t = R_t + \gamma^N * Q'_t+1
    # where,
    #   Q'_t+1 = \argmax_a Q(S_t+1, a)
    #          (or) 0 if S_t is a terminal state,
    # and
    #   N is the update horizon (by default, N=1).
    return self.reward_ph + self.cumulative_gamma * replay_next_qt_max * (
        1. - tf.cast(self.done_ph, tf.float32))

  def _build_train_op(self):
    """Builds a training op.

    Returns:
      train_op: An op performing one step of training from replay data.
    """
    replay_action_one_hot = tf.one_hot(
        self.action_ph, self.num_actions, 1., 0., name='action_one_hot')
    replay_chosen_q = tf.reduce_sum(
        self._replay_net_outputs.q_values * replay_action_one_hot,
        axis=1,
        name='replay_chosen_q')

    target = tf.stop_gradient(self._build_target_q_op())
    loss = tf.losses.huber_loss(
        target, replay_chosen_q, reduction=tf.losses.Reduction.NONE)
    if self.summary_writer is not None:
      with tf.variable_scope('Losses'):
        tf.summary.scalar('HuberLoss', tf.reduce_mean(loss))
    return self.optimizer.minimize(tf.reduce_mean(loss))

  def _build_sync_op(self):
    """Builds ops for assigning weights from online to target network.

    Returns:
      ops: A list of ops assigning weights from online to target network.
    """
    # Get trainable variables from online and target DQNs
    sync_qt_ops = []
    scope = tf.get_default_graph().get_name_scope()
    trainables_online = tf.get_collection(
        tf.GraphKeys.TRAINABLE_VARIABLES,
        scope=os.path.join(scope, 'Online'))
    trainables_target = tf.get_collection(
        tf.GraphKeys.TRAINABLE_VARIABLES,
        scope=os.path.join(scope, 'Target'))

    for (w_online, w_target) in zip(trainables_online, trainables_target):
      # Assign weights from online to target network.
      sync_qt_ops.append(w_target.assign(w_online, use_locking=True))
    return sync_qt_ops

  def begin_episode(self, observation):
    """Returns the agent's first action for this episode.

    Args:
      observation: numpy array, the environment's initial observation.

    Returns:
      int, the selected action.
    """
    self._reset_state()
    self._record_observation(observation)

    if not self.eval_mode:
      self._train_step()

    self.action = self._select_action()
    return self.action

  def step(self, reward, observation):
    """Records the most recent transition and returns the agent's next action.

    We store the observation of the last time step since we want to store it
    with the reward.

    Args:
      reward: float, the reward received from the agent's most recent action.
      observation: numpy array, the most recent observation.

    Returns:
      int, the selected action.
    """
    self._last_observation = self._observation
    self._record_observation(observation)

    if not self.eval_mode:
      self._store_transition(self._last_observation, self.action, reward, False)
      self._train_step()

    self.action = self._select_action()
    return self.action

  def end_episode(self, reward):
    """Signals the end of the episode to the agent.

    We store the observation of the current time step, which is the last
    observation of the episode.

    Args:
      reward: float, the last reward from the environment.
    """
    if not self.eval_mode:
      self._store_transition(self._observation, self.action, reward, True)

  def _select_action(self):
    """Select an action from the set of available actions.

    Chooses an action randomly with probability self._calculate_epsilon(), and
    otherwise acts greedily according to the current Q-value estimates.

    Returns:
       int, the selected action.
    """
    if self.eval_mode:
      epsilon = self.epsilon_eval
    else:
      epsilon = self.epsilon_fn(
          self.epsilon_decay_period,
          self.training_steps,
          self.min_replay_history,
          self.epsilon_train)
    if random.random() <= epsilon:
      # Choose a random action with probability epsilon.
      return random.randint(0, self.num_actions - 1)
    else:
      # Choose the action with highest Q-value at the current state.
      return self._sess.run(self._q_argmax, {self.state_ph: self.state})

  def _train_step(self):
    """Runs a single training step.

    Runs a training op if both:
      (1) A minimum number of frames have been added to the replay buffer.
      (2) `training_steps` is a multiple of `update_period`.

    Also, syncs weights from online to target network if training steps is a
    multiple of target update period.
    """
    # Run a train op at the rate of self.update_period if enough training steps
    # have been run. This matches the Nature DQN behaviour.
    if self._replay.memory.add_count > self.min_replay_history:
      if self.training_steps % self.update_period == 0:
        self._sess.run(self._train_op)
        if (self.summary_writer is not None and
            self.training_steps > 0 and
            self.training_steps % self.summary_writing_frequency == 0):
          summary = self._sess.run(self._merged_summaries)
          self.summary_writer.add_summary(summary, self.training_steps)

      if self.training_steps % self.target_update_period == 0:
        self._sess.run(self._sync_qt_ops)

    self.training_steps += 1

  def _record_observation(self, observation):
    """Records an observation and update state.

    Extracts a frame from the observation vector and overwrites the oldest
    frame in the state buffer.

    Args:
      observation: numpy array, an observation from the environment.
    """
    # Set current observation. We do the reshaping to handle environments
    # without frame stacking.
    self._observation = np.reshape(observation, self.observation_shape)
    # Swap out the oldest frame with the current frame.
    self.state = np.roll(self.state, -1, axis=-1)
    self.state[0, ..., -1] = self._observation

  def _store_transition(self, last_observation, action, reward, is_terminal):
    """Stores an experienced transition.

    Executes a tf session and executes replay buffer ops in order to store the
    following tuple in the replay buffer:
      (last_observation, action, reward, is_terminal).

    Pedantically speaking, this does not actually store an entire transition
    since the next state is recorded on the following time step.

    Args:
      last_observation: numpy array, last observation.
      action, the action taken.
      reward: float, the reward.
      is_terminal: bool, indicating if the current state is a terminal state.
    """
    self._replay.add(last_observation, action, reward, is_terminal)

  def _reset_state(self):
    """Resets the agent state by filling it with zeros."""
    self.state.fill(0)

  # def bundle_and_checkpoint(self, checkpoint_dir, iteration_number):
  #   """Returns a self-contained bundle of the agent's state.
  #
  #   This is used for checkpointing. It will return a dictionary containing all
  #   non-TensorFlow objects (to be saved into a file by the caller), and it saves
  #   all TensorFlow objects into a checkpoint file.
  #
  #   Args:
  #     checkpoint_dir: str, directory where TensorFlow objects will be saved.
  #     iteration_number, iteration number to use for naming the checkpoint
  #       file.
  #
  #   Returns:
  #     A dict containing additional Python objects to be checkpointed by the
  #       experiment. If the checkpoint directory does not exist, returns None.
  #   """
  #   # if not gfile.Exists(checkpoint_dir):
  #   #   return None
  #   # Call the Tensorflow saver to checkpoint the graph.
  #   self._saver.save(
  #       self._sess,
  #       os.path.join(checkpoint_dir, 'tf_ckpt'),
  #       global_step=iteration_number)
  #   # Checkpoint the out-of-graph replay buffer.
  #   self._replay.save(checkpoint_dir, iteration_number)
  #   bundle_dictionary = {}
  #   bundle_dictionary['state'] = self.state
  #   bundle_dictionary['training_steps'] = self.training_steps
  #   return bundle_dictionary
  #
  # def unbundle(self, checkpoint_dir, iteration_number, bundle_dictionary):
  #   """Restores the agent from a checkpoint.
  #
  #   Restores the agent's Python objects to those specified in bundle_dictionary,
  #   and restores the TensorFlow objects to those specified in the
  #   checkpoint_dir. If the checkpoint_dir does not exist, will not reset the
  #     agent's state.
  #
  #   Args:
  #     checkpoint_dir: str, path to the checkpoint saved by tf.Save.
  #     iteration_number, checkpoint version, used when restoring the replay
  #       buffer.
  #     bundle_dictionary: dict, containing additional Python objects owned by
  #       the agent.
  #
  #   Returns:
  #     bool, True if unbundling was successful.
  #   """
  #   try:
  #     # self._replay.load() will throw a NotFoundError if it does not find all
  #     # the necessary files.
  #     self._replay.load(checkpoint_dir, iteration_number)
  #   except tf.errors.NotFoundError:
  #     if not self.allow_partial_reload:
  #       # If we don't allow partial reloads, we will return False.
  #       return False
  #     logging.warning('Unable to reload replay buffer!')
  #   if bundle_dictionary is not None:
  #     for key in self.__dict__:
  #       if key in bundle_dictionary:
  #         self.__dict__[key] = bundle_dictionary[key]
  #   elif not self.allow_partial_reload:
  #     return False
  #   else:
  #     logging.warning("Unable to reload the agent's parameters!")
  #   # Restore the agent's TensorFlow graph.
  #   self._saver.restore(self._sess,
  #                       os.path.join(checkpoint_dir,
  #                                    'tf_ckpt-{}'.format(iteration_number)))
  #   return True
