import torch
from torch import nn
from torch.nn.utils.rnn import pad_sequence
import PytorchBoot.stereotype as stereotype


@stereotype.module("transformer_pose_n_pts_seq_encoder")
class TransformerSequenceEncoder(nn.Module):
    def __init__(self, config):
        super(TransformerSequenceEncoder, self).__init__()
        self.config = config
        embed_dim = config["pts_embed_dim"] + config["pose_embed_dim"]
        encoder_layer = nn.TransformerEncoderLayer(
            d_model=embed_dim,
            nhead=config["num_heads"],
            dim_feedforward=config["ffn_dim"],
            batch_first=True,
        )
        self.transformer_encoder = nn.TransformerEncoder(
            encoder_layer, num_layers=config["num_layers"]
        )
        self.fc = nn.Linear(embed_dim, config["output_dim"])

    def encode_sequence(self, pts_embedding_list_batch, pose_embedding_list_batch):
        combined_features_batch = []
        lengths = []

        for pts_embedding_list, pose_embedding_list in zip(pts_embedding_list_batch, pose_embedding_list_batch):
            combined_features = [
                torch.cat((pts_embed, pose_embed), dim=-1)
                for pts_embed, pose_embed in zip(pts_embedding_list, pose_embedding_list)
            ]
            combined_features_batch.append(torch.stack(combined_features))
            lengths.append(len(combined_features))

        combined_tensor = pad_sequence(combined_features_batch, batch_first=True)  # Shape: [batch_size, max_seq_len, embed_dim]

        max_len = max(lengths)
        padding_mask = torch.tensor([([0] * length + [1] * (max_len - length)) for length in lengths], dtype=torch.bool).to(combined_tensor.device)

        transformer_output = self.transformer_encoder(combined_tensor, src_key_padding_mask=padding_mask)
        final_feature = transformer_output.mean(dim=1)
        final_output = self.fc(final_feature)

        return final_output


if __name__ == "__main__":
    config = {
        "pts_embed_dim": 1024, 
        "pose_embed_dim": 256,
        "num_heads": 4,  
        "ffn_dim": 256,
        "num_layers": 3, 
        "output_dim": 2048, 
    }

    encoder = TransformerSequenceEncoder(config)
    seq_len = [5, 8, 9, 4]
    batch_size = 4

    pts_embedding_list_batch = [
        torch.randn(seq_len[idx], config["pts_embed_dim"]) for idx in range(batch_size)
    ]
    pose_embedding_list_batch = [
        torch.randn(seq_len[idx], config["pose_embed_dim"]) for idx in range(batch_size)
    ]
    output_feature = encoder.encode_sequence(
        pts_embedding_list_batch, pose_embedding_list_batch
    )
    print("Encoded Feature:", output_feature)
    print("Feature Shape:", output_feature.shape)