import torch.nn as nn
import torch.nn.functional as F
from torch.nn.init import xavier_uniform_, constant_
import torch
import math
from layers import general_conv3d_prenorm, fusion_prenorm
from mamba_ssm import Mamba
from torch.cuda.amp import autocast
from utils.initialization import InitWeights_He

basic_dims = 8 #5 for SS
transformer_basic_dims = 512
mlp_dim = 4096
num_heads = 8
depth = 1
num_modals = 4
patch_size = 8
input_patch_size = 128

class MambaTrans(nn.Module):
    def __init__(self, channels):
        super(MambaTrans, self).__init__()
        self.mamba = Mamba(
            d_model=channels,
            d_state=min(channels, 256),
            d_conv=4,
            expand=2,
        )
        self.norm1 = nn.LayerNorm(channels,)
        self.norm2 = nn.LayerNorm(channels,)
        self.head = nn.Linear(channels, channels)

    def forward(self, x):
        x = self.mamba(self.norm1(x)) + x
        x = self.head(self.norm2(x)) + x
        return x

class MambaLayer(nn.Module):
    def __init__(self, dim, d_state = 16, d_conv = 4, expand = 2):
        super().__init__()
        self.dim = dim
        # self.norm = nn.LayerNorm(dim)
        self.mamba = MambaTrans(dim)
    
    @autocast(enabled=False)
    def forward(self, x):
        if x.dtype == torch.float16:
            x = x.type(torch.float32)        
        # x = self.norm(x)
        x_mamba = self.mamba(x)

        return x_mamba
    
class MambaFusionLayer(nn.Module):
    def __init__(self, dim, num_tokens_fused_representation=None):
        super().__init__()
        self.dim = dim
        self.num_tokens_fused_representation = num_tokens_fused_representation
        self.fused_tokens = nn.Parameter(torch.randn(1, self.num_tokens_fused_representation, dim))
        self.mamba_layer = MambaLayer(dim)
    def forward(self, x): # (B, 2048, 512)
        B = x.size(0)
        fused_tokens = self.fused_tokens.repeat(B, 1, 1)
        x_fused = torch.cat([x, fused_tokens], dim=1) # (B, 2048+512, 512)
        x_mamba = self.mamba_layer(x_fused)
        x_mamba = x_mamba[:, -self.num_tokens_fused_representation:, :] # (B, 512, 512)    
        return x_mamba

class MambaFusionCatLayer(nn.Module):
    def __init__(self, dim, num_tokens_fused_representation=None):
        super().__init__()
        self.dim = dim
        self.num_tokens_fused_representation = num_tokens_fused_representation
        self.fused_tokens = nn.Parameter(torch.randn(1, self.num_tokens_fused_representation, dim))
        self.mamba_layer = MambaLayer(dim)
    def forward(self, x): # [(B, 512, 512)]*4
        B = x[0].size(0)
        fused_tokens = self.fused_tokens.repeat(B, 1, 1)
        x = torch.stack([*x, fused_tokens], dim=2) # (B, 512, 5, 512)
        x = x.view(B, -1, self.dim) # (B, 2048+512, 512)
        x = self.mamba_layer(x)
        x = x[:, 4::5, :]
        return x # (B, 512, 512)    

class Tokenize(nn.Module):
    def __init__(self, dims, num_modals=4):
        super(Tokenize, self).__init__()
        self.dims = dims
        self.num_modals = num_modals
    def forward(self, x):
        flair_intra_x, t1ce_intra_x, t1_intra_x, t2_intra_x = torch.chunk(x, self.num_modals, dim=1) #(B, 512, 8, 8, 8)
        multimodal_token_x = torch.cat((flair_intra_x.permute(0, 2, 3, 4, 1).contiguous().view(x.size(0), -1, self.dims), #(B, 512, 512)
                                         t1ce_intra_x.permute(0, 2, 3, 4, 1).contiguous().view(x.size(0), -1, self.dims),
                                         t1_intra_x.permute(0, 2, 3, 4, 1).contiguous().view(x.size(0), -1, self.dims),
                                         t2_intra_x.permute(0, 2, 3, 4, 1).contiguous().view(x.size(0), -1, self.dims),
                                        ), dim=1)
        return multimodal_token_x
    
class TokenizeSep(nn.Module):
    def __init__(self, dims, num_modals=4):
        super(TokenizeSep, self).__init__()
        self.dims = dims
        self.num_modals = num_modals
    def forward(self, x):
        flair_intra_x, t1ce_intra_x, t1_intra_x, t2_intra_x = torch.chunk(x, self.num_modals, dim=1) #(B, 512, 8, 8, 8)
        return  flair_intra_x.permute(0, 2, 3, 4, 1).contiguous().view(x.size(0), -1, self.dims), \
                t1ce_intra_x.permute(0, 2, 3, 4, 1).contiguous().view(x.size(0), -1, self.dims), \
                t1_intra_x.permute(0, 2, 3, 4, 1).contiguous().view(x.size(0), -1, self.dims), \
                t2_intra_x.permute(0, 2, 3, 4, 1).contiguous().view(x.size(0), -1, self.dims)        


class Encoder(nn.Module):
    def __init__(self):
        super(Encoder, self).__init__()

        self.e1_c1 = nn.Conv3d(in_channels=1, out_channels=basic_dims, kernel_size=3, stride=1, padding=1, padding_mode='reflect', bias=True)
        self.e1_c2 = general_conv3d_prenorm(basic_dims, basic_dims, pad_type='reflect')
        self.e1_c3 = general_conv3d_prenorm(basic_dims, basic_dims, pad_type='reflect')

        self.e2_c1 = general_conv3d_prenorm(basic_dims, basic_dims*2, stride=2, pad_type='reflect')
        self.e2_c2 = general_conv3d_prenorm(basic_dims*2, basic_dims*2, pad_type='reflect')
        self.e2_c3 = general_conv3d_prenorm(basic_dims*2, basic_dims*2, pad_type='reflect')

        self.e3_c1 = general_conv3d_prenorm(basic_dims*2, basic_dims*4, stride=2, pad_type='reflect')
        self.e3_c2 = general_conv3d_prenorm(basic_dims*4, basic_dims*4, pad_type='reflect')
        self.e3_c3 = general_conv3d_prenorm(basic_dims*4, basic_dims*4, pad_type='reflect')

        self.e4_c1 = general_conv3d_prenorm(basic_dims*4, basic_dims*8, stride=2, pad_type='reflect')
        self.e4_c2 = general_conv3d_prenorm(basic_dims*8, basic_dims*8, pad_type='reflect')
        self.e4_c3 = general_conv3d_prenorm(basic_dims*8, basic_dims*8, pad_type='reflect')

        self.e5_c1 = general_conv3d_prenorm(basic_dims*8, basic_dims*16, stride=2, pad_type='reflect')
        self.e5_c2 = general_conv3d_prenorm(basic_dims*16, basic_dims*16, pad_type='reflect')
        self.e5_c3 = general_conv3d_prenorm(basic_dims*16, basic_dims*16, pad_type='reflect')

    def forward(self, x):
        x1 = self.e1_c1(x)
        x1 = x1 + self.e1_c3(self.e1_c2(x1)) #(B, 8, 128, 128, 128)

        x2 = self.e2_c1(x1)
        x2 = x2 + self.e2_c3(self.e2_c2(x2)) #(B, 64, 64, 64, 64)

        x3 = self.e3_c1(x2)
        x3 = x3 + self.e3_c3(self.e3_c2(x3)) #(B, 32, 32, 32, 32)

        x4 = self.e4_c1(x3)
        x4 = x4 + self.e4_c3(self.e4_c2(x4)) #(B, 64, 16, 16, 16)

        x5 = self.e5_c1(x4)
        x5 = x5 + self.e5_c3(self.e5_c2(x5)) #(B, 128, 8, 8, 8)

        return x1, x2, x3, x4, x5

class Decoder_sep(nn.Module):
    def __init__(self, num_cls=4):
        super(Decoder_sep, self).__init__()

        self.d4 = nn.Upsample(scale_factor=2, mode='trilinear', align_corners=True)
        self.d4_c1 = general_conv3d_prenorm(basic_dims*16, basic_dims*8, pad_type='reflect')
        self.d4_c2 = general_conv3d_prenorm(basic_dims*16, basic_dims*8, pad_type='reflect')
        self.d4_out = general_conv3d_prenorm(basic_dims*8, basic_dims*8, k_size=1, padding=0, pad_type='reflect')

        self.d3 = nn.Upsample(scale_factor=2, mode='trilinear', align_corners=True)
        self.d3_c1 = general_conv3d_prenorm(basic_dims*8, basic_dims*4, pad_type='reflect')
        self.d3_c2 = general_conv3d_prenorm(basic_dims*8, basic_dims*4, pad_type='reflect')
        self.d3_out = general_conv3d_prenorm(basic_dims*4, basic_dims*4, k_size=1, padding=0, pad_type='reflect')

        self.d2 = nn.Upsample(scale_factor=2, mode='trilinear', align_corners=True)
        self.d2_c1 = general_conv3d_prenorm(basic_dims*4, basic_dims*2, pad_type='reflect')
        self.d2_c2 = general_conv3d_prenorm(basic_dims*4, basic_dims*2, pad_type='reflect')
        self.d2_out = general_conv3d_prenorm(basic_dims*2, basic_dims*2, k_size=1, padding=0, pad_type='reflect')

        self.d1 = nn.Upsample(scale_factor=2, mode='trilinear', align_corners=True)
        self.d1_c1 = general_conv3d_prenorm(basic_dims*2, basic_dims, pad_type='reflect')
        self.d1_c2 = general_conv3d_prenorm(basic_dims*2, basic_dims, pad_type='reflect')
        self.d1_out = general_conv3d_prenorm(basic_dims, basic_dims, k_size=1, padding=0, pad_type='reflect')

        self.seg_layer = nn.Conv3d(in_channels=basic_dims, out_channels=num_cls, kernel_size=1, stride=1, padding=0, bias=True)
        self.softmax = nn.Softmax(dim=1)

    def forward(self, x1, x2, x3, x4, x5):
        de_x5 = self.d4_c1(self.d4(x5)) #(B, 64, 16, 16, 16)

        cat_x4 = torch.cat((de_x5, x4), dim=1)
        de_x4 = self.d4_out(self.d4_c2(cat_x4))
        de_x4 = self.d3_c1(self.d3(de_x4)) #(B, 32, 32, 32, 32)

        cat_x3 = torch.cat((de_x4, x3), dim=1)
        de_x3 = self.d3_out(self.d3_c2(cat_x3))
        de_x3 = self.d2_c1(self.d2(de_x3)) #(B, 16, 64, 64, 64)

        cat_x2 = torch.cat((de_x3, x2), dim=1)
        de_x2 = self.d2_out(self.d2_c2(cat_x2))
        de_x2 = self.d1_c1(self.d1(de_x2))

        cat_x1 = torch.cat((de_x2, x1), dim=1)
        de_x1 = self.d1_out(self.d1_c2(cat_x1)) #(B, 8, 128, 128, 128)

        logits = self.seg_layer(de_x1)
        pred = self.softmax(logits) #(B, C, 128, 128, 128)

        return pred 

class Decoder_fuse(nn.Module):
    def __init__(self, num_cls=4, mamba_skip=False):
        super(Decoder_fuse, self).__init__()

        self.d4_c1 = general_conv3d_prenorm(basic_dims*16, basic_dims*8, pad_type='reflect')
        self.d4_c2 = general_conv3d_prenorm(basic_dims*16, basic_dims*8, pad_type='reflect')
        self.d4_out = general_conv3d_prenorm(basic_dims*8, basic_dims*8, k_size=1, padding=0, pad_type='reflect')

        self.d3_c1 = general_conv3d_prenorm(basic_dims*8, basic_dims*4, pad_type='reflect')
        self.d3_c2 = general_conv3d_prenorm(basic_dims*8, basic_dims*4, pad_type='reflect')
        self.d3_out = general_conv3d_prenorm(basic_dims*4, basic_dims*4, k_size=1, padding=0, pad_type='reflect')

        self.d2_c1 = general_conv3d_prenorm(basic_dims*4, basic_dims*2, pad_type='reflect')
        self.d2_c2 = general_conv3d_prenorm(basic_dims*4, basic_dims*2, pad_type='reflect')
        self.d2_out = general_conv3d_prenorm(basic_dims*2, basic_dims*2, k_size=1, padding=0, pad_type='reflect')

        self.d1_c1 = general_conv3d_prenorm(basic_dims*2, basic_dims, pad_type='reflect')
        self.d1_c2 = general_conv3d_prenorm(basic_dims, basic_dims, pad_type='reflect')
        self.d1_out = general_conv3d_prenorm(basic_dims, basic_dims, k_size=1, padding=0, pad_type='reflect')

        self.seg_d4 = nn.Conv3d(in_channels=basic_dims*16, out_channels=num_cls, kernel_size=1, stride=1, padding=0, bias=True)
        self.seg_d3 = nn.Conv3d(in_channels=basic_dims*8, out_channels=num_cls, kernel_size=1, stride=1, padding=0, bias=True)
        self.seg_d2 = nn.Conv3d(in_channels=basic_dims*4, out_channels=num_cls, kernel_size=1, stride=1, padding=0, bias=True)
        self.seg_d1 = nn.Conv3d(in_channels=basic_dims*2, out_channels=num_cls, kernel_size=1, stride=1, padding=0, bias=True)
        self.seg_layer = nn.Conv3d(in_channels=basic_dims, out_channels=num_cls, kernel_size=1, stride=1, padding=0, bias=True)
        self.softmax = nn.Softmax(dim=1)

        self.up2 = nn.Upsample(scale_factor=2, mode='trilinear', align_corners=True)
        self.up4 = nn.Upsample(scale_factor=4, mode='trilinear', align_corners=True)
        self.up8 = nn.Upsample(scale_factor=8, mode='trilinear', align_corners=True)
        self.up16 = nn.Upsample(scale_factor=16, mode='trilinear', align_corners=True)

        self.RFM5 = fusion_prenorm(in_channel=basic_dims*16, num_cls=num_cls)
        self.RFM4 = fusion_prenorm(in_channel=basic_dims*8, num_cls= 1 if mamba_skip else num_cls)
        self.RFM3 = fusion_prenorm(in_channel=basic_dims*4, num_cls= 1 if mamba_skip else num_cls)
        self.RFM2 = fusion_prenorm(in_channel=basic_dims*2, num_cls= 1 if mamba_skip else num_cls)
        # self.RFM1 = fusion_prenorm(in_channel=basic_dims*1, num_cls=1 if mamba_skip else num_cls)
        self.mamba_skip = mamba_skip


    def forward(self, x1, x2, x3, x4, x5):
        de_x5 = self.RFM5(x5)                       #(B, 128, 8, 8, 8)
        pred4 = self.softmax(self.seg_d4(de_x5))
        de_x5 = self.d4_c1(self.up2(de_x5))         #(B, 64, 16, 16, 16)

        de_x4 = self.RFM4(x4)                       #(B, 64, 16, 16, 16)
        de_x4 = torch.cat((de_x4, de_x5), dim=1)    #(B, 128, 16, 16, 16)    
        de_x4 = self.d4_out(self.d4_c2(de_x4))      #(B, 64, 16, 16, 16) 
        pred3 = self.softmax(self.seg_d3(de_x4))
        de_x4 = self.d3_c1(self.up2(de_x4))         #(B, 32, 32, 32, 32) 

        de_x3 = self.RFM3(x3)                       #(B, 32, 32, 32, 32)
        de_x3 = torch.cat((de_x3, de_x4), dim=1)    #(B, 64, 32, 32, 32)
        de_x3 = self.d3_out(self.d3_c2(de_x3))      #(B, 32, 32, 32, 32)
        pred2 = self.softmax(self.seg_d2(de_x3))        
        de_x3 = self.d2_c1(self.up2(de_x3))         #(B, 16, 64, 64, 64)

        de_x2 = self.RFM2(x2)                       #(B, 16, 64, 64, 64)
        de_x2 = torch.cat((de_x2, de_x3), dim=1)    #(B, 32, 64, 64, 64)
        de_x2 = self.d2_out(self.d2_c2(de_x2))      #(B, 16, 64, 64, 64)
        pred1 = self.softmax(self.seg_d1(de_x2))    
        de_x2 = self.d1_c1(self.up2(de_x2))         #(B, 8, 128, 128, 128)

        # de_x1 = self.RFM1(x1)                     #(B, 8, 128, 128, 128)
        # de_x1 = torch.cat((de_x1, de_x2), dim=1)  #(B, 16, 128, 128, 128)
        de_x1 = self.d1_out(self.d1_c2(de_x2))      #(B, 8, 128, 128, 128)

        logits = self.seg_layer(de_x1)              #(B, 4, 128, 128, 128)
        pred = self.softmax(logits)                 #(B, 4, 128, 128, 128)

        return pred, (self.up2(pred1), self.up4(pred2), self.up8(pred3), self.up16(pred4))

class SelfAttention(nn.Module):
    def __init__(
        self, dim, heads=8, qkv_bias=False, qk_scale=None, dropout_rate=0.0
    ):
        super().__init__()
        self.num_heads = heads
        head_dim = dim // heads
        self.scale = qk_scale or head_dim ** -0.5

        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
        self.attn_drop = nn.Dropout(dropout_rate)
        self.proj = nn.Linear(dim, dim)
        self.proj_drop = nn.Dropout(dropout_rate)

    def forward(self, x):
        B, N, C = x.shape
        qkv = (
            self.qkv(x)
            .reshape(B, N, 3, self.num_heads, C // self.num_heads)
            .permute(2, 0, 3, 1, 4)
        )
        q, k, v = (
            qkv[0],
            qkv[1],
            qkv[2],
        )  # make torchscript happy (cannot use tensor as tuple)

        attn = (q @ k.transpose(-2, -1)) * self.scale
        attn = attn.softmax(dim=-1)
        attn = self.attn_drop(attn)

        x = (attn @ v).transpose(1, 2).reshape(B, N, C)
        x = self.proj(x)
        x = self.proj_drop(x)
        return x

class Residual(nn.Module):
    def __init__(self, fn):
        super().__init__()
        self.fn = fn

    def forward(self, x):
        return self.fn(x) + x

class PreNorm(nn.Module):
    def __init__(self, dim, fn):
        super().__init__()
        self.norm = nn.LayerNorm(dim)
        self.fn = fn

    def forward(self, x):
        return self.fn(self.norm(x))

class PreNormDrop(nn.Module):
    def __init__(self, dim, dropout_rate, fn):
        super().__init__()
        self.norm = nn.LayerNorm(dim)
        self.dropout = nn.Dropout(p=dropout_rate)
        self.fn = fn

    def forward(self, x):
        return self.dropout(self.fn(self.norm(x)))

class GELU(nn.Module):
    def __init__(self):
        super().__init__()
    
    def forward(self, x):
        return F.gelu(x)

class FeedForward(nn.Module):
    def __init__(self, dim, hidden_dim, dropout_rate):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(dim, hidden_dim),
            GELU(),
            nn.Dropout(p=dropout_rate),
            nn.Linear(hidden_dim, dim),
            nn.Dropout(p=dropout_rate),
        )

    def forward(self, x):
        return self.net(x)

class Transformer(nn.Module):
    def __init__(self, embedding_dim, depth, heads, mlp_dim, dropout_rate=0.1, n_levels=1, n_points=4):
        super(Transformer, self).__init__()
        self.cross_attention_list = []
        self.cross_ffn_list = []
        self.depth = depth
        for j in range(self.depth):
            self.cross_attention_list.append(
                Residual(
                    PreNormDrop(
                        embedding_dim,
                        dropout_rate,
                        SelfAttention(embedding_dim, heads=heads, dropout_rate=dropout_rate),
                    )
                )
            )
            self.cross_ffn_list.append(
                Residual(
                    PreNorm(embedding_dim, FeedForward(embedding_dim, mlp_dim, dropout_rate))
                )
            )

        self.cross_attention_list = nn.ModuleList(self.cross_attention_list)
        self.cross_ffn_list = nn.ModuleList(self.cross_ffn_list)


    def forward(self, x, pos):
        for j in range(self.depth):
            x = x + pos
            x = self.cross_attention_list[j](x)
            x = self.cross_ffn_list[j](x)
        return x


class MaskModal(nn.Module):
    def __init__(self):
        super(MaskModal, self).__init__()
    
    def forward(self, x, mask):
        B, K, C, H, W, Z = x.size()
        y = torch.zeros_like(x)
        y[mask, ...] = x[mask, ...]
        x = y.view(B, -1, H, W, Z)
        return x
    

                                          
class Model(nn.Module):
    def __init__(self, num_cls=4, interleaved_tokenization=False, mamba_skip=False):
        super(Model, self).__init__()
        self.interleaved_tokenization = interleaved_tokenization

        self.flair_encoder = Encoder()
        self.t1ce_encoder = Encoder()
        self.t1_encoder = Encoder()
        self.t2_encoder = Encoder()

        if self.interleaved_tokenization:
            TokenizerClass = TokenizeSep
            MambaFusionLayerClass = MambaFusionCatLayer
        else:
            TokenizerClass = Tokenize
            MambaFusionLayerClass = MambaFusionLayer


        ########### IntraFormer
        self.flair_encode_conv = nn.Conv3d(basic_dims*16, transformer_basic_dims, kernel_size=1, stride=1, padding=0)
        self.t1ce_encode_conv = nn.Conv3d(basic_dims*16, transformer_basic_dims, kernel_size=1, stride=1, padding=0)
        self.t1_encode_conv = nn.Conv3d(basic_dims*16, transformer_basic_dims, kernel_size=1, stride=1, padding=0)
        self.t2_encode_conv = nn.Conv3d(basic_dims*16, transformer_basic_dims, kernel_size=1, stride=1, padding=0)
        
        #self.flair_decode_conv = nn.Conv3d(transformer_basic_dims, basic_dims*16, kernel_size=1, stride=1, padding=0)
        #self.t1ce_decode_conv = nn.Conv3d(transformer_basic_dims, basic_dims*16, kernel_size=1, stride=1, padding=0)
        #self.t1_decode_conv = nn.Conv3d(transformer_basic_dims, basic_dims*16, kernel_size=1, stride=1, padding=0)
        #self.t2_decode_conv = nn.Conv3d(transformer_basic_dims, basic_dims*16, kernel_size=1, stride=1, padding=0)

        self.flair_pos = nn.Parameter(torch.zeros(1, patch_size**3, transformer_basic_dims))
        self.t1ce_pos = nn.Parameter(torch.zeros(1, patch_size**3, transformer_basic_dims))
        self.t1_pos = nn.Parameter(torch.zeros(1, patch_size**3, transformer_basic_dims))
        self.t2_pos = nn.Parameter(torch.zeros(1, patch_size**3, transformer_basic_dims))
        self.fused_pos = nn.Parameter(torch.zeros(1, patch_size**3, transformer_basic_dims))

        self.flair_transformer = Transformer(embedding_dim=transformer_basic_dims, depth=depth, heads=num_heads, mlp_dim=mlp_dim)
        self.t1ce_transformer = Transformer(embedding_dim=transformer_basic_dims, depth=depth, heads=num_heads, mlp_dim=mlp_dim)
        self.t1_transformer = Transformer(embedding_dim=transformer_basic_dims, depth=depth, heads=num_heads, mlp_dim=mlp_dim)
        self.t2_transformer = Transformer(embedding_dim=transformer_basic_dims, depth=depth, heads=num_heads, mlp_dim=mlp_dim)
        ########### IntraFormer

        ########### InterFormer
        self.mamba_fusion_layer = MambaFusionLayer(dim=transformer_basic_dims, num_tokens_fused_representation=patch_size**3)
        self.multimodal_transformer = Transformer(embedding_dim=transformer_basic_dims, depth=depth, heads=num_heads, mlp_dim=mlp_dim, n_levels=num_modals)
        self.multimodal_decode_conv = nn.Conv3d(transformer_basic_dims, basic_dims*16*num_modals, kernel_size=1, padding=0)
        ########### InterFormer

        self.masker = MaskModal()

        ######## Skip Connections
        if not mamba_skip:
            self.tokenize = nn.ModuleList([
                TokenizerClass(dims=8, num_modals=num_modals), #(B, 8, 128, 128, 128)->(B, 128**3, 8)
                TokenizerClass(dims=16, num_modals=num_modals),  #(B, 16, 64, 64, 64)->(B, 64**3, 16)
                TokenizerClass(dims=32, num_modals=num_modals),  #(B, 32, 32, 32, 32)->(B, 32**3, 32)
                TokenizerClass(dims=64, num_modals=num_modals),  #(B, 64, 16, 16, 16)->(B, 16**3, 64)
                TokenizerClass(dims=512, num_modals=num_modals),   #(B, 512, 8, 8, 8)->(B, 8**3, 512)
            ])
            self.mamba_fusion_layers = nn.ModuleList([
                MambaFusionLayerClass(dim=8, num_tokens_fused_representation=128**3), #(B, 128**3, 8)->(B, 128**3, 8)
                MambaFusionLayerClass(dim=16, num_tokens_fused_representation=64**3), #(B, 64**3, 16)->(B, 64**3, 16)
                MambaFusionLayerClass(dim=32, num_tokens_fused_representation=32**3), #(B, 32**3, 32)->(B, 32**3, 32)
                MambaFusionLayerClass(dim=64, num_tokens_fused_representation=16**3), #(B, 16**3, 64)->(B, 16**3, 64)
                MambaFusionLayerClass(dim=512, num_tokens_fused_representation=8**3), #(B, 8**3, 512)->(B, 8**3, 512)
            ])
        else:
            self.tokenize = nn.ModuleList([
                #TokenizerClass(dims=8, num_modals=num_modals), #(B, 8, 128, 128, 128)->(B, 128**3, 8)
                TokenizerClass(dims=16, num_modals=num_modals),  #(B, 16, 64, 64, 64)->(B, 64**3, 16)
                TokenizerClass(dims=32, num_modals=num_modals),  #(B, 32, 32, 32, 32)->(B, 32**3, 32)
                TokenizerClass(dims=64, num_modals=num_modals),  #(B, 64, 16, 16, 16)->(B, 16**3, 64)
                TokenizerClass(dims=512, num_modals=num_modals),   #(B, 512, 8, 8, 8)->(B, 8**3, 512)
            ])
            self.mamba_fusion_layers = nn.ModuleList([
                #MambaFusionLayerClass(dim=8, num_tokens_fused_representation=128**3), #(B, 128**3, 8)->(B, 128**3, 8)
                MambaFusionLayerClass(dim=16, num_tokens_fused_representation=64**3), #(B, 64**3, 16)->(B, 64**3, 16)
                MambaFusionLayerClass(dim=32, num_tokens_fused_representation=32**3), #(B, 32**3, 32)->(B, 32**3, 32)
                MambaFusionLayerClass(dim=64, num_tokens_fused_representation=16**3), #(B, 16**3, 64)->(B, 16**3, 64)
                MambaFusionLayerClass(dim=512, num_tokens_fused_representation=8**3), #(B, 8**3, 512)->(B, 8**3, 512)
            ])
        ########

        self.decoder_fuse = Decoder_fuse(num_cls=num_cls, mamba_skip=mamba_skip)
        self.decoder_sep = Decoder_sep(num_cls=num_cls)

        self.is_training = False
        self.mamba_skip = mamba_skip

        # for m in self.modules():
        #     if isinstance(m, nn.Conv3d):
        #         torch.nn.init.kaiming_normal_(m.weight) #
        self.apply(InitWeights_He(1e-2))

    def forward(self, x, mask):
        #extract feature from different layers
        flair_x1, flair_x2, flair_x3, flair_x4, flair_x5 = self.flair_encoder(x[:, 0:1, :, :, :])
        t1ce_x1, t1ce_x2, t1ce_x3, t1ce_x4, t1ce_x5 = self.t1ce_encoder(x[:, 1:2, :, :, :])
        t1_x1, t1_x2, t1_x3, t1_x4, t1_x5 = self.t1_encoder(x[:, 2:3, :, :, :])
        t2_x1, t2_x2, t2_x3, t2_x4, t2_x5 = self.t2_encoder(x[:, 3:4, :, :, :])

        ########### IntraFormer
        flair_token_x5 = self.flair_encode_conv(flair_x5).permute(0, 2, 3, 4, 1).contiguous().view(x.size(0), -1, transformer_basic_dims) #(B, 512, 512)
        t1ce_token_x5 = self.t1ce_encode_conv(t1ce_x5).permute(0, 2, 3, 4, 1).contiguous().view(x.size(0), -1, transformer_basic_dims)
        t1_token_x5 = self.t1_encode_conv(t1_x5).permute(0, 2, 3, 4, 1).contiguous().view(x.size(0), -1, transformer_basic_dims)
        t2_token_x5 = self.t2_encode_conv(t2_x5).permute(0, 2, 3, 4, 1).contiguous().view(x.size(0), -1, transformer_basic_dims)

        flair_intra_token_x5 = self.flair_transformer(flair_token_x5, self.flair_pos)
        t1ce_intra_token_x5 = self.t1ce_transformer(t1ce_token_x5, self.t1ce_pos)
        t1_intra_token_x5 = self.t1_transformer(t1_token_x5, self.t1_pos)
        t2_intra_token_x5 = self.t2_transformer(t2_token_x5, self.t2_pos)

        flair_intra_x5 = flair_intra_token_x5.view(x.size(0), patch_size, patch_size, patch_size, transformer_basic_dims).permute(0, 4, 1, 2, 3).contiguous() #(B, 512, 8, 8, 8)
        t1ce_intra_x5 = t1ce_intra_token_x5.view(x.size(0), patch_size, patch_size, patch_size, transformer_basic_dims).permute(0, 4, 1, 2, 3).contiguous()
        t1_intra_x5 = t1_intra_token_x5.view(x.size(0), patch_size, patch_size, patch_size, transformer_basic_dims).permute(0, 4, 1, 2, 3).contiguous()
        t2_intra_x5 = t2_intra_token_x5.view(x.size(0), patch_size, patch_size, patch_size, transformer_basic_dims).permute(0, 4, 1, 2, 3).contiguous()

        if self.is_training:
            flair_pred = self.decoder_sep(flair_x1, flair_x2, flair_x3, flair_x4, flair_x5) #(B, C, 128, 128, 128)
            t1ce_pred = self.decoder_sep(t1ce_x1, t1ce_x2, t1ce_x3, t1ce_x4, t1ce_x5)
            t1_pred = self.decoder_sep(t1_x1, t1_x2, t1_x3, t1_x4, t1_x5)
            t2_pred = self.decoder_sep(t2_x1, t2_x2, t2_x3, t2_x4, t2_x5)
        ########### IntraFormer

        x1 = self.masker(torch.stack((flair_x1, t1ce_x1, t1_x1, t2_x1), dim=1), mask) #Bx4xCxHWZ = (B, 32, 128, 128, 128)
        x2 = self.masker(torch.stack((flair_x2, t1ce_x2, t1_x2, t2_x2), dim=1), mask)
        x3 = self.masker(torch.stack((flair_x3, t1ce_x3, t1_x3, t2_x3), dim=1), mask)
        x4 = self.masker(torch.stack((flair_x4, t1ce_x4, t1_x4, t2_x4), dim=1), mask)
        x5_intra = self.masker(torch.stack((flair_intra_x5, t1ce_intra_x5, t1_intra_x5, t2_intra_x5), dim=1), mask) #(1, 2048, 8, 8, 8)

        if self.mamba_skip:
            #x1 = self.tokenize[-5](x1) #(B, 128**3, 8)*4
            #x1 = self.mamba_fusion_layers[-5](x1) # (B, 128**3, 8)
            #x1 = x1.view(x.size(0), input_patch_size, input_patch_size, input_patch_size, basic_dims).permute(0, 4, 1, 2, 3).contiguous() #(B, 8, 128, 128, 128)

            x2 = self.tokenize[-4](x2) # (B, 64**3, 16)*4
            x2 = self.mamba_fusion_layers[-4](x2) # (B, 64**3, 16)
            x2 = x2.view(x.size(0), input_patch_size//2, input_patch_size//2, input_patch_size//2, basic_dims*2).permute(0, 4, 1, 2, 3).contiguous() #(B, 16, 64, 64, 64)

            x3 = self.tokenize[-3](x3) # (B, 32**3, 32)*4
            x3 = self.mamba_fusion_layers[-3](x3) # (B, 32**3, 32)
            x3 = x3.view(x.size(0), input_patch_size//4, input_patch_size//4, input_patch_size//4, basic_dims*4).permute(0, 4, 1, 2, 3).contiguous() #(B, 32, 32, 32, 32)

            x4 = self.tokenize[-2](x4) # (B, 16**3, 64)*4
            x4 = self.mamba_fusion_layers[-2](x4) # (B, 16**3, 64)
            x4 = x4.view(x.size(0), input_patch_size//8, input_patch_size//8, input_patch_size//8, basic_dims*8).permute(0, 4, 1, 2, 3).contiguous() #(B, 64, 16, 16, 16)
        #######        

        ########### InterFormer
        # flair_intra_x5, t1ce_intra_x5, t1_intra_x5, t2_intra_x5 = torch.chunk(x5_intra, num_modals, dim=1) #(B, 512, 8, 8, 8)
        # multimodal_token_x5 = torch.cat((flair_intra_x5.permute(0, 2, 3, 4, 1).contiguous().view(x.size(0), -1, transformer_basic_dims), #(B, 512, 512)
        #                                  t1ce_intra_x5.permute(0, 2, 3, 4, 1).contiguous().view(x.size(0), -1, transformer_basic_dims),
        #                                  t1_intra_x5.permute(0, 2, 3, 4, 1).contiguous().view(x.size(0), -1, transformer_basic_dims),
        #                                  t2_intra_x5.permute(0, 2, 3, 4, 1).contiguous().view(x.size(0), -1, transformer_basic_dims),
        #                                 ), dim=1)
        # multimodal_pos = torch.cat((self.flair_pos, self.t1ce_pos, self.t1_pos, self.t2_pos), dim=1) #(B, 2048, 512)
        multimodal_token_x5 = self.tokenize[-1](x5_intra) # (B, 4*8**3, 128)
        fused_multimodal = self.mamba_fusion_layers[-1](multimodal_token_x5)
        multimodal_pos = self.fused_pos.repeat(x.size(0), 1, 1)
        multimodal_inter_token_x5 = self.multimodal_transformer(fused_multimodal, multimodal_pos)
        multimodal_inter_x5 = self.multimodal_decode_conv(multimodal_inter_token_x5.view(multimodal_inter_token_x5.size(0), patch_size, patch_size, patch_size, transformer_basic_dims).permute(0, 4, 1, 2, 3).contiguous()) #(B, 512, 8, 8, 8) -> (B, 512, 8, 8, 8)
        x5_inter = multimodal_inter_x5

        fuse_pred, preds = self.decoder_fuse(x1, x2, x3, x4, x5_inter)
        ########### InterFormer
        
        if self.is_training:
            return fuse_pred, (flair_pred, t1ce_pred, t1_pred, t2_pred), preds
        return fuse_pred
