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

# %% auto 0
__all__ = ['CriticNetworkStateAction', 'CriticNetworkState', 'CriticNetworkStateActionRNN',
           'CriticNetworkStateActionRNN_MLP_hybrid']

# %% ../../../nbs/agents/networks/01_critics.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

# Other libraries
from abc import ABC, abstractmethod

# %% ../../../nbs/agents/networks/01_critics.ipynb 6
class CriticNetworkStateAction(BaseNetwork):
    """
    Simple 3-Layer network for the critic function approximator where both state and action are inputs to the network.
    The forward function takes the state and action as input, concatenates them and passes them through the network.
    
    """
    # TODO make the network architecture more flexible

    def __init__(self, input_shape, output_shape, n_features, squeeze_output = False, **kwargs):
        super().__init__(input_shape, output_shape, n_features, **kwargs)
        self.squeeze = squeeze_output

    def forward(self, state, action):
        state_action = torch.cat((state.float(), action.float()), dim=1)
        features1 = F.relu(self._h1(state_action))
        features2 = F.relu(self._h2(features1))
        q = self._h3(features2)

        if self.squeeze:
            return torch.squeeze(q)
        else:
            return q

# %% ../../../nbs/agents/networks/01_critics.ipynb 7
class CriticNetworkState(BaseNetwork):
    
    """
    Simple 3-Layer network for the critic function approximator for discrete actions where only the state is the input to the network.
    Like in Q-Learning, the critic network outputs one Q-value for each discrete action.
    
    """

    # TODO make the network architecture more flexible
    # TODO put squeeze into base network

    def __init__(self, input_shape, output_shape, n_features, squeeze_output = False, **kwargs):
        super().__init__(input_shape, output_shape, n_features, **kwargs)
        self.squeeze = squeeze_output

    def forward(self, state, action = None):

        state = state.float()

        features1 = F.relu(self._h1(state))
        features2 = F.relu(self._h2(features1))
        q = self._h3(features2)

        if action is None:
            if self.squeeze:
                return torch.squeeze(q)
            else:
                return q
        else:
            action = action.long()
            if len(action.shape) < len(q.shape):
                action=action.unsqueeze(1)
            q_acted = torch.squeeze(q.gather(1, action))

            return q_acted

# %% ../../../nbs/agents/networks/01_critics.ipynb 8
class CriticNetworkStateActionRNN(nn.Module):
    """
    #TODO update docstring
    Base class for simple 3-Layer network.
    """

    def __init__(self, input_shape, output_shape, input_shape_RNN, input_shape_MLP, n_features, input_processor = None, squeeze_output = False, use_cuda=False, **kwargs):
     
        super(CriticNetworkStateActionRNN, self).__init__()

        # print("setting critic network")
        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

        # print("input shape RNN: ", self._n_input_RNN)
        # print("input shape MLP: ", self._n_input_MLP)

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

        # print("input mlp: ", self._n_input_MLP)
        # print("input mlp2: ", int(self._n_features/2))

        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.squeeze = squeeze_output #! check what option is correct here
        
        if use_cuda:
            self.device = torch.device("cuda")
        else:
            self.device = torch.device("cpu")   
        print("device: ", self.device)

        # print("done setting critic network")

    def forward(self, state, action):

        state = self.input_processor(state)

        state_ts = state[0]
        state_static = state[1]
        state_static = torch.cat((state_static, action), dim = 1)

        self._h1_recurrent.flatten_parameters()

        # init hidden state
        h0 = torch.zeros(1 , state_ts.size(0), int(self._n_features/2)).requires_grad_().to(self.device)
    
        # get first layers and concatenate
        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))
        features1_mpl = F.relu(self._h1_mlp(state_static.float().to(self.device)))
        features1 = torch.cat((features1_rnn, features1_mpl), dim = 1)

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

        q = self._h3(features2)

        if self.squeeze:
            return torch.squeeze(q)
        else:
            return q

# %% ../../../nbs/agents/networks/01_critics.ipynb 9
class CriticNetworkStateActionRNN_MLP_hybrid(nn.Module):
    """
    #TODO update docstring
    Base class for simple 3-Layer network.
    """

    def __init__(self, input_shape, output_shape, input_shape_RNN, input_shape_MLP, n_features, input_processor = None, squeeze_output = False, use_cuda=False, **kwargs):
     
        super(CriticNetworkStateActionRNN, 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

        # print("input mlp: ", self._n_input_MLP)
        # print("input mlp2: ", int(self._n_features/2))

        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.squeeze = squeeze_output #! check what option is correct here

        if use_cuda:
            self.device = torch.device("cuda")
        else:
            self.device = torch.device("cpu")   
        print("device: ", self.device)

    def forward(self, state, action):

        # print("in critic")

        state = self.input_processor(state)

        state_ts = state[0]
        state_static = state[1]
        state_static = torch.cat((state_static, action), dim = 1)

        self._h1_recurrent.flatten_parameters()

        # init hidden state
        h0 = torch.zeros(1 , state_ts.size(0), int(self._n_features/2)).requires_grad_().to(self.device)
    
        # get first layers and concatenate
        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))
        features1_mpl = F.relu(self._h1_mlp(state_static.float().to(self.device)))
        features1 = torch.cat((features1_rnn, features1_mpl), dim = 1)

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

        q = self._h3(features2)

        if self.squeeze:
            return torch.squeeze(q)
        else:
            return q
