import torch
import numpy as np
import torch.nn.functional as F
from torch.autograd import Variable
from basenets.MLP import MLP
from basenets.Conv import Conv
from torch import nn

class FCPG_Gaussian(MLP):
    def __init__(self,
                 n_inputfeats,
                 n_actions,
                 sigma,
                 n_hiddens = [30],
                 nonlinear = F.tanh,
                 usebn = False,
                 outactive = None,
                 outscaler = None,
                 initializer = "orthogonal",
                 initializer_param = {"gain":np.sqrt(2), "last_gain": 0.1}
                 ):
        self.n_actions = n_actions
        super(FCPG_Gaussian, self).__init__(
            n_inputfeats,  # input dim
            n_actions,  # output dim
            n_hiddens,  # hidden unit number list
            nonlinear,
            usebn,
            outactive,
            outscaler,
            initializer,
            initializer_param=initializer_param,
        )
        self.logstd = nn.Parameter(torch.log(sigma * torch.ones(n_actions) + 1e-8))

    def forward(self,x, other_data = None):
        x = MLP.forward(self, x, other_data)
        # for exploration, we need to make sure that the std is not too low.
        logstd = torch.clamp(self.logstd, min = np.log(0.1))
        return x, logstd.expand_as(x), torch.exp(logstd).expand_as(x)

    def cuda(self, device = None):
        self.logstd.cuda()
        return self._apply(lambda t: t.cuda(device))

class FCPG_Softmax(MLP):
    def __init__(self,
                 n_inputfeats,    # input dim
                 n_actions,   # output dim
                 n_hiddens = [10],  # hidden unit number list
                 nonlinear = F.tanh,
                 usebn = False,
                 outactive = F.softmax,
                 outscaler = None,
                 initializer = "orthogonal",
                 initializer_param = {"gain":np.sqrt(2), "last_gain": 0.1}
                 ):
        self.n_actions = n_actions
        super(FCPG_Softmax, self).__init__(
                 n_inputfeats,    # input dim
                 n_actions,   # output dim
                 n_hiddens,  # hidden unit number list
                 nonlinear,
                 usebn,
                 outactive,
                 outscaler,
                 initializer,
                 initializer_param=initializer_param,
                 )

    def forward(self, x, other_data=None):
        x = MLP.forward(self, x, other_data)
        # for exploration, and similar to e-greedy
        x = x + 0.01 / self.n_actions
        x = x / torch.sum(x, dim = -1, keepdim=True).detach()
        return x

class ConvPG_Softmax(Conv):
    def __init__(self,
                 n_inputfeats,    # input dim
                 n_actions,   # output dim
                 k_sizes = [8, 4, 3],
                 channels = [8, 16, 16],
                 strides = [4, 2, 2],
                 fcs = [32, 32, 32],  # hidden unit number list
                 nonlinear = F.relu,
                 usebn = False,
                 outactive = F.softmax,
                 outscaler = None,
                 initializer="xavier",
                 initializer_param={}
                 ):
        self.n_actions = n_actions
        super(ConvPG_Softmax, self).__init__(
                 n_inputfeats,    # input dim
                 n_actions,   # output dim
                 k_sizes,
                 channels,
                 strides,
                 fcs,
                 nonlinear,
                 usebn,
                 outactive,
                 outscaler,
                 initializer,
                 initializer_param=initializer_param,
                 )

    def forward(self, x, other_data=None):
        x = Conv.forward(self, x, other_data)
        # for exploration, and similar to e-greedy
        x = x + 0.01 / self.n_actions
        x = x / torch.sum(x, dim=-1, keepdim=True).detach()
        return x

# TODO: support multi-layer value function in which action is concat before the final layer
class FCVALUE(MLP):
    def __init__(self,
                 n_inputfeats,
                 n_hiddens = [30],
                 nonlinear = F.tanh,
                 usebn = False,
                 outactive = None,
                 outscaler = None,
                 initializer="orthogonal",
                 initializer_param={"gain":np.sqrt(2), "last_gain": 0.1}
                 ):
        super(FCVALUE, self).__init__(
                 n_inputfeats,
                 1,
                 n_hiddens,
                 nonlinear,
                 usebn,
                 outactive,
                 outscaler,
                 initializer,
                 initializer_param=initializer_param,
                 )


[((37, 12, 37, 44), 'basenets.MLP.MLP.forward', 'MLP.forward', ({(37, 24, 37, 28): 'self', (37, 30, 37, 31): 'x', (37, 33, 37, 43): 'other_data'}, {}), '(self, x, other_data)', False, 'from basenets.MLP import MLP\n'), ((72, 12, 72, 44), 'basenets.MLP.MLP.forward', 'MLP.forward', ({(72, 24, 72, 28): 'self', (72, 30, 72, 31): 'x', (72, 33, 72, 43): 'other_data'}, {}), '(self, x, other_data)', False, 'from basenets.MLP import MLP\n'), ((110, 12, 110, 45), 'basenets.Conv.Conv.forward', 'Conv.forward', ({(110, 25, 110, 29): 'self', (110, 31, 110, 32): 'x', (110, 34, 110, 44): 'other_data'}, {}), '(self, x, other_data)', False, 'from basenets.Conv import Conv\n'), ((20, 45, 20, 55), 'numpy.sqrt', 'np.sqrt', ({(20, 53, 20, 54): '(2)'}, {}), '(2)', True, 'import numpy as np\n'), ((56, 45, 56, 55), 'numpy.sqrt', 'np.sqrt', ({(56, 53, 56, 54): '(2)'}, {}), '(2)', True, 'import numpy as np\n'), ((126, 43, 126, 53), 'numpy.sqrt', 'np.sqrt', ({(126, 51, 126, 52): '(2)'}, {}), '(2)', True, 'import numpy as np\n'), ((39, 48, 39, 59), 'numpy.log', 'np.log', ({(39, 55, 39, 58): '0.1'}, {}), '(0.1)', True, 'import numpy as np\n'), ((40, 39, 40, 56), 'torch.exp', 'torch.exp', ({(40, 49, 40, 55): 'logstd'}, {}), '(logstd)', False, 'import torch\n'), ((75, 16, 75, 52), 'torch.sum', 'torch.sum', (), '', False, 'import torch\n'), ((113, 16, 113, 50), 'torch.sum', 'torch.sum', (), '', False, 'import torch\n'), ((34, 53, 34, 74), 'torch.ones', 'torch.ones', ({(34, 64, 34, 73): 'n_actions'}, {}), '(n_actions)', False, 'import torch\n')]