"""DQN Agent Class

Heavily inspired by:
- https://github.com/deepmind/dqn_zoo/blob/master/dqn_zoo/dqn/agent.py
"""
from typing import Callable, Optional, Mapping, Union

import jax
import jax.numpy as jnp
import chex
import rlax
import optax
import numpy as np
from absl import logging

from rljaxff import parts
from rljaxff import loss as loss_lib
from rljaxff import replay as replay_lib

Numeric = Union[jnp.ndarray, float, int]


class DQN(parts.Agent):
    """DQN Agent """

    def __init__(self,
                 sample_network_input,
                 network: parts.Network,
                 optimizer: optax.GradientTransformation,
                 transition_accumulator: replay_lib.TransitionAccumulator,
                 replay: replay_lib.TransitionReplay,
                 batch_size: int,
                 exploration_epsilon: Callable[[int], float],
                 min_replay_capacity_fraction: float,
                 learn_period: int,
                 target_network_update_period: int,
                 grad_error_bound: float,
                 rng_key: parts.PRNGKey):
        self._transition_accumulator = transition_accumulator
        self._replay = replay
        self._batch_size = batch_size
        self._exploration_epsilon = exploration_epsilon
        self._min_replay_capacity = (
            replay.capacity * min_replay_capacity_fraction
        )
        self._learn_period = learn_period
        self._target_network_update_period = target_network_update_period

        # Initialize network parameters and optimizer
        self._rng_key, net_rng_key = jax.random.split(rng_key)  # type: ignore
        self._online_params = network.init(
            net_rng_key, sample_network_input[None, ...]
        )
        self._target_params = self._online_params
        self._opt_state = optimizer.init(self._online_params)

        # Other agent state
        self._action: Optional[parts.Action] = None
        self._frame_t = -1
        self._statistics = {
            'state_value': np.nan,
            'loss_value': np.nan
        }

        def double_q_learning(q_tm1: jnp.ndarray,
                              a_tm1: Numeric,
                              r_t: Numeric,
                              discount_t: Numeric,
                              q_target_t: jnp.ndarray,
                              q_t: jnp.ndarray):
            target_tm1 = r_t + discount_t * q_target_t[q_t.argmax()]
            target_tm1 = jax.lax.stop_gradient(target_tm1)
            return target_tm1 - q_tm1[a_tm1]

        self._batched_q_learning = jax.vmap(double_q_learning)

        def loss_fn(online_params: parts.NetworkParams,
                    target_params: parts.NetworkParams,
                    transitions,
                    rng_key: parts.PRNGKey):
            _, *apply_keys = jax.random.split(rng_key, 4)
            q_tm1 = network.apply(
                online_params, rng=apply_keys[0], inputs=transitions.s_tm1
            )
            q_t = network.apply(
                online_params, rng=apply_keys[1], inputs=transitions.s_t
            )
            q_target_t = network.apply(
                target_params, rng=apply_keys[2], inputs=transitions.s_t
            )

            td_errors = self._batched_q_learning(
                q_tm1,
                transitions.a_tm1,
                transitions.r_t,
                transitions.discount_t,
                q_target_t,
                q_t
            )

            # td_errors = rlax.clip_gradient(
            #     td_errors, -grad_error_bound, grad_error_bound
            # )
            # losses = rlax.l2_loss(td_errors)
            losses = loss_lib.huber_loss(td_errors)

            chex.assert_shape(losses, (self._batch_size, ))
            loss = jnp.mean(losses)
            return loss

        def update(rng_key: parts.PRNGKey,
                   opt_state,
                   online_params: parts.NetworkParams,
                   target_params: parts.NetworkParams,
                   transitions):
            """Computes learning update from batch of transitions """
            rng_key, update_key = jax.random.split(rng_key)  # type: ignore
            loss_value, d_loss_d_params = jax.value_and_grad(loss_fn)(
                online_params, target_params, transitions, update_key
            )

            updates, new_opt_state = optimizer.update(
                d_loss_d_params, opt_state
            )
            new_online_params = optax.apply_updates(online_params, updates)

            return rng_key, new_opt_state, new_online_params, loss_value

        self._update = jax.jit(update)

        def select_action(rng_key: parts.PRNGKey,
                          network_params: parts.NetworkParams,
                          s_t: parts.Observation,
                          exploration_epsilon: float):
            """Samples action from eps-greedy policy """
            rng_keys = jax.random.split(rng_key, 3)
            q_t = network.apply(
                network_params, rng_keys[1], s_t[None, ...]
            )[0]
            a_t = rlax.epsilon_greedy().sample(
                rng_keys[2], q_t, exploration_epsilon
            )
            v_t = jnp.max(q_t, axis=-1)
            return rng_keys[0], a_t, v_t, q_t

        self._select_action = jax.jit(select_action)

    def step(self, timestep: parts.Timestep) -> parts.Action:
        """Selects action and potentially learns. """
        self._frame_t += 1

        action = self._action = self._act(timestep)

        transitions = self._transition_accumulator.step(timestep, action)
        for transition in transitions:
            self._replay.add(transition)

        if self._replay.size < self._min_replay_capacity:
            return action

        if self._frame_t % self._learn_period == 0:
            self._learn()

        if self._frame_t % self._target_network_update_period == 0:
            logging.log_first_n(
                logging.INFO,
                'First target network update frame %d',
                1,
                self._frame_t
            )
            self._target_params = self._online_params

        return action

    def reset(self) -> None:
        """Resets the agent's episodic state

        Should be called at the start of each episode.
        """
        self._action = None

    def _act(self, timestep: parts.Timestep) -> parts.Action:
        """Select action given timestep """
        s_t = timestep.observation
        self._rng_key, a_t, v_t, q_t = self._select_action(
            self._rng_key, self._online_params, s_t, self.exploration_epsilon
        )
        a_t, v_t, q_t = jax.device_get((a_t, v_t, q_t))

        self._statistics['state_value'] = v_t
        return parts.Action(a_t)

    def _learn(self):
        """Perform a step of learning from batch of transitions from replay """
        logging.log_first_n(logging.INFO, 'Begin learning', 1)
        transitions = self._replay.sample(self._batch_size)

        self._rng_key, self._opt_state, self._online_params, loss_value = \
            self._update(
                self._rng_key,
                self._opt_state,
                self._online_params,
                self._target_params,
                transitions
            )
        self._statistics['loss_value'] = jax.device_get(loss_value)

    @property
    def statistics(self) -> Mapping[str, float]:
        assert all(
            not isinstance(x, jnp.DeviceArray)
            for x in self._statistics.values()
        )
        return self._statistics

    @property
    def exploration_epsilon(self) -> float:
        """The epsilon for eps-greedy policy given current frame """
        return self._exploration_epsilon(self._frame_t)
