from typing import Union, Dict, Optional, List
from easydict import EasyDict
import numpy as np
import torch
import torch.nn as nn

from ding.utils import SequenceType, squeeze, MODEL_REGISTRY
from ..common import RegressionHead, ReparameterizationHead
from .vae import VanillaVAE


@MODEL_REGISTRY.register('bcq')
class BCQ(nn.Module):
    """
    Overview:
        Model of BCQ (Batch-Constrained deep Q-learning).
        Off-Policy Deep Reinforcement Learning without Exploration.
        https://arxiv.org/abs/1812.02900
    Interface:
        ``forward``, ``compute_actor``, ``compute_critic``, ``compute_vae``, ``compute_eval``
    Property:
        ``mode``
    """

    mode = ['compute_actor', 'compute_critic', 'compute_vae', 'compute_eval']

    def __init__(
            self,
            obs_shape: Union[int, SequenceType],
            action_shape: Union[int, SequenceType, EasyDict],
            actor_head_hidden_size: List = [400, 300],
            critic_head_hidden_size: List = [400, 300],
            activation: Optional[nn.Module] = nn.ReLU(),
            vae_hidden_dims: List = [750, 750],
            phi: float = 0.05
    ) -> None:
        """
        Overview:
            Initialize neural network, i.e. agent Q network and actor.
        Arguments:
            - obs_shape (:obj:`int`): the dimension of observation state
            - action_shape (:obj:`int`): the dimension of action shape
            - actor_hidden_size (:obj:`list`): the list of hidden size of actor
            - critic_hidden_size (:obj:'list'): the list of hidden size of critic
            - activation (:obj:`nn.Module`): Activation function in network, defaults to nn.ReLU().
            - vae_hidden_dims (:obj:`list`): the list of hidden size of vae
        """
        super(BCQ, self).__init__()
        obs_shape: int = squeeze(obs_shape)
        action_shape = squeeze(action_shape)
        self.action_shape = action_shape
        self.input_size = obs_shape
        self.phi = phi

        critic_input_size = self.input_size + action_shape
        self.critic = nn.ModuleList()
        for _ in range(2):
            net = []
            d = critic_input_size
            for dim in critic_head_hidden_size:
                net.append(nn.Linear(d, dim))
                net.append(activation)
                d = dim
            net.append(nn.Linear(d, 1))
            self.critic.append(nn.Sequential(*net))

        net = []
        d = critic_input_size
        for dim in actor_head_hidden_size:
            net.append(nn.Linear(d, dim))
            net.append(activation)
            d = dim
        net.append(nn.Linear(d, 1))
        self.actor = nn.Sequential(*net)

        self.vae = VanillaVAE(action_shape, obs_shape, action_shape * 2, vae_hidden_dims)

    def forward(self, inputs: Dict[str, torch.Tensor], mode: str) -> Dict[str, torch.Tensor]:
        """
        Overview:
            The unique execution (forward) method of BCQ method, and one can indicate different modes to implement \
            different computation graph, including ``compute_actor`` and ``compute_critic`` in BCQ.
        Mode compute_actor:
            Arguments:
                - inputs (:obj:`Dict`): Input dict data, including obs and action tensor.
            Returns:
                - output (:obj:`Dict`): Output dict data, including action tensor.
        Mode compute_critic:
            Arguments:
                - inputs (:obj:`Dict`): Input dict data, including obs and action tensor.
            Returns:
                - output (:obj:`Dict`): Output dict data, including q_value tensor.
        Mode compute_vae:
            Arguments:
                - inputs (:obj:`Dict`): Input dict data, including obs and action tensor.
            Returns:
                - outputs (:obj:`Dict`): Dict containing keywords ``recons_action`` \
                (:obj:`torch.Tensor`), ``prediction_residual`` (:obj:`torch.Tensor`), \
                ``input`` (:obj:`torch.Tensor`), ``mu`` (:obj:`torch.Tensor`), \
                ``log_var`` (:obj:`torch.Tensor`) and ``z`` (:obj:`torch.Tensor`).
        Mode compute_eval:
            Arguments:
                - inputs (:obj:`Dict`): Input dict data, including obs and action tensor.
            Returns:
                - output (:obj:`Dict`): Output dict data, including action tensor.
        Examples:
            >>> inputs = {'obs': torch.randn(4, 32), 'action': torch.randn(4, 6)}
            >>> model = BCQ(32, 6)
            >>> outputs = model(inputs, mode='compute_actor')
            >>> outputs = model(inputs, mode='compute_critic')
            >>> outputs = model(inputs, mode='compute_vae')
            >>> outputs = model(inputs, mode='compute_eval')

        .. note::
            For specific examples, one can refer to API doc of ``compute_actor`` and ``compute_critic`` respectively.
        """
        assert mode in self.mode, "not support forward mode: {}/{}".format(mode, self.mode)
        return getattr(self, mode)(inputs)

    def compute_critic(self, inputs: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
        """
        Overview:
            Use critic network to compute q value.
        Arguments:
            - inputs (:obj:`Dict`): Input dict data, including obs and action tensor.
        Returns:
            - outputs (:obj:`Dict`): Dict containing keywords ``q_value`` (:obj:`torch.Tensor`).
        Shapes:
            - inputs (:obj:`Dict`): :math:`(B, N, D)`, where B is batch size, N is sample number, D is input dimension.
            - outputs (:obj:`Dict`): :math:`(B, N)`.
        Examples:
            >>> inputs = {'obs': torch.randn(4, 32), 'action': torch.randn(4, 6)}
            >>> model = BCQ(32, 6)
            >>> outputs = model.compute_critic(inputs)
        """
        obs, action = inputs['obs'], inputs['action']
        if len(action.shape) == 1:  # (B, ) -> (B, 1)
            action = action.unsqueeze(1)
        x = torch.cat([obs, action], dim=-1)
        x = [m(x).squeeze() for m in self.critic]
        return {'q_value': x}

    def compute_actor(self, inputs: Dict[str, torch.Tensor]) -> Dict[str, Union[torch.Tensor, Dict[str, torch.Tensor]]]:
        """
        Overview:
            Use actor network to compute action.
        Arguments:
            - inputs (:obj:`Dict`): Input dict data, including obs and action tensor.
        Returns:
            - outputs (:obj:`Dict`): Dict containing keywords ``action`` (:obj:`torch.Tensor`).
        Shapes:
            - inputs (:obj:`Dict`): :math:`(B, N, D)`, where B is batch size, N is sample number, D is input dimension.
            - outputs (:obj:`Dict`): :math:`(B, N)`.
        Examples:
            >>> inputs = {'obs': torch.randn(4, 32), 'action': torch.randn(4, 6)}
            >>> model = BCQ(32, 6)
            >>> outputs = model.compute_actor(inputs)
        """
        input = torch.cat([inputs['obs'], inputs['action']], -1)
        x = self.actor(input)
        action = self.phi * 1 * torch.tanh(x)
        action = (action + inputs['action']).clamp(-1, 1)
        return {'action': action}

    def compute_vae(self, inputs: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
        """
        Overview:
            Use vae network to compute action.
        Arguments:
            - inputs (:obj:`Dict`): Input dict data, including obs and action tensor.
        Returns:
            - outputs (:obj:`Dict`): Dict containing keywords ``recons_action`` (:obj:`torch.Tensor`), \
                ``prediction_residual`` (:obj:`torch.Tensor`), ``input`` (:obj:`torch.Tensor`), \
                ``mu`` (:obj:`torch.Tensor`), ``log_var`` (:obj:`torch.Tensor`) and ``z`` (:obj:`torch.Tensor`).
        Shapes:
            - inputs (:obj:`Dict`): :math:`(B, N, D)`, where B is batch size, N is sample number, D is input dimension.
            - outputs (:obj:`Dict`): :math:`(B, N)`.
        Examples:
            >>> inputs = {'obs': torch.randn(4, 32), 'action': torch.randn(4, 6)}
            >>> model = BCQ(32, 6)
            >>> outputs = model.compute_vae(inputs)
        """
        return self.vae.forward(inputs)

    def compute_eval(self, inputs: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
        """
        Overview:
            Use actor network to compute action.
        Arguments:
            - inputs (:obj:`Dict`): Input dict data, including obs and action tensor.
        Returns:
            - outputs (:obj:`Dict`): Dict containing keywords ``action`` (:obj:`torch.Tensor`).
        Shapes:
            - inputs (:obj:`Dict`): :math:`(B, N, D)`, where B is batch size, N is sample number, D is input dimension.
            - outputs (:obj:`Dict`): :math:`(B, N)`.
        Examples:
            >>> inputs = {'obs': torch.randn(4, 32), 'action': torch.randn(4, 6)}
            >>> model = BCQ(32, 6)
            >>> outputs = model.compute_eval(inputs)
        """
        obs = inputs['obs']
        obs_rep = obs.clone().unsqueeze(0).repeat_interleave(100, dim=0)
        z = torch.randn((obs_rep.shape[0], obs_rep.shape[1], self.action_shape * 2)).to(obs.device).clamp(-0.5, 0.5)
        sample_action = self.vae.decode_with_obs(z, obs_rep)['reconstruction_action']
        action = self.compute_actor({'obs': obs_rep, 'action': sample_action})['action']
        q = self.compute_critic({'obs': obs_rep, 'action': action})['q_value'][0]
        idx = q.argmax(dim=0).unsqueeze(0).unsqueeze(-1)
        idx = idx.repeat_interleave(action.shape[-1], dim=-1)
        action = action.gather(0, idx).squeeze()
        return {'action': action}