# -*- coding: utf-8 -*-
r"""
Feed Forward 
==============
    Simple Feed Forward Neural Network module that can be used for classification or regression
"""
import torch
from torch import nn

from polos.modules.activations import build_activation


class FeedForward(nn.Module):
    """
    Feed Forward Neural Network.

    :param in_dim: Number input features.
    :param out_dim: Number of output features. Default is just a score.
    :param hidden_sizes: list with the size of the hidden layers.
        This parameter can also be a string with the sizes splited by a comma.
    :param activations: Name of the activation function to be used in the hidden layers.
    :param final_activation: Name of the final activation function or False if we dont
        want a final activation.
    :param dropout: dropout to be used in the hidden layers.
    """

    def __init__(
        self,
        in_dim: int,
        out_dim: int = 1,
        hidden_sizes: str = "3072,1536,768",
        activations: str = "Sigmoid",
        final_activation: str = "Sigmoid",
        dropout: float = 0.0,
    ) -> None:
        super().__init__()

        if isinstance(hidden_sizes, str):
            hidden_sizes = [int(x) for x in hidden_sizes.split(",")]

        if isinstance(hidden_sizes, int):
            hidden_sizes = [hidden_sizes]

        activation_func = build_activation(activations)

        self.ff = torch.nn.Sequential()
        self.ff.add_module("linear_1", nn.Linear(in_dim, hidden_sizes[0]))
        self.ff.add_module("activation_1", activation_func)
        self.ff.add_module("dropout_1", nn.Dropout(dropout))

        for i in range(1, len(hidden_sizes)):
            self.ff.add_module(
                "linear_{}".format(i + 1),
                nn.Linear(int(hidden_sizes[i - 1]), hidden_sizes[i]),
            )
            self.ff.add_module("activation_{}".format(i + 1), activation_func)
            self.ff.add_module("dropout_{}".format(i + 1), nn.Dropout(dropout))

        self.ff.add_module(
            "linear_{}".format(len(hidden_sizes) + 1),
            nn.Linear(hidden_sizes[-1], int(out_dim)),
        )
        if final_activation:
            final_activation = build_activation(final_activation)
            self.ff.add_module(
                "activation_{}".format(len(hidden_sizes) + 1), final_activation
            )

    def forward(self, in_features: torch.Tensor) -> torch.Tensor:
        return self.ff(in_features)