import torch
from torch import nn

import PytorchBoot.stereotype as stereotype

@stereotype.module("transformer_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']
        self.positional_encoding = nn.Parameter(torch.zeros(1, config['max_seq_len'], embed_dim))
        encoder_layer = nn.TransformerEncoderLayer(d_model=embed_dim, nhead=config['num_heads'], dim_feedforward=config['ffn_dim'])
        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):
        batch_size = len(pts_embedding_list_batch)
        combined_features_batch = []
        
        for i in range(batch_size):
            combined_features = [torch.cat((pts_embed, pose_embed), dim=-1)
                                 for pts_embed, pose_embed in zip(pts_embedding_list_batch[i][:-1], pose_embedding_list_batch[i][:-1])]
            combined_features_batch.append(torch.stack(combined_features))
        
        combined_tensor = torch.stack(combined_features_batch)  # Shape: [batch_size, seq_len-1, embed_dim]
        
        # Adjust positional encoding to match batch size
        pos_encoding = self.positional_encoding[:, :combined_tensor.size(1), :].repeat(batch_size, 1, 1)
        combined_tensor = combined_tensor + pos_encoding

        # Transformer encoding
        transformer_output = self.transformer_encoder(combined_tensor)
        
        # Mean pooling
        final_feature = transformer_output.mean(dim=1)
        
        # Fully connected layer
        final_output = self.fc(final_feature)
        
        return final_output

if __name__ == "__main__":
    config = {
        'pts_embed_dim': 1024,  # 每个点云embedding的维度
        'pose_embed_dim': 256,  # 每个姿态embedding的维度
        'num_heads': 4,   # 多头注意力机制的头数
        'ffn_dim': 256,   # 前馈神经网络的维度
        'num_layers': 3,  # Transformer 编码层数
        'max_seq_len': 10, # 最大序列长度
        'output_dim': 2048, # 输出特征维度
    }

    encoder = TransformerSequenceEncoder(config)
    seq_len = 5
    batch_size = 4
    
    pts_embedding_list_batch = [torch.randn(seq_len, config['pts_embed_dim']) for _ in range(batch_size)]
    pose_embedding_list_batch = [torch.randn(seq_len, config['pose_embed_dim']) for _ 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)