# AUTOGENERATED! DO NOT EDIT! File to edit: ../../../nbs/agents/networks/03_actors.ipynb.

# %% auto 0
__all__ = ['ActorNetwork', 'ActorNetworkDiscrete', 'ActorNetworkRNN', 'ActorNetworkRNN_MLP_hybrid']

# %% ../../../nbs/agents/networks/03_actors.ipynb 4
# Networks
from .base import BaseNetwork

# General libraries
import numpy as np

# Pytorch
import torch
import torch.nn as nn
import torch.nn.functional as F

# %% ../../../nbs/agents/networks/03_actors.ipynb 6
class ActorNetwork(BaseNetwork):
    """
    Simple 3-Layer network for the actor function approximator.
    
    """
    # TODO make the network architecture more flexible

    def __init__(self, input_shape, output_shape, n_features, **kwargs):
        super().__init__(input_shape, output_shape, n_features, **kwargs)

    def forward(self, state):
        
        #features1 = F.relu(self._h1(torch.squeeze(state, 1).float()))
        features1 = F.relu(self._h1(state.float()))
        features2 = F.relu(self._h2(features1))
        a = self._h3(features2)

        return a

# %% ../../../nbs/agents/networks/03_actors.ipynb 7
class ActorNetworkDiscrete(BaseNetwork):
    """
    Simple 3-Layer network for the actor function approximator. Adding a softmax layer to get a probability distribution over the actions.
    
    """
    # TODO make the network architecture more flexible

    def __init__(self, input_shape, output_shape, n_features, **kwargs):
        super().__init__(input_shape, output_shape, n_features, **kwargs)

    def forward(self, state):
        #features1 = F.relu(self._h1(torch.squeeze(state, 1).float()))
        features1 = F.relu(self._h1(state.float()))
        features2 = F.relu(self._h2(features1))
        a = self._h3(features2)
        a = F.softmax(a, dim=1)

        return a

# %% ../../../nbs/agents/networks/03_actors.ipynb 8
class ActorNetworkRNN(nn.Module):
    """
    Simple 3-Layer network for the actor function approximator.
    
    """
    # TODO make the network architecture more flexible

    def __init__(self, input_shape, output_shape, input_shape_RNN, n_features, input_processor = None, device = None, **kwargs):
        super(ActorNetworkRNN, self).__init__()

        # print("setting up actor")
        self._n_input_RNN = input_shape_RNN # first dimension: number of features; second dimension: sequence length
        self._n_output = output_shape[0]
        self._n_features = n_features

        # print("setting up input processor")
        self.input_processor = input_processor

        # print("setting up hi recurrent")
        self._h1_recurrent = nn.GRU(self._n_input_RNN [0], self._n_features, 1, batch_first = True)
        # print("setting up linear layers")
        self._h2 = nn.Linear(self._n_features, self._n_features)
        self._h3 = nn.Linear(self._n_features, self._n_output)

        # print("setting up weights")
        nn.init.xavier_uniform_(self._h2.weight,
                                gain=nn.init.calculate_gain('relu'))
        nn.init.xavier_uniform_(self._h3.weight,
                                gain=nn.init.calculate_gain('linear'))

        self.device = device

        if "cuda" in self.device:
            self.use_cuda = True
        else:
            self.use_cuda = False
        
        # print("device: ", self.device)

        # print("actor set up")
    

    def forward(self, state):

        # print("state before transform:",  state.device)

        state = self.input_processor(state)

        # print("state after transform:",  state.device)

        state_ts = state[0]
        state_static = state[1]

        h0 = torch.zeros(1 , state_ts.size(0), self._n_features).requires_grad_().to(self.device)
        
        features1_rnn, _ = self._h1_recurrent(state_ts.float().to(self.device), h0.detach())
        features1_rnn = features1_rnn[:, -1, :]
        #features1_rnn = torch.zeros(state_ts.size(0), int(self._n_features/2)


        features2 = F.relu(self._h2(features1_rnn)) 

        a = self._h3(features2)

        return a

# %% ../../../nbs/agents/networks/03_actors.ipynb 9
class ActorNetworkRNN_MLP_hybrid(nn.Module):
    """
    Simple 3-Layer network for the actor function approximator.
    
    """
    # TODO make the network architecture more flexible

    def __init__(self, input_shape, output_shape, input_shape_RNN, input_shape_MLP, n_features, input_processor = None, device=None, **kwargs):
        super(ActorNetworkRNN_MLP_hybrid, self).__init__()

        self._n_input_RNN = input_shape_RNN # first dimension: number of features; second dimension: sequence length
        self._n_input_MLP = input_shape_MLP[0]
        self._n_output = output_shape[0]
        self._n_features = n_features

        assert n_features/2 == int(n_features/2), "n_features must be even"

        self.input_processor = input_processor

        self._h1_recurrent = nn.GRU(self._n_input_RNN [0], int(self._n_features/2), 1, batch_first = True)
        self._h1_mlp = nn.Linear(self._n_input_MLP, int(self._n_features/2))
        self._h2 = nn.Linear(self._n_features, self._n_features)
        self._h3 = nn.Linear(self._n_features, self._n_output)

        nn.init.xavier_uniform_(self._h1_mlp.weight,
                                gain=nn.init.calculate_gain('relu'))
        nn.init.xavier_uniform_(self._h2.weight,
                                gain=nn.init.calculate_gain('relu'))
        nn.init.xavier_uniform_(self._h3.weight,
                                gain=nn.init.calculate_gain('linear')) 

        self.device = device

        if "cuda" in self.device:
            self.use_cuda = True
        else:
            self.use_cuda = False
        
        print("device: ", self.device)
        print("use_cuda: ", self.use_cuda)
    
    def forward(self, state):

        # print("state before transform:",  state.device)

        state = self.input_processor(state)

        # print("state after transform:",  state.device)

        state_ts = state[0].to(self.device)
        state_static = state[1].to(self.device)

        self._h1_recurrent.flatten_parameters()

        h0 = torch.zeros(1 , state_ts.size(0), int(self._n_features/2)).requires_grad_().to(self.device)
        
        features1_rnn, _ = self._h1_recurrent(state_ts.float(), h0.detach())
        features1_rnn = features1_rnn[:, -1, :]
        #features1_rnn = torch.zeros(state_ts.size(0), int(self._n_features/2))
        features1_mpl = F.relu(self._h1_mlp(state_static.float()))
        features1 = torch.cat((features1_rnn, features1_mpl), dim = 1)

        features2 = F.relu(self._h2(features1)) 

        a = self._h3(features2)

        return a
