import inspect

import numpy as np
import torch
from gymnasium import spaces
from rlberry import types
from rlberry.agents import AgentWithSimplePolicy, AgentTorch
from rlberry.agents.torch.utils.training import (
    loss_function_factory,
    model_factory,
    optimizer_factory,
    size_model_config,
)
from rlberry.agents.torch.dqn.dqn_utils import (
    lambda_returns,
    polynomial_schedule,
    stable_scaled_log_softmax,
    stable_softmax,
)
from rlberry.agents.utils import replay
from rlberry.utils.torch import choose_device
from rlberry.utils.factory import load
from typing import Callable, Optional, Union


import rlberry

logger = rlberry.logger


def default_q_net_fn(env, **kwargs):
    """
    Returns a default Q value network.
    """
    del kwargs
    model_config = {
        "type": "MultiLayerPerceptron",
        "layer_sizes": (64, 64),
        "reshape": False,
    }
    model_config = size_model_config(env, **model_config)
    return model_factory(**model_config)


class MunchausenDQNAgent(AgentTorch, AgentWithSimplePolicy):
    """Munchausen DQN Agent based on PyTorch.

    Notes
    -----
    Uses Munchausen trick for DQN for computing targets by default.
    Compared to DQN, the scaled log-policy was added to the immediate
    reward. Slightly modifying DQN in that way provides an agent that
    is competitive with distributional methods on Atari games, without
    making use of distributional RL, n-step returns or prioritized replay.
    See more: https://arxiv.org/pdf/2007.14430.pdf

    Parameters
    ----------
    env: :class:`~rlberry.types.Env`
        Environment, can be a tuple (constructor, kwargs)
    gamma: float, default = 0.99
        Discount factor.
    batch_size: int, default=32
        Batch size.
    chunk_size: int, default=8
        Length of sub-trajectories sampled from the replay buffer.
    lambda_: float, default=0.5
        Q(lambda) parameter.
    tau: float, default=0.03
        softmax temperature for the policy
    alpha: float, default=0.9
        Munchausen coefficient
    target_update_parameter : int or float
        If int: interval (in number total number of online updates) between updates of the target network.
        If float: soft update coefficient
    device: str
        Torch device, see :func:`~rlberry.utils.torch.choose_device`
    learning_rate : float, default = 1e-3
        Optimizer learning rate.
    clip_value_min: float, default = -1,
        minimum value for munchausen term
    loss_function: {"l1", "l2", "smooth_l1"}, default: "l2"
        Loss function used to compute Bellman error.
    epsilon_init: float, default = 1.0
        Initial epsilon value for epsilon-greedy exploration.
    epsilon_final: float, default = 0.1
        Final epsilon value for epsilon-greedy exploration.
    epsilon_decay_interval : int
        After :code:`epsilon_decay` timesteps, epsilon approaches :code:`epsilon_final`.
    optimizer_type : {"ADAM", "RMS_PROP"}
        Optimization algorithm.
    q_net_constructor : Callable, str or None
        Function/constructor that returns a torch module for the Q-network:
        :code:`qnet = q_net_constructor(env, **kwargs)`.

        Module (Q-network) requirements:

        * Input shape = (batch_dim, chunk_size, obs_dims)

        * Ouput shape = (batch_dim, chunk_size, number_of_actions)

        Example: use `rlberry.agents.torch.utils.training.model_factory_from_env`,
         and `q_net_kwargs`
        parameter to modify the neural network::

            model_configs = {
                "type": "MultiLayerPerceptron",
                "layer_sizes": (5, 5),
                "reshape": False,
            }

            agent = MunchausenDQNAgent(env,
                q_net_constructor=model_factory_from_env,
                q_net_kwargs=model_configs
                )
        If str then it should correspond to the full path to the constructor function,
        e.g.::
            agent = MunchausenDQNAgent(env,
                q_net_constructor='rlberry.agents.torch.utils.training.model_factory_from_env',
                q_net_kwargs=model_configs
                )

        If None then it is set to MultiLayerPerceptron with 2 hidden layers
        of size 64

    q_net_kwargs : optional, dict
        Parameters for q_net_constructor.
    use_prioritized_replay : bool, default = False
        If True, use Prioritized Experience Replay.
    train_interval: int
        Update the model every :code:`train_interval` steps.
        If -1, train only at the end of the episodes.
    gradient_steps: int
        How many gradient steps to do at each update.
        If -1, take the number of timesteps since last update.
    max_replay_size : int
        Maximum number of transitions in the replay buffer.
    learning_starts : int
        How many steps of the model to collect transitions for before learning starts
    eval_interval : int, default = None
        Interval (in number of transitions) between agent evaluations in fit().
        If None, never evaluate.
    """

    name = "Munchausen DQN"

    def __init__(
        self,
        env: types.Env,
        gamma: float = 0.99,
        batch_size: int = 32,
        chunk_size: int = 8,
        lambda_: float = 0.5,
        tau: float = 0.03,
        alpha: float = 0.9,
        target_update_parameter: Union[int, float] = 0.005,
        # tardet_update_freq: int = 8000,
        device: str = "cuda:best",
        learning_rate: float = 5e-5,
        clip_value_min: float = -1.0,
        epsilon_init: float = 1.0,
        epsilon_final: float = 0.1,
        epsilon_decay_interval: int = 20_000,
        loss_function: str = "l2",
        optimizer_type: str = "ADAM",
        q_net_constructor: Optional[Callable[..., torch.nn.Module]] = None,
        q_net_kwargs: Optional[dict] = None,
        use_prioritized_replay: bool = False,
        train_interval: int = 4,
        gradient_steps: int = -1,
        max_replay_size: int = 1_000_000,
        learning_starts: int = 5_000,
        eval_interval: Optional[int] = None,
        **kwargs,
    ):
        # For all parameters, define self.param = param
        _, _, _, values = inspect.getargvalues(inspect.currentframe())
        values.pop("self")
        for arg, val in values.items():
            setattr(self, arg, val)

        AgentWithSimplePolicy.__init__(self, env, **kwargs)
        env = self.env
        assert isinstance(env.observation_space, spaces.Box)
        assert isinstance(env.action_space, spaces.Discrete)

        # M-DQN parameters

        # Online and target Q networks, torch device
        self._device = choose_device(device)
        if isinstance(q_net_constructor, str):
            q_net_ctor = load(q_net_constructor)
        elif q_net_constructor is None:
            q_net_ctor = default_q_net_fn
        else:
            q_net_ctor = q_net_constructor
        q_net_kwargs = q_net_kwargs or dict()
        self._qnet_online = q_net_ctor(env, **q_net_kwargs).to(self._device)
        self._qnet_target = q_net_ctor(env, **q_net_kwargs).to(self._device)

        # Optimizer and loss
        optimizer_kwargs = {"optimizer_type": optimizer_type, "lr": learning_rate}
        self._optimizer = optimizer_factory(
            self._qnet_online.parameters(), **optimizer_kwargs
        )
        self._loss_function = loss_function_factory(loss_function, reduction="none")

        # Training params
        self._train_interval = train_interval
        self._gradient_steps = gradient_steps
        self._learning_starts = learning_starts
        self._learning_starts = learning_starts
        self._eval_interval = eval_interval

        # Setup replay buffer
        if hasattr(self.env, "_max_episode_steps"):
            max_episode_steps = self.env._max_episode_steps
        else:
            max_episode_steps = np.inf
        self._max_episode_steps = max_episode_steps

        self._replay_buffer = replay.ReplayBuffer(
            max_replay_size=max_replay_size,
            rng=self.rng,
            max_episode_steps=self._max_episode_steps,
            enable_prioritized=use_prioritized_replay,
        )
        self._replay_buffer.setup_entry("observations", np.float32)
        self._replay_buffer.setup_entry("next_observations", np.float32)
        self._replay_buffer.setup_entry("actions", np.int32)
        self._replay_buffer.setup_entry("rewards", np.float32)
        self._replay_buffer.setup_entry("dones", bool)

        # Counters
        self._total_timesteps = 0
        self._total_episodes = 0
        self._total_updates = 0
        self._timesteps_since_last_update = 0

        # epsilon scheduling
        self._epsilon_schedule = polynomial_schedule(
            self.epsilon_init,
            self.epsilon_final,
            power=1.0,
            transition_steps=self.epsilon_decay_interval,
            transition_begin=0,
        )

    @property
    def total_timesteps(self):
        return self._total_timesteps

    def _must_update(self, is_end_of_episode):
        """Returns true if the model must be updated in the current timestep,
        and the number of gradient steps to take"""
        total_timesteps = self._total_timesteps
        n_gradient_steps = self._gradient_steps

        if total_timesteps < self._learning_starts:
            return False, -1

        if n_gradient_steps == -1:
            n_gradient_steps = self._timesteps_since_last_update

        run_update = False
        if self._train_interval == -1:
            run_update = is_end_of_episode
        else:
            run_update = total_timesteps % self._train_interval == 0
        return run_update, n_gradient_steps

    def _update(self, n_gradient_steps):
        """Update networks."""
        if self.use_prioritized_replay:
            sampling_mode = "prioritized"
        else:
            sampling_mode = "uniform"

        for _ in range(n_gradient_steps):
            # Sample a batch
            sampled = self._replay_buffer.sample(
                self.batch_size, self.chunk_size, sampling_mode=sampling_mode
            )
            if not sampled:
                return

            # Update counters
            self._timesteps_since_last_update = 0
            self._total_updates += 1

            batch = sampled.data
            batch_info = sampled.info
            assert batch["rewards"].shape == (self.batch_size, self.chunk_size)

            # Get batched tensors
            batch_observations = torch.FloatTensor(batch["observations"]).to(
                self._device
            )
            batch_rewards = torch.FloatTensor(batch["rewards"]).to(self._device)
            batch_next_observations = torch.FloatTensor(batch["next_observations"]).to(
                self._device
            )
            batch_actions = torch.LongTensor(batch["actions"]).to(self._device)
            batch_dones = torch.LongTensor(batch["dones"]).to(self._device)

            # Get target Q estimates
            target_q_values_tp1 = self._qnet_target(batch_next_observations).detach()
            target_q_values = self._qnet_target(batch_observations).detach()

            # Compute softmax policies for the current and next step
            log_pi = stable_scaled_log_softmax(target_q_values, self.tau, -1)
            log_pi_tp1 = stable_scaled_log_softmax(target_q_values_tp1, self.tau, -1)
            pi_tp1 = stable_softmax(target_q_values_tp1, self.tau, -1)

            # Compute the "next step" part of the target
            target_v_tp1 = (
                torch.sum((target_q_values_tp1 - log_pi_tp1) * pi_tp1, -1).cpu().numpy()
            )

            # Compute the Munchausen term
            munchausen_term = torch.gather(
                log_pi, dim=-1, index=batch_actions[:, :, None]
            )[:, :, 0]
            clipped_munchausen_term = torch.clip(
                munchausen_term, self.clip_value_min, 0
            )
            final_munchausen_term = self.alpha * clipped_munchausen_term

            # Compute the final target
            batch_lambda_returns = lambda_returns(
                (batch_rewards + final_munchausen_term).cpu().numpy(),
                self.gamma * (1.0 - np.array(batch["dones"], dtype=np.float32)),
                target_v_tp1,
                np.array(self.lambda_, dtype=np.float32),
            )
            targets = torch.tensor(batch_lambda_returns, device=self._device)

            # Compute loss
            batch_q_values = self._qnet_online(batch_observations)
            batch_values = torch.gather(
                batch_q_values, dim=-1, index=batch_actions[:, :, None]
            )[
                :, :, 0
            ]  # shape (batch, chunk)

            assert batch_values.shape == targets.shape
            per_element_loss = self._loss_function(batch_values, targets)
            per_batch_element_loss = per_element_loss.mean(dim=1)
            weights = torch.FloatTensor(batch_info["weights"]).to(self._device)
            loss = torch.sum(per_batch_element_loss * weights) / torch.sum(weights)

            self._optimizer.zero_grad()
            loss.backward()
            self._optimizer.step()

            if self.writer:
                self.writer.add_scalar(
                    "losses/q_loss", loss.item(), self._total_updates
                )

            # update priorities
            if self.use_prioritized_replay:
                new_priorities = per_element_loss.abs().detach().cpu().numpy() + 1e-6
                self._replay_buffer.update_priorities(
                    batch_info["indices"], new_priorities
                )

            # target update

            if self.target_update_parameter > 1:
                if self._total_updates % self.target_update_parameter == 0:
                    self._qnet_target.load_state_dict(self._qnet_online.state_dict())
            else:
                tau = self.target_update_parameter
                for param, target_param in zip(
                    self._qnet_online.parameters(), self._qnet_target.parameters()
                ):
                    target_param.data.copy_(
                        tau * param.data + (1 - tau) * target_param.data
                    )

    def fit(self, budget: int, **kwargs):
        """
        Train the agent using the provided environment.

        Parameters
        ----------
        budget: int
            Number of timesteps to train the agent for.
            One step = one transition in the environment.
        """
        del kwargs
        timesteps_counter = 0
        episode_rewards = 0.0
        episode_timesteps = 0
        observation, info = self.env.reset()
        while timesteps_counter < budget:
            if self.total_timesteps < self._learning_starts:
                action = self.env.action_space.sample()
            else:
                self._timesteps_since_last_update += 1
                action = self._policy(observation, evaluation=False)
            next_observation, reward, terminated, truncated, info = self.env.step(
                action
            )
            done = terminated or truncated

            # store data
            episode_rewards += reward
            self._replay_buffer.append(
                {
                    "observations": observation,
                    "actions": action,
                    "rewards": reward,
                    "dones": done,
                    "next_observations": next_observation,
                }
            )

            # counters and next obs
            self._total_timesteps += 1
            timesteps_counter += 1
            episode_timesteps += 1
            observation = next_observation

            # update
            run_update, n_gradient_steps = self._must_update(done)
            if run_update:
                self._update(n_gradient_steps)

            # eval
            total_timesteps = self._total_timesteps
            if (
                self._eval_interval is not None
                and total_timesteps % self._eval_interval == 0
            ):
                eval_rewards = self.eval(
                    eval_horizon=self._max_episode_steps, gamma=1.0
                )
                if self.writer:
                    buffer_size = len(self._replay_buffer)
                    self.writer.add_scalar(
                        "eval_rewards", eval_rewards, total_timesteps
                    )
                    self.writer.add_scalar("buffer_size", buffer_size, total_timesteps)

            # check if episode ended
            if done:
                self._total_episodes += 1
                self._replay_buffer.end_episode()
                if self.writer:
                    self.writer.add_scalar(
                        "episode_rewards", episode_rewards, total_timesteps
                    )
                    self.writer.add_scalar(
                        "total_episodes", self._total_episodes, total_timesteps
                    )
                episode_rewards = 0.0
                episode_timesteps = 0
                observation, info = self.env.reset()

    def _policy(self, observation, evaluation=False):
        epsilon = self._epsilon_schedule(self.total_timesteps)
        if (not evaluation) and self.rng.uniform() < epsilon:
            if self.writer:
                self.writer.add_scalar("epsilon", epsilon, self.total_timesteps)
            return self.env.action_space.sample()
        else:
            with torch.no_grad():
                observation = (
                    torch.FloatTensor(observation).to(self._device).unsqueeze(0)
                )
                qvals_tensor = self._qnet_online(observation)[0]
                action = qvals_tensor.argmax().item()
                return action

    def policy(self, observation):
        return self._policy(observation, evaluation=True)
