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

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

# %% ../../../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
