import torch
import torch.nn as nn
from models.utils import get_1d_sincos_pos_embed
from models.transformer_modules import BlockFlashAttn, freq_pos_enc

class Decoder(nn.Module):
    def __init__(self, num_patches=75, encoder_embed_dim=128, decoder_embed_dim=128, decoder_depth=8, decoder_num_heads=4,
                 mlp_ratio=4, norm_layer=nn.LayerNorm, decoder_pred_dim=40, fusion_method=None,  pos_enc='sinusoidal'):
        super().__init__()

        # --------------------------------------------------------------------------
        self.num_patches = num_patches        
        # ConvMAE decoder specifics
        self.decoder_embed = nn.Linear(encoder_embed_dim, decoder_embed_dim, bias=True)

        self.mask_token = nn.Parameter(torch.zeros(1, 1, decoder_embed_dim))        
        
        self.pos_enc = pos_enc
        self.decoder_pos_embed = nn.Parameter(torch.zeros(1, self.num_patches, decoder_embed_dim), requires_grad=False)  # fixed sin-cos embedding
        if pos_enc == 'rope':
            self.rope = freq_pos_enc(dim=decoder_embed_dim // decoder_num_heads, max_seq_len=num_patches)
        
        self.decoder_blocks = nn.ModuleList([
            BlockFlashAttn(decoder_embed_dim, decoder_num_heads, mlp_ratio, qkv_bias=True, qk_scale=None, norm_layer=norm_layer)
            for i in range(decoder_depth)])

        self.decoder_norm = norm_layer(decoder_embed_dim)        
        
        self.decoder_pred = nn.Linear(decoder_embed_dim, decoder_pred_dim, bias=True) # decoder to patch
        # --------------------------------------------------------------------------        

        self.initialize_weights()

    def initialize_weights(self):
        # initialization
        # initialize (and freeze) pos_embed by sin-cos embedding
        decoder_pos_embed = get_1d_sincos_pos_embed(self.decoder_pos_embed.shape[-1], self.num_patches, cls_token=False)
        self.decoder_pos_embed.data.copy_(torch.from_numpy(decoder_pos_embed).float().unsqueeze(0))

                # timm's trunc_normal_(std=.02) is effectively normal_(std=0.02) as cutoff is too big (2.)
#        torch.nn.init.normal_(self.cls_token, std=.02)
        torch.nn.init.normal_(self.mask_token, std=.02)    

    def forward_features(self, x, ids_keep, idx=0):
        # embed tokens
        x = self.decoder_embed(x)
        N, L, d = x.shape
        
        if self.num_patches > x.shape[1]:
            # append mask tokens to sequence
            masked = self.mask_token.repeat(N, self.num_patches, 1).to(dtype=x.dtype, device=x.device)
            masked.scatter_(dim=1, index=ids_keep.unsqueeze(-1).repeat(1, 1, d), src=x)  # fill the mask tokens with the original tokens
            x_ = masked
            # if idx == 0:
            #     x_ = torch.cat([x, mask_tokens], dim=1)
            # else:
            #     x_ = torch.cat([mask_tokens, x], dim=1)
        else:   
            x_ = x
        # x = torch.gather(x_, dim=1, index=ids_restore.unsqueeze(-1).repeat(1, 1, x.shape[2]))  # unshuffle

        # add pos embed
        if self.pos_enc == 'sinusoidal':            
            x = x_ + self.decoder_pos_embed
        else:
            x = x_

        # apply Transformer blocks
        for blk in self.decoder_blocks:
            if self.pos_enc == 'rope':
                pe = self.rope.unsqueeze(0).to(dtype=x.dtype, device=x.device)  # (1, max_seq_len, dim)
                pe = pe.repeat(N, 1, 1)  # (N, max_seq_len, dim)
            else:
                pe = None
            x, _ = blk(x, rope=pe)
        x = self.decoder_norm(x)

        return x

    def forward(self, x, ids_keep):
        x = self.forward_features(x, ids_keep)

        # predictor projection
        x = self.decoder_pred(x)
        
        return x

class SpecificDecoder(Decoder):
    def __init__(self, num_patches=75, encoder_embed_dim=128, 
                decoder_embed_dim=128, decoder_depth=8, decoder_num_heads=4,
                mlp_ratio=4, norm_layer=nn.LayerNorm, decoder_pred_dim=40,
                fusion_method='concat'):
        super().__init__(num_patches=num_patches, encoder_embed_dim=encoder_embed_dim, 
                decoder_embed_dim=decoder_embed_dim, decoder_depth=decoder_depth, decoder_num_heads=decoder_num_heads,
                mlp_ratio=mlp_ratio, norm_layer=norm_layer, decoder_pred_dim=decoder_pred_dim, fusion_method=fusion_method)  
        
        self.fusion_method = fusion_method
        if self.fusion_method == 'concat':
            self.prenorm = norm_layer(encoder_embed_dim * 2)  # prenorm for the weighted sum features
            self.merge = nn.Linear(encoder_embed_dim * 2, encoder_embed_dim)        
        elif self.fusion_method == 'cross_attn':
            self.cross_attention = BlockFlashAttn(encoder_embed_dim, decoder_num_heads, mlp_ratio, norm_layer=norm_layer)
            self.post_norm = norm_layer(encoder_embed_dim)

    def forward(self, x, x_shared, ids_keep, weight_shared=0.5, idx=0):

        # x = x * (1 - weight_shared) + x_shared * weight_shared
        # x = x_shared
        if self.fusion_method == 'concat':
            x = self.prenorm(x)  # apply prenorm before the forward_features
            x = torch.cat([x, x_shared], dim=-1)
            x = self.merge(x)
        elif self.fusion_method == 'cross_attn':
            x = self.cross_attention(x, x_shared)  
            x = self.post_norm(x)     
        
        
        x = self.forward_features(x, ids_keep, idx=idx)
        # x = x * (1 - weight_shared) + x_shared * weight_shared

        # predictor projection
        x = self.decoder_pred(x)

        return x