# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
"""
DETR Transformer class.

Copy-paste from torch.nn.Transformer with modifications:
    * positional encodings are passed in MHattention
    * extra LN at the end of encoder is removed
    * decoder returns a stack of activations from all decoding layers
"""
import copy
from typing import Optional, List

import torch
import torch.nn.functional as F
from torch import nn, Tensor
import math
from matplotlib import pyplot as plt
import seaborn as sns

import IPython
e = IPython.embed

def plot_attention_weights(attention_weights, input_tokens, layer_name):
    sns.heatmap(attention_weights, xticklabels=input_tokens, yticklabels=input_tokens, cmap='viridis')
    plt.title(f'Attention Weights at {layer_name}')
    plt.xlabel('Input Tokens')
    plt.ylabel('Attention to Tokens')
    plt.show()

class HATransformer(nn.Module):

    def __init__(self, d_model=512, nhead=8, num_blocks=3, num_encoder_layers=6,
                 num_decoder_layers=6, dim_feedforward=2048, dropout=0.1,
                 activation="relu", normalize_before=False,
                 return_intermediate_dec=False):
        super().__init__()

        # cvae_encoder_layer = TransformerEncoderLayer(d_model, nhead, dim_feedforward,
        #                                         dropout, activation, normalize_before)
        # cvae_encoder_norm = nn.LayerNorm(d_model) if normalize_before else None
        # self.cvae_encoder = TransformerEncoder(cvae_encoder_layer, num_encoder_layers, cvae_encoder_norm)

        self.encoder = HATransformerEncoder(d_model, nhead, num_blocks, 
                                            dim_feedforward, dropout, activation)

        decoder_layer = TransformerDecoderLayer(d_model, nhead, dim_feedforward,
                                                dropout, activation, normalize_before)
        decoder_norm = nn.LayerNorm(d_model)
        # self.decoder1 = TransformerDecoder(decoder_layer, num_decoder_layers, decoder_norm,
        #                                   return_intermediate=return_intermediate_dec)
        # self.decoder2 = TransformerDecoder(decoder_layer, num_decoder_layers, decoder_norm,
        #                                   return_intermediate=return_intermediate_dec)

        self.multiarm_decoder = MultiArmDecoder(decoder_layer, num_decoder_layers, decoder_norm,
                                          return_intermediate=return_intermediate_dec)
        self._reset_parameters()

        self.d_model = d_model
        self.nhead = nhead

    def _reset_parameters(self):
        for p in self.parameters():
            if p.dim() > 1:
                nn.init.xavier_uniform_(p)

    def forward(self, src, mask, query_embed, pos_embed, cls_tokens, proprio_input=None, num_cameras=None, additional_pos_embed=None):
        # TODO flatten only when input has H and W
        if len(src.shape) == 4: # has H and W
            # flatten NxCxHxW to HWxNxC
            bs, c, h, w = src.shape
            src = src.flatten(2).permute(2, 0, 1)
            pos_embed = pos_embed.flatten(2).permute(2, 0, 1).repeat(1, bs, 1)
            query_embed = query_embed.unsqueeze(1).repeat(1, bs, 1)
            # mask = mask.flatten(1)

            additional_pos_embed = additional_pos_embed.unsqueeze(1).repeat(1, bs, 1) # seq, bs, dim
            pos_embed = torch.cat([additional_pos_embed, pos_embed], axis=0)

            proprio_input = proprio_input.permute(1,0,2)
            # addition_input = torch.stack([latent_input, proprio_input], axis=0)
            src = torch.cat([proprio_input, src], axis=0)
        else:
            assert len(src.shape) == 3
            # flatten NxHWxC to HWxNxC
            bs, hw, c = src.shape
            src = src.permute(1, 0, 2)
            pos_embed = pos_embed.unsqueeze(1).repeat(1, bs, 1)
            query_embed = query_embed.unsqueeze(1).repeat(1, bs, 1)

        tgt1 = torch.zeros_like(query_embed)
        tgt2 = torch.zeros_like(query_embed)
        # cls_tokens = torch.zeros((3, bs, self.d_model), device=src.device)

        memory = self.encoder(src, cls_tokens, num_cameras, device=src.device)
        chunk1_index = 7 + 1
        chunk2_index = 2*(7 + 1)
        arm1_chunk = memory[0:9]  # CLS + 7 for arm1 + latent_input_arm1
        arm2_chunk = memory[9:18]  # CLS + 7 for arm2 + latent_input_arm2
        obs_chunk = memory[18:]  # CLS + observation
    
        # Memory1 for decoder1: CLS from arm2, chunk from arm1, chunk from obs
        memory1 = torch.cat([arm2_chunk[:3], arm1_chunk, obs_chunk], dim=0)  # CLS_arm2, all tokens of arm1, all tokens of obs

        # Memory2 for decoder2: CLS from arm1, chunk from arm2, chunk from obs
        memory2 = torch.cat([arm1_chunk[:3], arm2_chunk, obs_chunk], dim=0)  # CLS_arm1, all tokens of arm2, all tokens of obs

        hs1, hs2, attn_weights, attn_weights2 = self.multiarm_decoder(tgt1, tgt2, memory, memory_key_padding_mask=mask,
                          pos=pos_embed, query_pos=query_embed)
        # hs1 = hs1.transpose(1, 2)

        # hs2 = self.decoder2(tgt2, memory, memory_key_padding_mask=mask,
        #                   pos=pos_embed, query_pos=query_embed)
        # hs2 = hs2.transpose(1, 2)

        hs = torch.cat([hs1, hs2], axis=0)

        return hs1, hs2, attn_weights, attn_weights2
    
    

class TransformerEncoder(nn.Module):

    def __init__(self, encoder_layer, num_layers, norm=None):
        super().__init__()
        self.layers = _get_clones(encoder_layer, num_layers)
        self.num_layers = num_layers
        self.norm = norm

    def forward(self, src,
                mask: Optional[Tensor] = None,
                src_key_padding_mask: Optional[Tensor] = None,
                pos: Optional[Tensor] = None):
        output = src

        for layer in self.layers:
            output = layer(output, src_mask=mask,
                           src_key_padding_mask=src_key_padding_mask, pos=pos)

        if self.norm is not None:
            output = self.norm(output)

        return output

class MultiArmDecoder(nn.Module):

    def __init__(self, decoder_layer, num_layers, norm=None, return_intermediate=False, dim_feedforward=2048, dropout=0.1,
                 activation="relu"):
        super().__init__()
        # self.layers = _get_clones(decoder_layer, num_layers)
        self.num_layers = num_layers
        self.norm = norm
        self.return_intermediate = return_intermediate
        self.arm1_decoder = _get_clones(decoder_layer, num_layers)
        self.arm2_decoder = _get_clones(decoder_layer, num_layers)
        self.self_attn = nn.MultiheadAttention(512, num_heads=8)

        self.linear1 = nn.Linear(512, dim_feedforward)
        self.dropout = nn.Dropout(dropout)
        self.linear2 = nn.Linear(dim_feedforward, 512)

        self.norm1 = nn.LayerNorm(512)
        self.norm2 = nn.LayerNorm(512)
        self.dropout1 = nn.Dropout(dropout)
        self.dropout2 = nn.Dropout(dropout)

        self.activation = _get_activation_fn(activation)

        device = next(self.parameters()).device  # Assuming 'self' is your model instance
        self.pos_encoder = PositionalEncoding(512, 5000, device=device)


    def forward(self, tgt1, tgt2, memory,
                tgt_mask: Optional[Tensor] = None,
                memory_mask: Optional[Tensor] = None,
                tgt_key_padding_mask: Optional[Tensor] = None,
                memory_key_padding_mask: Optional[Tensor] = None,
                pos: Optional[Tensor] = None,
                query_pos: Optional[Tensor] = None):
        # output = tgt

        # intermediate = []
        arm1_cls = memory[:7]
        arm1_chunk = memory[7:14]
        arm2_cls = memory[14:21]
        arm2_chunk = memory[21:28]
        obs_chunk = memory[28:]



        memory1 = torch.cat([arm1_cls, arm1_chunk, arm2_cls, obs_chunk], dim=0)  # CLS_arm2, all tokens of arm1, all tokens of obs
        memory2 = torch.cat([arm1_cls, arm2_cls, arm2_chunk, obs_chunk], dim=0)

        for layer in range(self.num_layers):
            tgt1, attn_weights = self.arm1_decoder[layer](tgt1, memory1, tgt_mask=tgt_mask,
                           memory_mask=memory_mask,
                           tgt_key_padding_mask=tgt_key_padding_mask,
                           memory_key_padding_mask=memory_key_padding_mask,
                           pos=pos, query_pos=query_pos)
            tgt2, attn_weights2 = self.arm2_decoder[layer](tgt2, memory2, tgt_mask=tgt_mask,
                            memory_mask=memory_mask,
                            tgt_key_padding_mask=tgt_key_padding_mask,
                            memory_key_padding_mask=memory_key_padding_mask,
                            pos=pos, query_pos=query_pos)
            # output = layer(output, memory, tgt_mask=tgt_mask,
            #                memory_mask=memory_mask,
            #                tgt_key_padding_mask=tgt_key_padding_mask,
            #                memory_key_padding_mask=memory_key_padding_mask,
            #                pos=pos, query_pos=query_pos)
            
            if layer <= self.num_layers - 2:
                combined = torch.cat((tgt1, tgt2), dim=0)
                k = q = self.pos_encoder(combined)
                combined2, _ = self.self_attn(k, q, value=combined)

                combined = combined + self.dropout1(combined2)
                combined = self.norm1(combined)
                combined2 = self.linear2(self.dropout(self.activation(self.linear1(combined))))
                combined = combined + self.dropout2(combined2)
                combined = self.norm2(combined)


                tgt1, tgt2 = torch.split(combined, tgt1.size(0), dim=0)
                tgt1 = self.norm(tgt1)
                tgt2 = self.norm(tgt2)
            else:
                tgt1 = self.norm(tgt1)
                tgt2 = self.norm(tgt2)
            # if self.return_intermediate:
            #     intermediate.append(self.norm(output))

        # if self.norm is not None:
        #     output = self.norm(output)
        #     if self.return_intermediate:
        #         intermediate.pop()
        #         intermediate.append(output)

        # if self.return_intermediate:
        #     return torch.stack(intermediate)

        return tgt1, tgt2, attn_weights, attn_weights2

class PositionalEncoding(nn.Module):
    def __init__(self, d_model, max_len=3, device=None):
        super(PositionalEncoding, self).__init__()
        self.encoding = torch.zeros(max_len, d_model, device=device)
        position = torch.arange(0, max_len, dtype=torch.float, device=device).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))

        self.encoding[:, 0::2] = torch.sin(position * div_term)
        self.encoding[:, 1::2] = torch.cos(position * div_term)
        self.encoding = self.encoding.unsqueeze(0)

    def forward(self, x):
        # Ensure positional encoding is on the same device as x
        if self.encoding.device != x.device:
            self.encoding = self.encoding.to(x.device)
        return x + self.encoding[:, :x.size(1)].detach()
    
class TransformerEncoderLayer(nn.Module):

    def __init__(self, d_model, nhead, dim_feedforward=2048, dropout=0.1,
                 activation="relu", normalize_before=False):
        super().__init__()
        self.self_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout)
        # Implementation of Feedforward model
        self.linear1 = nn.Linear(d_model, dim_feedforward)
        self.dropout = nn.Dropout(dropout)
        self.linear2 = nn.Linear(dim_feedforward, d_model)

        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.dropout1 = nn.Dropout(dropout)
        self.dropout2 = nn.Dropout(dropout)

        self.activation = _get_activation_fn(activation)
        self.normalize_before = normalize_before

    def with_pos_embed(self, tensor, pos: Optional[Tensor]):
        return tensor if pos is None else tensor + pos

    def forward_post(self,
                     src,
                     src_mask: Optional[Tensor] = None,
                     src_key_padding_mask: Optional[Tensor] = None,
                     pos: Optional[Tensor] = None):
        q = k = self.with_pos_embed(src, pos)
        src2 = self.self_attn(q, k, value=src, attn_mask=src_mask,
                              key_padding_mask=src_key_padding_mask)[0]
        src = src + self.dropout1(src2)
        src = self.norm1(src)
        src2 = self.linear2(self.dropout(self.activation(self.linear1(src))))
        src = src + self.dropout2(src2)
        src = self.norm2(src)
        return src

    def forward_pre(self, src,
                    src_mask: Optional[Tensor] = None,
                    src_key_padding_mask: Optional[Tensor] = None,
                    pos: Optional[Tensor] = None):
        src2 = self.norm1(src)
        q = k = self.with_pos_embed(src2, pos)
        src2 = self.self_attn(q, k, value=src2, attn_mask=src_mask,
                              key_padding_mask=src_key_padding_mask)[0]
        src = src + self.dropout1(src2)
        src2 = self.norm2(src)
        src2 = self.linear2(self.dropout(self.activation(self.linear1(src2))))
        src = src + self.dropout2(src2)
        return src

    def forward(self, src,
                src_mask: Optional[Tensor] = None,
                src_key_padding_mask: Optional[Tensor] = None,
                pos: Optional[Tensor] = None):
        if self.normalize_before:
            return self.forward_pre(src, src_mask, src_key_padding_mask, pos)
        return self.forward_post(src, src_mask, src_key_padding_mask, pos)



#### SEGMENT-WISE ENCODER
class SegmentWiseEncoder(nn.Module):
    def __init__(self, d_model, nhead, dim_feedforward=2048, dropout=0.1, activation="relu"):
        super().__init__()
        self.self_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout)
        self.norm1 = nn.LayerNorm(d_model)
        self.dropout1 = nn.Dropout(dropout)

        self.linear1 = nn.Linear(d_model, dim_feedforward)
        self.activation = getattr(nn.functional, activation)
        self.dropout = nn.Dropout(dropout)
        self.linear2 = nn.Linear(dim_feedforward, d_model)
        self.dropout2 = nn.Dropout(dropout)
        self.norm2 = nn.LayerNorm(d_model)

    def forward(self, src, src_key_padding_mask=None, pos_encoder=None):
        # k = q = pos_encoder(src)
        src2 = self.self_attn(src, src, src, key_padding_mask=src_key_padding_mask)[0]
        src = src + self.dropout1(src2)
        src = self.norm1(src)

        src2 = self.linear2(self.dropout(self.activation(self.linear1(src))))
        src = src + self.dropout2(src2)
        src = self.norm2(src)
        return src

#### CROSS-SEGMENT ENCODER
class CrossSegmentEncoder(nn.Module):
    def __init__(self, d_model, nhead, dim_feedforward=2048, dropout=0.1, activation="relu"):
        super().__init__()
        self.self_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout)
        self.norm1 = nn.LayerNorm(d_model)
        self.dropout1 = nn.Dropout(dropout)

        self.linear1 = nn.Linear(d_model, dim_feedforward)
        self.activation = getattr(nn.functional, activation)
        self.dropout = nn.Dropout(dropout)
        self.linear2 = nn.Linear(dim_feedforward, d_model)
        self.dropout2 = nn.Dropout(dropout)
        self.norm2 = nn.LayerNorm(d_model)

    def forward(self, cls_tokens, pos_encoder=None):
        # Processing only the CLS tokens
        # q = k = pos_encoder(cls_tokens)
        cls_tokens2 = self.self_attn(cls_tokens, cls_tokens, cls_tokens)[0]
        cls_tokens = cls_tokens + self.dropout1(cls_tokens2)
        cls_tokens = self.norm1(cls_tokens)

        cls_tokens2 = self.linear2(self.dropout(self.activation(self.linear1(cls_tokens))))
        cls_tokens = cls_tokens + self.dropout2(cls_tokens2)
        cls_tokens = self.norm2(cls_tokens)
        return cls_tokens
    
class HATransformerEncoder(nn.Module): ## TODO: NEED TO CONVER TO TENSOR OPERATION
    def __init__(self, d_model, nhead, num_blocks=3, dim_feedforward=2048, dropout=0.1, activation="relu"):
        super().__init__()
        self.num_blocks = num_blocks
        self.segment_encoders = nn.ModuleList([SegmentWiseEncoder(d_model, nhead, dim_feedforward, dropout, activation) for _ in range(num_blocks)])
        self.cross_segment_encoders = nn.ModuleList([CrossSegmentEncoder(d_model, nhead, dim_feedforward, dropout, activation) for _ in range(num_blocks)])
        device = next(self.parameters()).device  # Assuming 'self' is your model instance
        self.pos_encoder = PositionalEncoding(d_model, 5000, device=device)
        self.arm_cls = 7
        self.cam_length = 300
        self.cam_cls = 5
        self.cam_with_cls = 300 + self.cam_cls
        # self.layer_norm = nn.LayerNorm(d_model)

    def forward(self, segments, cls_tokens, num_cameras, device):
        segment_1 = torch.cat([cls_tokens[0:self.arm_cls,:,:], segments[0:7,:,:]], dim=0)
        segment_2 = torch.cat([cls_tokens[self.arm_cls:2*self.arm_cls,:,:], segments[7:14,:,:]], dim=0)
        segment_3_1 = torch.cat([cls_tokens[2*self.arm_cls:2*self.arm_cls+self.cam_cls,:,:], segments[14:14+self.cam_length,:,:]], dim=0)
        if num_cameras == 4:
            segment_3_2 = torch.cat([cls_tokens[2*self.arm_cls+self.cam_cls:2*self.arm_cls+2*self.cam_cls,:,:], segments[14+self.cam_length:14+2*self.cam_length,:,:]], dim=0)
            segment_3_3 = torch.cat([cls_tokens[2*self.arm_cls+2*self.cam_cls:2*self.arm_cls+3*self.cam_cls,:,:], segments[14+2*self.cam_length:14+3*self.cam_length,:,:]], dim=0)
            segment_3_4 = torch.cat([cls_tokens[2*self.arm_cls+3*self.cam_cls:2*self.arm_cls+4*self.cam_cls,:,:], segments[14+3*self.cam_length:14+4*self.cam_length,:,:]], dim=0)
            segments = torch.cat([segment_1, segment_2, segment_3_1, segment_3_2, segment_3_3, segment_3_4], dim=0)
        else:
            segments = torch.cat([segment_1, segment_2, segment_3_1], dim=0)
        # segments = [segment_1, segment_2, segment_3]
        segments = self.pos_encoder(segments.to(device))

        for i in range(self.num_blocks):
            segment_1 = segments[0:self.arm_cls+7]
            segment_2 = segments[self.arm_cls+7:2*(self.arm_cls+7)]
            segment_3_1 = segments[2*(self.arm_cls+7):2*(self.arm_cls+7)+self.cam_with_cls]

            updated_segment_1 = self.segment_encoders[i](segment_1)
            updated_segment_2 = self.segment_encoders[i](segment_2)
            updated_segment_3_1 = self.segment_encoders[i](segment_3_1)
            if num_cameras == 4:
                segment_3_2 = segments[2*(self.arm_cls+7)+self.cam_with_cls:2*(self.arm_cls+7)+2*self.cam_with_cls]
                segment_3_3 = segments[2*(self.arm_cls+7)+2*self.cam_with_cls:2*(self.arm_cls+7)+3*self.cam_with_cls]
                segment_3_4 = segments[2*(self.arm_cls+7)+3*self.cam_with_cls:2*(self.arm_cls+7)+4*self.cam_with_cls]
                updated_segment_3_2 = self.segment_encoders[i](segment_3_2)
                updated_segment_3_3 = self.segment_encoders[i](segment_3_3)
                updated_segment_3_4 = self.segment_encoders[i](segment_3_4)
                updated_cls_tokens = self.cross_segment_encoders[i](torch.cat([updated_segment_1[0:self.arm_cls], updated_segment_2[0:self.arm_cls], updated_segment_3_1[0:self.cam_cls], updated_segment_3_2[0:self.cam_cls], updated_segment_3_3[0:self.cam_cls], updated_segment_3_4[0:self.cam_cls]], axis=0))
            else:
                updated_cls_tokens = self.cross_segment_encoders[i](torch.cat([updated_segment_1[0:self.arm_cls], updated_segment_2[0:self.arm_cls], updated_segment_3_1[0:self.cam_cls]], axis=0))


            if num_cameras == 4:
                segments = torch.cat([
                    updated_cls_tokens[0:self.arm_cls],                                                segment_1[self.arm_cls:], 
                    updated_cls_tokens[self.arm_cls:2*self.arm_cls],                                   segment_2[self.arm_cls:], 
                    updated_cls_tokens[2*self.arm_cls:2*self.arm_cls+self.cam_cls],                    segment_3_1[self.cam_cls:],
                    updated_cls_tokens[2*self.arm_cls+self.cam_cls:2*self.arm_cls+2*self.cam_cls],     segment_3_2[self.cam_cls:],
                    updated_cls_tokens[2*self.arm_cls+2*self.cam_cls:2*self.arm_cls+3*self.cam_cls],   segment_3_3[self.cam_cls:],
                    updated_cls_tokens[2*self.arm_cls+3*self.cam_cls:2*self.arm_cls+4*self.cam_cls],   segment_3_4[self.cam_cls:]
                ])
            else:
                segments = torch.cat([
                    updated_cls_tokens[0:self.arm_cls], segment_1[self.arm_cls:], 
                    updated_cls_tokens[self.arm_cls:2*self.arm_cls], segment_2[self.arm_cls:], 
                    updated_cls_tokens[2*self.arm_cls:2*self.arm_cls+self.cam_cls], segment_3_1[self.cam_cls:]
                ])


        # if i == self.num_blocks - 1:
        #     plot_attention_weights(attn_weights, [0, 1, 2], 'Encoder Layer X')
        return segments 

# class HATransformerEncoder(nn.Module):
#     def __init__(self, d_model, nhead, num_blocks=3, dim_feedforward=2048, dropout=0.1, activation="relu", num_segment_layers=1, num_cross_layers=1, final_layers=1,):
#         super().__init__()
#         self.num_segment_layers = num_segment_layers
#         self.num_cross_layers = num_cross_layers
#         # Create multiple segment-wise encoders
#         # Use the same segment-wise encoder for the first two segments
#         self.shared_segment_encoder = nn.ModuleList([
#             SegmentWiseEncoder(d_model, nhead, dim_feedforward, dropout, activation) for _ in range(num_segment_layers)
#         ])
        
#         # Use a separate segment-wise encoder for the third segment
#         self.unique_segment_encoder = nn.ModuleList([
#             SegmentWiseEncoder(d_model, nhead, dim_feedforward, dropout, activation) for _ in range(num_segment_layers)
#         ])
        
#         # Create cross-segment encoders
#         self.cross_segment_encoders = nn.ModuleList([
#             CrossSegmentEncoder(d_model, nhead, dim_feedforward, dropout, activation) for _ in range(num_cross_layers)
#         ])

#         self.final_segment_encoders = nn.ModuleList([SegmentWiseEncoder(d_model, nhead, dim_feedforward, dropout, activation) for _ in range(final_layers)])

#         self.pos_encoder = PositionalEncoding(d_model, 5000)  # Device will be set during the forward pass

#     def forward(self, segments, cls_tokens, num_cameras, device):

#         segment_1 = torch.cat([cls_tokens[0:1,:,:], segments[0:7,:,:]], dim=0)
#         segment_2 = torch.cat([cls_tokens[1:2,:,:], segments[7:14,:,:]], dim=0)
#         segment_3_1 = torch.cat([cls_tokens[2:3,:,:], segments[14:14+300,:,:]], dim=0)

#         segments = torch.cat([segment_1, segment_2, segment_3_1], dim=0)

#         # Positional encoding is applied once at the start
#         segments = self.pos_encoder(segments.to(device))
#         segments = [segments[0:8], segments[8:16], segments[16:16+301]]
#         for i in range(self.num_cross_layers):
            
#             # Apply multiple segment-wise encoders before each cross-segment encoding
#             for j in range(self.num_segment_layers):
#                 segments[0] = self.shared_segment_encoder[j](segments[0])  # First segment
#                 segments[1] = self.shared_segment_encoder[j](segments[1])  # Second segment
#                 segments[2] = self.unique_segment_encoder[j](segments[2]) 
#             # Concatenate the cls tokens and updated segments
#             cls_tokens_concat = torch.cat([seg[0:1] for seg in segments], dim=0)
#             cls_tokens = self.cross_segment_encoders[i](cls_tokens_concat)
#             # Apply cross-segment encoder
#             segments = torch.cat([cls_tokens[0:1], segments[0][1:], cls_tokens[1:2], segments[1][1:], cls_tokens[2:3], segments[2][1:]], dim=0)
        
#             segments = [segments[0:8], segments[8:16], segments[16:16+301]]
        
#         for final_encoder in self.final_segment_encoders:
#             segments = [final_encoder(seg) for seg in segments]

#         segments = torch.cat(segments, dim=0)
#         return segments

class TransformerDecoderLayer(nn.Module):

    def __init__(self, d_model, nhead, dim_feedforward=2048, dropout=0.1,
                 activation="relu", normalize_before=False):
        super().__init__()
        self.self_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout)
        self.multihead_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout)
        # Implementation of Feedforward model
        self.linear1 = nn.Linear(d_model, dim_feedforward)
        self.dropout = nn.Dropout(dropout)
        self.linear2 = nn.Linear(dim_feedforward, d_model)

        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.norm3 = nn.LayerNorm(d_model)
        self.dropout1 = nn.Dropout(dropout)
        self.dropout2 = nn.Dropout(dropout)
        self.dropout3 = nn.Dropout(dropout)

        self.activation = _get_activation_fn(activation)
        self.normalize_before = normalize_before

    def with_pos_embed(self, tensor, pos: Optional[Tensor]):
        return tensor if pos is None else tensor + pos

    def forward_post(self, tgt, memory,
                     tgt_mask: Optional[Tensor] = None,
                     memory_mask: Optional[Tensor] = None,
                     tgt_key_padding_mask: Optional[Tensor] = None,
                     memory_key_padding_mask: Optional[Tensor] = None,
                     pos: Optional[Tensor] = None,
                     query_pos: Optional[Tensor] = None):
        q = k = self.with_pos_embed(tgt, query_pos)
        tgt2 = self.self_attn(q, k, value=tgt, attn_mask=tgt_mask,
                              key_padding_mask=tgt_key_padding_mask)[0]
        tgt = tgt + self.dropout1(tgt2)
        tgt = self.norm1(tgt)
        tgt2,  attn_weights = self.multihead_attn(query=self.with_pos_embed(tgt, query_pos),
                                   key=self.with_pos_embed(memory, pos),
                                   value=memory, attn_mask=memory_mask,
                                   key_padding_mask=memory_key_padding_mask)
        tgt = tgt + self.dropout2(tgt2)
        tgt = self.norm2(tgt)
        tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt))))
        tgt = tgt + self.dropout3(tgt2)
        tgt = self.norm3(tgt)
        return tgt, attn_weights

    def forward_pre(self, tgt, memory,
                    tgt_mask: Optional[Tensor] = None,
                    memory_mask: Optional[Tensor] = None,
                    tgt_key_padding_mask: Optional[Tensor] = None,
                    memory_key_padding_mask: Optional[Tensor] = None,
                    pos: Optional[Tensor] = None,
                    query_pos: Optional[Tensor] = None):
        tgt2 = self.norm1(tgt)
        q = k = self.with_pos_embed(tgt2, query_pos)
        tgt2 = self.self_attn(q, k, value=tgt2, attn_mask=tgt_mask,
                              key_padding_mask=tgt_key_padding_mask)[0]
        tgt = tgt + self.dropout1(tgt2)
        tgt2 = self.norm2(tgt)
        tgt2, attn_weights = self.multihead_attn(query=self.with_pos_embed(tgt2, query_pos),
                                   key=self.with_pos_embed(memory, pos),
                                   value=memory, attn_mask=memory_mask,
                                   key_padding_mask=memory_key_padding_mask)
        tgt = tgt + self.dropout2(tgt2)
        tgt2 = self.norm3(tgt)
        tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt2))))
        tgt = tgt + self.dropout3(tgt2)
        return tgt, attn_weights

    def forward(self, tgt, memory,
                tgt_mask: Optional[Tensor] = None,
                memory_mask: Optional[Tensor] = None,
                tgt_key_padding_mask: Optional[Tensor] = None,
                memory_key_padding_mask: Optional[Tensor] = None,
                pos: Optional[Tensor] = None,
                query_pos: Optional[Tensor] = None):
        if self.normalize_before:
            return self.forward_pre(tgt, memory, tgt_mask, memory_mask,
                                    tgt_key_padding_mask, memory_key_padding_mask, pos, query_pos)
        return self.forward_post(tgt, memory, tgt_mask, memory_mask,
                                 tgt_key_padding_mask, memory_key_padding_mask, pos, query_pos)


def _get_clones(module, N):
    return nn.ModuleList([copy.deepcopy(module) for i in range(N)])


def build_hatransformer(args):
    return HATransformer(
        d_model=args.hidden_dim,
        dropout=args.dropout,
        nhead=args.nheads,
        num_blocks=args.num_blocks,
        dim_feedforward=args.dim_feedforward,
        num_encoder_layers=args.enc_layers,
        num_decoder_layers=args.dec_layers,
        normalize_before=args.pre_norm,
        return_intermediate_dec=True,
    )


def _get_activation_fn(activation):
    """Return an activation function given a string"""
    if activation == "relu":
        return F.relu
    if activation == "gelu":
        return F.gelu
    if activation == "glu":
        return F.glu
    raise RuntimeError(F"activation should be relu/gelu, not {activation}.")
