from typing import Any, Dict, List

import torch
import torch.nn as nn
from torch import Tensor


class ConcatLinear(nn.Module):
    """
    Concat different outputs and feed into a linear layer

    :param concat_output_keys: The keys of features to concat.
    :type concat_output_keys: list of str
    :param input_dim: The total sum of input dim.
    :type input_dim: int
    :param outpt_dim: The output dim.
    :type outpt_dim: int

    """

    def __init__(
        self, concat_output_keys: List[str], input_dim: int, outpt_dim: int
    ) -> None:
        super().__init__()

        self.concat_output_keys = concat_output_keys
        self.linear = nn.Linear(input_dim, outpt_dim)

    def forward(self, intermediate_output_dict: Dict[str, Any]) -> Tensor:
        input_feature = torch.cat(
            [intermediate_output_dict[key][0] for key in self.concat_output_keys], dim=1
        )
        return self.linear(input_feature)
