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

# %% auto 0
__all__ = ['BaseNetwork']

# %% ../../../nbs/agents/networks/00_base.ipynb 4
# 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/00_base.ipynb 6
class BaseNetwork(nn.Module, ABC):
    """
    Base class for simple 3-Layer network.
    """

    # TODO make the network architecture more flexible

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

        n_input = input_shape[-1]
        n_output = output_shape[0]

        self._h1 = nn.Linear(n_input, n_features)
        self._h2 = nn.Linear(n_features, n_features)
        self._h3 = nn.Linear(n_features, n_output)

        nn.init.xavier_uniform_(self._h1.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'))
    

    @abstractmethod
    def forward(self, *args, **kwargs):
        pass
