#Modified from https://github.com/Wenxuan-1119/TransBTS
import torch
import torch.nn as nn
from models.BiTrUnet.Transformer import TransformerModel
from models.BiTrUnet.PositionalEncoding import FixedPositionalEncoding,LearnedPositionalEncoding,LearnedPositionalEncoding_4
from models.BiTrUnet.Unet_skipconnection import Unet
from models.BiTrUnet.unet import Att_EquiUnet
from models.BiTrUnet.layers import ConvBnRelu, UBlock, conv1x1, UBlockCbam, CBAM

#define the encoder class
class BiTransformerUnet(nn.Module):
    def __init__(
        self,
        img_dim,
        patch_dim,
        num_channels,
        embedding_dim_4,
        embedding_dim,
        num_heads,
        num_layers,
        hidden_dim,
        dropout_rate=0.0,
        attn_dropout_rate=0.0,
        conv_patch_representation=True,
        positional_encoding_type="learned",
    ):
        super(BiTrUnet, self).__init__()

        assert embedding_dim % num_heads == 0
        assert img_dim % patch_dim == 0

        self.img_dim = img_dim
        self.embedding_dim_4 = embedding_dim_4
        self.embedding_dim = embedding_dim
        self.num_heads = num_heads
        self.patch_dim = patch_dim
        self.num_channels = num_channels
        self.dropout_rate = dropout_rate
        self.attn_dropout_rate = attn_dropout_rate
        self.conv_patch_representation = conv_patch_representation

        self.num_patches = int((img_dim // patch_dim) ** 3)
        self.seq_length = self.num_patches
        self.flatten_dim = 128 * num_channels

        self.linear_encoding = nn.Linear(self.flatten_dim, self.embedding_dim)
        if positional_encoding_type == "learned":
            self.position_encoding = LearnedPositionalEncoding(
                self.seq_length, self.embedding_dim, self.seq_length
            )
            self.position_encoding_4 = LearnedPositionalEncoding_4(
                self.seq_length, self.embedding_dim_4, self.seq_length
            )
        elif positional_encoding_type == "fixed":
            self.position_encoding = FixedPositionalEncoding(
                self.embedding_dim,
            )
            self.position_encoding_4 = FixedPositionalEncoding(
                self.embedding_dim_4,
            )

        self.pe_dropout = nn.Dropout(p=self.dropout_rate)

        self.transformer = TransformerModel(
            embedding_dim,
            num_layers,
            num_heads,
            hidden_dim,

            self.dropout_rate,
            self.attn_dropout_rate,
        )

        self.transformer_4 = TransformerModel(
            embedding_dim_4,
            num_layers,
            num_heads,
            hidden_dim,

            self.dropout_rate,
            self.attn_dropout_rate,
        )
        
        self.pre_head_ln = nn.LayerNorm(embedding_dim)
        self.pre_head_ln_4 = nn.LayerNorm(embedding_dim_4)

        if self.conv_patch_representation:

            self.conv_x = nn.Conv3d(
                256,
                self.embedding_dim,
                kernel_size=3,
                stride=1,
                padding=1
            )

        if self.conv_patch_representation:

            self.conv_x_4 = nn.Conv3d(
                128,
                self.embedding_dim,
                kernel_size=3,
                stride=1,
                padding=1
            )

        #Encoder 
        self.Unet = Att_EquiUnet(inplanes = 16, num_classes = 4 , width = 16)
        self.bn = nn.BatchNorm3d(256)
        self.bn_4 = nn.BatchNorm3d(128)
        self.relu = nn.ReLU(inplace=True)


    def encode(self, x):
        if self.conv_patch_representation:
            # combine embedding with conv patch distribution
            x1_1, x2_1, x3_1, x4_1, x = self.Unet(x)

            x4_t = self.bn_4(x4_1) #
            x4_t= self.relu(x4_t) #
            x4_t= self.conv_x_4(x4_t) #
            x4_t = x4_t.permute(0, 2, 3, 4, 1).contiguous() #
            x4_t= x4_t.view(x4_t.size(0), -1, self.embedding_dim_4) #
            
            x = self.bn(x)
            x = self.relu(x)
            x = self.conv_x(x)
            x = x.permute(0, 2, 3, 4, 1).contiguous()
            x = x.view(x.size(0), -1, self.embedding_dim)

        else:
            x = self.Unet(x)
            x = self.bn(x)
            x = self.relu(x)
            x = (
                x.unfold(2, 2, 2)
                .unfold(3, 2, 2)
                .unfold(4, 2, 2)
                .contiguous()
            )
            x = x.view(x.size(0), x.size(1), -1, 8)
            x = x.permute(0, 2, 3, 1).contiguous()
            x = x.view(x.size(0), -1, self.flatten_dim)
            x = self.linear_encoding(x)

        x4_t = self.position_encoding_4(x4_t) #
        x4_t = self.pe_dropout(x4_t) #
        
        x = self.position_encoding(x)
        x = self.pe_dropout(x)

        # apply transformer
        x4_t, intmd_x4_t = self.transformer_4(x4_t) #
        x4_t = self.pre_head_ln_4(x4_t) #

        x, intmd_x = self.transformer(x)
        x = self.pre_head_ln(x)

        return x1_1, x2_1, x3_1, x4_1, x, intmd_x, x4_t, intmd_x4_t

    def decode(self, x):
        raise NotImplementedError("Should be implemented in child class!!")

    def forward(self, x, auxillary_output_layers=[1, 2, 3, 4], auxillary_output_layers_4=[1, 2, 3, 4]):

        x1_1, x2_1, x3_1, x4_1, encoder_output, intmd_encoder_outputs, encoder_output_4, intmd_encoder_outputs_4 = self.encode(x)

        decoder_output = self.decode(
            x1_1, x2_1, x3_1, x4_1, encoder_output, intmd_encoder_outputs, encoder_output_4, intmd_encoder_outputs_4, auxillary_output_layers, auxillary_output_layers_4
        )

        if auxillary_output_layers is not None:
            auxillary_outputs = {}
            for i in auxillary_output_layers:
                val = str(2 * i - 1)
                _key = 'Z' + str(i)
                auxillary_outputs[_key] = intmd_encoder_outputs[val]

        if auxillary_output_layers_4 is not None:
            auxillary_outputs_4 = {}
            for i in auxillary_output_layers_4:
                val_4 = str(2 * i - 1)
                _key_4 = 'Z' + str(i)
                auxillary_outputs_4[_key_4] = intmd_encoder_outputs_4[val_4]


        return decoder_output

    def _get_padding(self, padding_type, kernel_size):
        assert padding_type in ['SAME', 'VALID']
        if padding_type == 'SAME':
            _list = [(k - 1) // 2 for k in kernel_size]
            return tuple(_list)
        return tuple(0 for _ in kernel_size)

    def _reshape_output(self, x):
        x = x.view(
            x.size(0),
            int(self.img_dim / self.patch_dim),
            int(self.img_dim / self.patch_dim),
            int(self.img_dim / self.patch_dim),
            self.embedding_dim,
        )
        x = x.permute(0, 4, 1, 2, 3).contiguous()

        return x

    def _reshape_output_4(self, x):
        x = x.view(
            x.size(0),
            int(self.img_dim*2 / self.patch_dim),
            int(self.img_dim*2 / self.patch_dim),
            int(self.img_dim *2/ self.patch_dim),
            self.embedding_dim,
        )
        x = x.permute(0, 4, 1, 2, 3).contiguous()

        return x


#Define the Decoder class, which takes the encoder part as the input
class BiTransUnet(BiTransformerUnet):
    def __init__(
        self,
        img_dim,
        patch_dim,
        num_channels,
        num_classes,
        embedding_dim_4,
        embedding_dim,
        num_heads,
        num_layers,
        hidden_dim,
        dropout_rate=0.0,
        attn_dropout_rate=0.0,
        conv_patch_representation=True,
        positional_encoding_type="learned",
    ):
        super(BTS, self).__init__(
            img_dim=img_dim,
            patch_dim=patch_dim,
            num_channels=num_channels,
            embedding_dim_4=embedding_dim_4,
            embedding_dim=embedding_dim,
            num_heads=num_heads,
            num_layers=num_layers,
            hidden_dim=hidden_dim,
            dropout_rate=dropout_rate,
            attn_dropout_rate=attn_dropout_rate,
            conv_patch_representation=conv_patch_representation,
            positional_encoding_type=positional_encoding_type,
        )

        self.num_classes = num_classes

        self.Softmax = nn.Softmax(dim=1)

        self.Enblock7_1 = EnBlock1_4(in_channels=self.embedding_dim_4 // 2) # 256
        self.Enblock7_2 = EnBlock2_4(in_channels=self.embedding_dim_4 // 4) # 128
        
        self.Enblock8_1 = EnBlock1(in_channels=self.embedding_dim) #256
        self.Enblock8_2 = EnBlock2(in_channels=self.embedding_dim // 2) #128
        
        self.DeUp5 = DeUp_Cat(in_channels=self.embedding_dim //2 , out_channels=self.embedding_dim//2) #128
        self.DeBlock5 = DeBlock(in_channels=self.embedding_dim //2)#128

        self.Cat4 = Cat(in_channels=self.embedding_dim_4//4, out_channels=self.embedding_dim_4//4)
        self.Block4 = DeBlock(in_channels=self.embedding_dim_4//4)
        
        self.DeUp4 = DeUp_Cat(in_channels=self.embedding_dim_4//4, out_channels=self.embedding_dim_4//8) #64
        self.DeBlock4 = DeBlock(in_channels=self.embedding_dim_4//8)

        self.DeUp3 = DeUp_Cat(in_channels=self.embedding_dim_4//8, out_channels=self.embedding_dim_4//16) #32
        self.DeBlock3 = DeBlock(in_channels=self.embedding_dim_4//16)

        self.DeUp2 = DeUp_Cat(in_channels=self.embedding_dim_4//16, out_channels=self.embedding_dim_4//32)#16
        self.DeBlock2 = DeBlock(in_channels=self.embedding_dim_4//32)

        self.endconv = nn.Conv3d(self.embedding_dim_4 // 32, 4, kernel_size=1) #16

        #self.decoder8_1 = UBlock(self.embedding_dim, self.embedding_dim // 2,  self.embedding_dim // 4, norm_layer, dropout=dropout)
        #self.decoder8_2 = UBlock(self.embedding_dim // 4, self.embedding_dim // 4,  self.embedding_dim // 4, norm_layer, dropout=dropout)

                
    def decode(self, x1_1, x2_1, x3_1, x4_1, x, intmd_x, x4_t, intmd_x4_t, intmd_layers=[1, 2, 3, 4], intmd_layers_4=[1, 2, 3, 4]):

        assert intmd_layers is not None, "pass the intermediate layers for MLA"
        encoder_outputs = {}
        all_keys = []
        for i in intmd_layers:
            val = str(2 * i - 1)
            _key = 'Z' + str(i)
            all_keys.append(_key)
            encoder_outputs[_key] = intmd_x[val]
        all_keys.reverse()

        encoder_outputs_4 = {} #
        all_keys_4= [] #
        for i in intmd_layers_4: #
            val_4 = str(2 * i - 1) #
            _key_4 = 'Z' + str(i) #
            all_keys_4.append(_key_4) #
            encoder_outputs_4[_key_4] = intmd_x4_t[val_4] #
        all_keys_4.reverse() #

        x7 = encoder_outputs_4[all_keys_4[0]] #
        x7 = self._reshape_output_4(x7) # 2 256 16 16 16 
        x7 = self.Enblock7_1(x7) #
        x7 = self.Enblock7_2(x7) # 2 128 32 32 32 

        x8 = encoder_outputs[all_keys[0]]
        x8 = self._reshape_output(x8)
        x8 = self.Enblock8_1(x8)
        x8 = self.Enblock8_2(x8)

        y5 = self.DeUp5(x8, x4_1)  # (1, 128, 32, 32, 32)
        y5 = self.DeBlock5(y5)

        y4 = self.Cat4(x7, y5) # (1, 64, 32, 32, 32)
        y4 = self.Block4(y4)
        
        y4 = self.DeUp4(y4, x3_1)  # (1, 64, 32, 32, 32)
        y4 = self.DeBlock4(y4)

        y3 = self.DeUp3(y4, x2_1)  # (1, 32, 64, 64, 64)
        y3 = self.DeBlock3(y3)

        y2 = self.DeUp2(y3, x1_1)  # (1, 16, 128, 128, 128)
        y2 = self.DeBlock2(y2)

        y = self.endconv(y2)      # (1, 4, 128, 128, 128)
        y = self.Softmax(y)
        return y

#The following classes are specific operations for each decoder layer
class EnBlock1(nn.Module):
    def __init__(self, in_channels):
        super(EnBlock1, self).__init__()

        self.bn1 = nn.BatchNorm3d(256 // 2)
        self.relu1 = nn.ReLU(inplace=True)
        self.bn2 = nn.BatchNorm3d(256 // 2)
        self.relu2 = nn.ReLU(inplace=True)
        self.conv1 = nn.Conv3d(in_channels, in_channels // 2, kernel_size=3, padding=1)
        self.conv2 = nn.Conv3d(in_channels // 2, in_channels // 2, kernel_size=3, padding=1)

    def forward(self, x):
        x1 = self.conv1(x)
        x1 = self.bn1(x1)
        x1 = self.relu1(x1)
        x1 = self.conv2(x1)
        x1 = self.bn2(x1)
        x1 = self.relu2(x1)

        return x1


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

        self.conv1 = nn.Conv3d(in_channels, in_channels, kernel_size=3, padding=1)
        self.bn1 = nn.BatchNorm3d(256 // 2)
        self.relu1 = nn.ReLU(inplace=True)
        self.bn2 = nn.BatchNorm3d(256 // 2)
        self.relu2 = nn.ReLU(inplace=True)
        self.conv2 = nn.Conv3d(in_channels, in_channels, kernel_size=3, padding=1)

    def forward(self, x):
        x1 = self.conv1(x)
        x1 = self.bn1(x1)
        x1 = self.relu1(x1)
        x1 = self.conv2(x1)
        x1 = self.bn2(x1)
        x1 = self.relu2(x1)
        x1 = x1 + x

        return x1

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

        self.bn1 = nn.BatchNorm3d(512 // 4)
        self.relu1 = nn.ReLU(inplace=True)
        self.bn2 = nn.BatchNorm3d(512 // 4)
        self.relu2 = nn.ReLU(inplace=True)
        self.conv1 = nn.Conv3d(in_channels, in_channels // 2 , kernel_size=3, padding=1)
        self.conv2 = nn.Conv3d(in_channels // 2 , in_channels // 2 , kernel_size=3, padding=1)

    def forward(self, x):
        x1 = self.conv1(x)
        x1 = self.bn1(x1)
        x1 = self.relu1(x1)
        x1 = self.conv2(x1)
        x1 = self.bn2(x1)
        x1 = self.relu2(x1)

        return x1


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

        self.conv1 = nn.Conv3d(in_channels, in_channels, kernel_size=3, padding=1)
        self.bn1 = nn.BatchNorm3d(512 // 4)
        self.relu1 = nn.ReLU(inplace=True)
        self.bn2 = nn.BatchNorm3d(512 // 4)
        self.relu2 = nn.ReLU(inplace=True)
        self.conv2 = nn.Conv3d(in_channels, in_channels, kernel_size=3, padding=1)

    def forward(self, x):
        x1 = self.conv1(x)
        x1 = self.bn1(x1)
        x1 = self.relu1(x1)
        x1 = self.conv2(x1)
        x1 = self.bn2(x1)
        x1 = self.relu2(x1)
        x1 = x1 + x

        return x1

class DeUp_Cat(nn.Module):
    def __init__(self, in_channels, out_channels):
        super(DeUp_Cat, self).__init__()
        self.conv1 = nn.Conv3d(in_channels, out_channels, kernel_size=1)
        self.conv2 = nn.ConvTranspose3d(out_channels, out_channels, kernel_size=2, stride=2) #decrease half size
        self.conv3 = nn.Conv3d(out_channels*2, out_channels, kernel_size=1)

    def forward(self, x, prev):
        x1 = self.conv1(x)
        y = self.conv2(x1)
        # y = y + prev
        y = torch.cat((prev, y), dim=1)
        y = self.conv3(y)
        return y

class Cat(nn.Module):
    def __init__(self, in_channels, out_channels):
        super(Cat, self).__init__()
        self.conv1 = nn.Conv3d(in_channels, out_channels, kernel_size=1)
        self.conv2 = nn.ConvTranspose3d(out_channels, out_channels, kernel_size=1)
        self.conv3 = nn.Conv3d(out_channels*2, out_channels, kernel_size=1)

    def forward(self, x, prev):
        x1 = self.conv1(x)
        y = self.conv2(x1)
        # y = y + prev
        y = torch.cat((prev, y), dim=1)
        y = self.conv3(y)
        return y
    
class DeBlock(nn.Module):
    def __init__(self, in_channels):
        super(DeBlock, self).__init__()

        self.bn1 = nn.BatchNorm3d(in_channels)
        self.relu1 = nn.ReLU(inplace=True)
        self.conv1 = nn.Conv3d(in_channels, in_channels, kernel_size=3, padding=1)
        self.conv2 = nn.Conv3d(in_channels, in_channels, kernel_size=3, padding=1)
        self.bn2 = nn.BatchNorm3d(in_channels)
        self.relu2 = nn.ReLU(inplace=True)

    def forward(self, x):
        x1 = self.conv1(x)
        x1 = self.bn1(x1)
        x1 = self.relu1(x1)
        x1 = self.conv2(x1)
        x1 = self.bn2(x1)
        x1 = self.relu2(x1)
        x1 = x1 + x

        return x1



#Combine the encoder and the decoder to run the whole BiTr-Unet
def BiTrUnet(dataset='brats', _conv_repr=True, _pe_type="learned"):

    if dataset.lower() == 'brats':
        img_dim = 128
        num_classes = 4

    num_channels = 4
    patch_dim = 16
    aux_layers = [1, 2, 3, 4]
    model = BiTransUnet(
        img_dim,
        patch_dim,
        num_channels,
        num_classes,
        embedding_dim_4=512,
        embedding_dim=256,
        num_heads=8,
        num_layers=4,
        hidden_dim=4096,
        dropout_rate=0.1,
        attn_dropout_rate=0.1,
        conv_patch_representation=_conv_repr,
        positional_encoding_type=_pe_type,
    )

    return aux_layers, model


if __name__ == '__main__':
    with torch.no_grad():
        import os
        os.environ['CUDA_VISIBLE_DEVICES'] = '0'
        cuda0 = torch.device('cuda:0')
        x = torch.rand((1, 4, 128, 128, 128), device=cuda0)
        _, model = BiTrUnet(dataset='brats', _conv_repr=True, _pe_type="learned")
        model.cuda()
        y = model(x)
        print(y.shape)
