import torch
import torch.nn as nn
from models.transformer_modules import BlockFlashAttn, Mlp, freq_pos_enc, PatchUpsample

class FusionModule(nn.Module):
    def __init__(self, encoder_embed_dim=128, 
                num_heads=4,  mlp_ratio=4, norm_layer=nn.LayerNorm, fusion_depth=2,
                fusion_method='concat', pos_enc='sinusoidal'):
        super().__init__()
        self.fusion_method = fusion_method
        self.num_heads = num_heads

        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)        
            self.merge = Mlp(in_features=encoder_embed_dim * 2, out_features=encoder_embed_dim)        
        elif self.fusion_method == 'cross_attn':
            self.cross_attention_blocks =  nn.ModuleList([
                BlockFlashAttn(encoder_embed_dim, num_heads, mlp_ratio, qkv_bias=True, qk_scale=None, norm_layer=norm_layer)
                # BlockFlashAttn(encoder_embed_dim, num_heads, mlp_ratio, qkv_bias=True, qk_scale=None, norm_layer=norm_layer)
            for i in range(fusion_depth)])
        self.post_norm = norm_layer(encoder_embed_dim)
        
        self.pos_enc = pos_enc
        if pos_enc == 'rope':
            self.rope = freq_pos_enc(dim=encoder_embed_dim // num_heads, max_seq_len=128)

    def forward(self, x, y):

        if self.fusion_method == 'concat':            
            x = torch.cat([x, y], dim=-1)
            x = self.prenorm(x)  # apply prenorm before the forward_features
            x = self.merge(x)
            
        elif self.fusion_method == 'cross_attn':
            
            if self.pos_enc == 'rope':
                
                pe = self.rope.unsqueeze(0).to(dtype=x.dtype, device=x.device)  # (1, max_seq_len, dim)
                pe = pe[:, :x.shape[1], :]  # (1, seq_len, dim)
                pe = pe.repeat(x.shape[0], 1, 1)  # (N, max_seq_len, dim)
                pe2 = None
                if not x.shape[1] == y.shape[1]:                    
                    pe2 = freq_pos_enc(dim=y.shape[2] // self.num_heads, max_seq_len=y.shape[1])
                    pe2 = pe2.unsqueeze(0).to(dtype=x.dtype, device=x.device)  # (1, max_seq_len, dim)
                    pe2 = pe2.repeat(x.shape[0], 1, 1)  # (N, max_seq_len, dim)
                
            else:
                pe = None
            
            for blk in self.cross_attention_blocks:
                x, _ = blk(x, y, rope=pe, rope2=pe2)
        x = self.post_norm(x)
        return x

class FusionUpsampleBlock(nn.Module):
    def __init__(self, dim_list, fusion_method='cross_attn', fusion_depth=2, fusion_heads=4,
                 mlp_ratio=4., norm_layer=nn.LayerNorm, pos_enc='sinusoidal', 
                 use_shared=True, use_upsample=True, modalities=["EEG","EOG"]):
        super().__init__()
        
        self.fusion_method = fusion_method
        self.use_shared = use_shared
        self.use_upsample = use_upsample
        self.fusion = nn.ModuleDict()
        self.upsampler = nn.ModuleDict()
        self.dim_list = dim_list
        self.modalities = modalities
        
        for m in self.modalities:                  
            if use_shared:
                self.fusion[m] = FusionModule(encoder_embed_dim=dim_list[-1],
                                            mlp_ratio=mlp_ratio, norm_layer=norm_layer, 
                                            fusion_method=fusion_method, fusion_depth=fusion_depth, num_heads=fusion_heads, pos_enc=pos_enc)
            
            if use_upsample:
                upsampler_list = []
                for i in range(len(dim_list) - 1):
                    up = 2 if i > 0 else 5 
                    upsampler_list.append(PatchUpsample(embed_dim=dim_list[1+i], up_dim=dim_list[i], upsample=up, norm_layer=norm_layer))                    
                    
                self.upsampler[m] = nn.ModuleList(upsampler_list)
        

    def forward(self, x_shared, x_specific):
        
        latents_fused = []
        upsampled = []
        for i, m in enumerate(self.modalities):             
            if self.use_shared:                
                x_fused = self.fusion[m](x_specific[-1][i], x_shared[-1][i])
            else:
                x_fused = x_specific[-1][i]            
            latents_fused.append(x_fused)
            
            if self.use_upsample:
                for j in range(len(self.dim_list)-2, -1, -1):                                                        
                    x_fused = self.upsampler[m][j](x_fused, x_specific[j][i])                    
                upsampled.append(x_fused)
            
        return latents_fused, upsampled