from typing import Callable, List, Optional

import torch
import torch.nn as nn


class FeatureEmbedder(nn.Module):
    def __init__(self, cardinalities: List[int], embedding_dims: List[int],) -> None:
        super().__init__()

        self.__num_features = len(cardinalities)

        def create_embedding(c: int, d: int) -> nn.Embedding:
            embedding = nn.Embedding(c, d)
            return embedding

        self.__embedders = nn.ModuleList(
            [create_embedding(c, d) for c, d in zip(cardinalities, embedding_dims)]
        )

    def forward(self, features: torch.Tensor) -> torch.Tensor:
        if self.__num_features > 1:
            # we slice the last dimension, giving an array of length
            # self.__num_features with shape (N,T) or (N)
            cat_feature_slices = torch.chunk(features, self.__num_features, dim=-1)
        else:
            cat_feature_slices = [features]

        return torch.cat(
            [
                embed(cat_feature_slice.squeeze(-1))
                for embed, cat_feature_slice in zip(
                    self.__embedders, cat_feature_slices
                )
            ],
            dim=-1,
        )


class FeatureAssembler(nn.Module):
    def __init__(
        self,
        T: int,
        embed_static: Optional[FeatureEmbedder] = None,
        embed_dynamic: Optional[FeatureEmbedder] = None,
    ) -> None:
        super().__init__()

        self.T = T
        self.embeddings = nn.ModuleDict(
            {"embed_static": embed_static, "embed_dynamic": embed_dynamic}
        )

    def forward(
        self,
        feat_static_cat: torch.Tensor,
        feat_static_real: torch.Tensor,
        feat_dynamic_cat: torch.Tensor,
        feat_dynamic_real: torch.Tensor,
    ) -> torch.Tensor:
        processed_features = [
            self.process_static_cat(feat_static_cat),
            self.process_static_real(feat_static_real),
            self.process_dynamic_cat(feat_dynamic_cat),
            self.process_dynamic_real(feat_dynamic_real),
        ]

        return torch.cat(processed_features, dim=-1)

    def process_static_cat(self, feature: torch.Tensor) -> torch.Tensor:
        if self.embeddings["embed_static"] is not None:
            feature = self.embeddings["embed_static"](feature)
        return feature.unsqueeze(1).expand(-1, self.T, -1).float()

    def process_dynamic_cat(self, feature: torch.Tensor) -> torch.Tensor:
        if self.embeddings["embed_dynamic"] is None:
            return feature.float()
        else:
            return self.embeddings["embed_dynamic"](feature)

    def process_static_real(self, feature: torch.Tensor) -> torch.Tensor:
        return feature.unsqueeze(1).expand(-1, self.T, -1)

    def process_dynamic_real(self, feature: torch.Tensor) -> torch.Tensor:
        return feature