import torch.nn as nn
from torch.nn import Dropout, Linear, Module, Sequential

from gdeep.topology_layers.pooling_layers import (
    AttentionPoolingLayer,
    MaxPoolingLayer,
    MeanPoolingLayer,
    SumPoolingLayer,
)
from gdeep.utility.enum_types import ActivationFunction, PoolerType
from .attention_factory import AttentionFactory
from .persformer_config import PersformerConfig


def get_pooling_layer(config: PersformerConfig) -> Module:
    """
    Get the pooling layer.

    Args:
        config: The configuration of the model.

    Returns:
        The pooling layer.
    """
    if config.pooler_type is PoolerType.ATTENTION:
        return AttentionPoolingLayer(config)  # type: ignore
    elif config.pooler_type is PoolerType.MAX:
        return MaxPoolingLayer(config)  # type: ignore
    elif config.pooler_type is PoolerType.MEAN:
        return MeanPoolingLayer(config)  # type: ignore
    elif config.pooler_type is PoolerType.SUM:
        return SumPoolingLayer(config)  # type: ignore
    else:
        raise ValueError(f"Pooler type {config.pooler_type} is not supported.")


def get_feed_forward_layer(config: PersformerConfig) -> Module:
    """
    Get the feed forward layer.
    """
    feed_forward_layer = Sequential()
    feed_forward_layer.add_module(
        "intermediate", Linear(config.hidden_size, config.intermediate_size)
    )
    feed_forward_layer.add_module(
        "activation", get_activation_function(config.hidden_act)
    )
    feed_forward_layer.add_module("dropout", Dropout(config.hidden_dropout_prob))

    feed_forward_layer.add_module(
        "output", Linear(config.intermediate_size, config.hidden_size)
    )
    feed_forward_layer.add_module("dropout", Dropout(config.hidden_dropout_prob))
    return feed_forward_layer


def get_attention_layer(config: PersformerConfig) -> Module:
    """
    Get the attention layer.
    """
    return AttentionFactory().build(config)  # type: ignore


def get_activation_function(activation_function: ActivationFunction) -> Module:
    """
    Get the activation function.
    """
    if activation_function is ActivationFunction.RELU:
        return nn.ReLU()
    elif activation_function is ActivationFunction.GELU:
        return nn.GELU()
    elif activation_function is ActivationFunction.SELU:
        return nn.SELU()
    elif activation_function is ActivationFunction.MISH:
        return nn.Mish()
    else:
        raise ValueError("Unknown activation function.")
