from typing import Optional

import torch
from torch.nn import Module

from ..persformer_config import PersformerConfig

# Type aliases
from gdeep.utility.custom_types import Tensor


class MeanPoolingLayer(Module):

    config: PersformerConfig

    def __init__(self, config: PersformerConfig):
        super().__init__()
        self.config = config

    def forward(
        self, input_batch: Tensor, attention_mask: Optional[Tensor] = None
    ) -> Tensor:
        """
        Forward pass of the model.

        Args:
            input_batch:
                The input batch. Of shape (batch_size, sequence_length, hidden_size)

        Returns:
            The pooled output. Of shape (batch_size, hidden_size)
        """
        # Initialize the output tensor
        if attention_mask is not None:
            output = input_batch * attention_mask.unsqueeze(dim=-1)
            output = output.sum(dim=-2) / attention_mask.sum(dim=-1, keepdim=True)
            return output
        else:
            output = input_batch
            return output.mean(dim=-2)
