#    Copyright 2020 Division of Medical Image Computing, German Cancer Research Center (DKFZ), Heidelberg, Germany
#
#    Licensed under the Apache License, Version 2.0 (the "License");
#    you may not use this file except in compliance with the License.
#    You may obtain a copy of the License at
#
#        http://www.apache.org/licenses/LICENSE-2.0
#
#    Unless required by applicable law or agreed to in writing, software
#    distributed under the License is distributed on an "AS IS" BASIS,
#    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#    See the License for the specific language governing permissions and
#    limitations under the License.


from copy import deepcopy
from utilities.nd_softmax import softmax_helper
from torch import nn
import torch
import numpy as np
from network_architecture.initialization import InitWeights_He
from network_architecture.neural_network import SegmentationNetwork
from network_architecture.generic_UNet import Generic_UNet, ConvDropoutNormNonlin
from network_architecture.custom_modules.conv_blocks import BasicResidualBlock, ResidualLayer
from network_architecture.generic_UNet import Upsample
from network_architecture.generic_modular_UNet import PlainConvUNetDecoder, get_default_network_config

import torch.nn.functional

class LinearUpSampling(nn.Module):
    '''
    Trilinear interpolate to upsampling
    '''
    def __init__(self, inChans, outChans, scale_factor=2, mode="trilinear", align_corners=True):
        super(LinearUpSampling, self).__init__()
        self.scale_factor = scale_factor
        self.mode = mode
        self.align_corners = align_corners
        self.conv1 = nn.Conv3d(in_channels=inChans, out_channels=outChans, kernel_size=1)
        self.conv2 = nn.Conv3d(in_channels=inChans, out_channels=outChans, kernel_size=1)
    
    def forward(self, x, skipx=None):
        out = self.conv1(x)
        # out = self.up1(out)
        out = nn.functional.interpolate(out, scale_factor=self.scale_factor, mode=self.mode, align_corners=self.align_corners)

        if skipx is not None:
            out = torch.cat((out, skipx), 1)
            out = self.conv2(out)
        
        return out

class VDResampling(nn.Module):
    '''
    Variational Auto-Encoder Resampling block
    '''
    def __init__(self, inChans=256, outChans=256, in_dense_features=(3,4,2), out_dense_features=(5,14,8), stride=2, kernel_size=3, padding=1, activation="LeakyReLU", normalization="group_normalization"):
        super(VDResampling, self).__init__()
        
        self.midChans = int(inChans / 2)
        self.inChans = inChans
        self.in_dense_features = in_dense_features
        self.out_dense_features = out_dense_features
        if normalization == "group_normalization":
            self.gn1 = nn.GroupNorm(num_groups=8,num_channels=inChans)
        if activation == "relu":
            self.actv1 = nn.ReLU(inplace=True)
            self.actv2 = nn.ReLU(inplace=True)
        elif activation == "LeakyReLU":
            self.actv1 = nn.LeakyReLU(negative_slope=1e-2, inplace=True)
            self.actv2 = nn.LeakyReLU(negative_slope=1e-2, inplace=True)
        self.actv_vd = nn.Sigmoid()
        self.conv1 = nn.Conv3d(in_channels=inChans, out_channels=16, kernel_size=kernel_size, stride=stride, padding=padding)
        self.dense1 = nn.Linear(in_features=16*in_dense_features[0]*in_dense_features[1]*in_dense_features[2], out_features=inChans)
        self.dense2 = nn.Linear(in_features=self.midChans, out_features=inChans*out_dense_features[0]*out_dense_features[1]*out_dense_features[2])
        
    def forward(self, x):
        out = self.gn1(x)
        out = self.actv1(out)
        # print("ACTIVATE SIZE:", out.size())
        out = self.conv1(out)
        # print("COV SIZE:", out.size())
        out = out.view(-1, self.num_flat_features(out))
        # print("FLAT SIZE:", out.size())
        out_vd = self.dense1(out)
        # print("DENSE1 SIZE:", out_vd.size())
        out_vd = self.actv_vd(out_vd)
        # ACTIVATE SIZE: torch.Size([2, 320, 5, 7, 4])
        # COV SIZE: torch.Size([2, 16, 3, 4, 2])
        # FLAT SIZE: torch.Size([2, 384])
        # DENSE1 SIZE: torch.Size([2, 320])
        distr = out_vd 
        out = VDraw(out_vd)
        # print("VDraw SIZE:", out.size())

    #     VDraw SIZE: torch.Size([2, 160])
    # DENSE2 SIZE: torch.Size([2, 3840])  match [2, 5,14,8]
    # ACTI2 SIZE: torch.Size([2, 3840])
    # VIEW SIZE: torch.Size([2, 160, 3, 4, 2])
    # UPSAMPL SIZE: torch.Size([2, 320, 6, 8, 4])
        out = self.dense2(out)
        # print("DENSE2 SIZE:", out.size())
        out = self.actv2(out)
        # print("ACTI2 SIZE:", out.size())
        out = out.view((-1, self.inChans, self.out_dense_features[0],self.out_dense_features[1],self.out_dense_features[2]))
        # print("VIEW SIZE:", out.size())
        
        return out, distr
        
    def num_flat_features(self, x):
        size = x.size()[1:]
        num_features = 1
        for s in size:
            num_features *= s
            
        return num_features

def VDraw(x):
    # Generate a Gaussian distribution with the given mean(128-d) and std(128-d)
    d = x.size(1) // 2
    return torch.distributions.Normal(x[:,:d], x[:,d:]).sample()

class VDecoderBlock(nn.Module):
    '''
    Variational Decoder block
    '''
    def __init__(self, inChans, outChans, num_groups=8, activation="LeakyReLU", normalization="group_normalization", mode="trilinear"):
        super(VDecoderBlock, self).__init__()

        self.up0 = LinearUpSampling(inChans, outChans, mode=mode)
        self.block = DecoderBlock(outChans, outChans, num_groups=num_groups, activation=activation, normalization=normalization)
    
    def forward(self, x):
        out = self.up0(x)
        out = self.block(out)

        return out

class DecoderBlock(nn.Module):
    '''
    Decoder block
    '''
    def __init__(self, inChans, outChans, stride=1, padding=1, num_groups=8, activation="LeakyReLU", normalization="group_normalization"):
        super(DecoderBlock, self).__init__()
        
        if normalization == "group_normalization":
            self.norm1 = nn.GroupNorm(num_groups=num_groups, num_channels=outChans)
            self.norm2 = nn.GroupNorm(num_groups=num_groups, num_channels=outChans)
        else:
            self.norm1 = None
            self.norm2 = None
        if activation == "relu":
            self.actv1 = nn.ReLU(inplace=True)
            self.actv2 = nn.ReLU(inplace=True)
        elif activation == "LeakyReLU":
            self.actv1 = nn.LeakyReLU(negative_slope=1e-2, inplace=True)
            self.actv2 = nn.LeakyReLU(negative_slope=1e-2, inplace=True)           
        self.conv1 = nn.Conv3d(in_channels=inChans, out_channels=outChans, kernel_size=3, stride=stride, padding=padding)
        self.conv2 = nn.Conv3d(in_channels=outChans, out_channels=outChans, kernel_size=3, stride=stride, padding=padding)
        
        
    def forward(self, x):
        residual = x
        
        out = x
        
        if self.norm1:
            out = self.norm1(out)
        out = self.actv1(out)
        out = self.conv1(out)
        if self.norm2:
            out = self.norm2(out)
        out = self.actv2(out)
        out = self.conv2(out)
        
        out += residual
        
        return out

class VAE_ResidualUNetEncoder(nn.Module):
    def __init__(self, input_channels, base_num_features, num_blocks_per_stage, feat_map_mul_on_downscale,
                 pool_op_kernel_sizes, conv_kernel_sizes, props, default_return_skips=True,
                 max_num_features=480, block=BasicResidualBlock):
        """
        Following UNet building blocks can be added by utilizing the properties this class exposes (TODO)

        this one includes the bottleneck layer!

        :param input_channels:
        :param base_num_features:
        :param num_blocks_per_stage:
        :param feat_map_mul_on_downscale:
        :param pool_op_kernel_sizes:
        :param conv_kernel_sizes:
        :param props:
        """
        super(VAE_ResidualUNetEncoder, self).__init__()

        self.default_return_skips = default_return_skips
        self.props = props

        self.stages = []
        self.stage_output_features = []
        self.stage_pool_kernel_size = []
        self.stage_conv_op_kernel_size = []

        assert len(pool_op_kernel_sizes) == len(conv_kernel_sizes)

        num_stages = len(conv_kernel_sizes)

        if not isinstance(num_blocks_per_stage, (list, tuple)):
            num_blocks_per_stage = [num_blocks_per_stage] * num_stages
        else:
            assert len(num_blocks_per_stage) == num_stages

        self.num_blocks_per_stage = num_blocks_per_stage  # decoder may need this

        

        # self.initial_conv = props['conv_op'](input_channels, base_num_features, 3, padding=1, **props['conv_op_kwargs'])
        # self.initial_norm = props['norm_op'](base_num_features, **props['norm_op_kwargs'])
        # self.initial_nonlin = props['nonlin'](**props['nonlin_kwargs'])

        self.conv1 = props['conv_op'](input_channels, base_num_features, 3, padding=1, **props['conv_op_kwargs'])
        self.norm1 = props['norm_op'](base_num_features, **props['norm_op_kwargs'])
        self.nonlin = props['nonlin'](**props['nonlin_kwargs'])
        

        current_input_features = base_num_features

        
        for stage in range(num_stages):
            current_output_features = min(base_num_features * feat_map_mul_on_downscale ** stage, max_num_features)
            current_kernel_size = conv_kernel_sizes[stage]
            current_pool_kernel_size = pool_op_kernel_sizes[stage]

            current_stage = ResidualLayer(current_input_features, current_output_features, current_kernel_size, props,
                                          self.num_blocks_per_stage[stage], current_pool_kernel_size, block)

            self.stages.append(current_stage)
            self.stage_output_features.append(current_output_features)
            self.stage_conv_op_kernel_size.append(current_kernel_size)
            self.stage_pool_kernel_size.append(current_pool_kernel_size)

            # update current_input_features
            current_input_features = current_output_features
        
        # cum_upsample = np.cumprod(np.vstack(self.stage_pool_kernel_size), axis=0).astype(int)
        

        self.stages = nn.ModuleList(self.stages)

    def forward(self, x, return_skips=None):
        """

        :param x:
        :param return_skips: if none then self.default_return_skips is used
        :return:
        """
        skips = []

        # x = self.initial_nonlin(self.initial_norm(self.initial_conv(x)))
        x = self.nonlin(self.norm1(self.conv1(x)))
    


        for s in self.stages:
            x = s(x)
            if self.default_return_skips:
                skips.append(x)

        if return_skips is None:
            return_skips = self.default_return_skips

        if return_skips:
            return skips
        else:
            return x

    @staticmethod
    def compute_approx_vram_consumption(patch_size, base_num_features, max_num_features,
                                        num_modalities, pool_op_kernel_sizes, num_conv_per_stage_encoder,
                                        feat_map_mul_on_downscale, batch_size):
        npool = len(pool_op_kernel_sizes) - 1

        current_shape = np.array(patch_size)

        tmp = (num_conv_per_stage_encoder[0] * 2 + 1) * np.prod(current_shape) * base_num_features \
              + num_modalities * np.prod(current_shape)

        num_feat = base_num_features

        for p in range(1, npool + 1):
            current_shape = current_shape / np.array(pool_op_kernel_sizes[p])
            num_feat = min(num_feat * feat_map_mul_on_downscale, max_num_features)
            num_convs = num_conv_per_stage_encoder[p] * 2 + 1  # + 1 for conv in skip in first block
            print(p, num_feat, num_convs, current_shape)
            tmp += num_convs * np.prod(current_shape) * num_feat
        return tmp * batch_size



class ResidualUNetDecoder(nn.Module):
    def __init__(self, previous, num_classes, num_blocks_per_stage=None, network_props=None, deep_supervision=False,
                 upscale_logits=False, block=BasicResidualBlock):
        super(ResidualUNetDecoder, self).__init__()
        self.num_classes = num_classes
        self.deep_supervision = deep_supervision
        """
        We assume the bottleneck is part of the encoder, so we can start with upsample -> concat here
        """
        previous_stages = previous.stages
        previous_stage_output_features = previous.stage_output_features
        previous_stage_pool_kernel_size = previous.stage_pool_kernel_size
        previous_stage_conv_op_kernel_size = previous.stage_conv_op_kernel_size

        if network_props is None:
            self.props = previous.props
        else:
            self.props = network_props

        if self.props['conv_op'] == nn.Conv2d:
            transpconv = nn.ConvTranspose2d
            upsample_mode = "bilinear"
        elif self.props['conv_op'] == nn.Conv3d:
            transpconv = nn.ConvTranspose3d
            upsample_mode = "trilinear"
        else:
            raise ValueError("unknown convolution dimensionality, conv op: %s" % str(self.props['conv_op']))

        if num_blocks_per_stage is None:
            num_blocks_per_stage = previous.num_blocks_per_stage[:-1][::-1]

        assert len(num_blocks_per_stage) == len(previous.num_blocks_per_stage) - 1

        self.stage_pool_kernel_size = previous_stage_pool_kernel_size
        self.stage_output_features = previous_stage_output_features
        self.stage_conv_op_kernel_size = previous_stage_conv_op_kernel_size

        num_stages = len(previous_stages) - 1  # we have one less as the first stage here is what comes after the
        # bottleneck

        self.tus = []
        self.stages = []
        self.deep_supervision_outputs = []

        # only used for upsample_logits
        cum_upsample = np.cumprod(np.vstack(self.stage_pool_kernel_size), axis=0).astype(int)

        for i, s in enumerate(np.arange(num_stages)[::-1]):
            features_below = previous_stage_output_features[s + 1]
            features_skip = previous_stage_output_features[s]

            self.tus.append(transpconv(features_below, features_skip, previous_stage_pool_kernel_size[s + 1],
                                       previous_stage_pool_kernel_size[s + 1], bias=False))
            # after we tu we concat features so now we have 2xfeatures_skip
            self.stages.append(ResidualLayer(2 * features_skip, features_skip, previous_stage_conv_op_kernel_size[s],
                                             self.props, num_blocks_per_stage[i], None, block))

            if deep_supervision and s != 0:
                seg_layer = self.props['conv_op'](features_skip, num_classes, 1, 1, 0, 1, 1, False)
                if upscale_logits:
                    upsample = Upsample(scale_factor=cum_upsample[s], mode=upsample_mode)
                    self.deep_supervision_outputs.append(nn.Sequential(seg_layer, upsample))
                else:
                    self.deep_supervision_outputs.append(seg_layer)

        self.segmentation_output = self.props['conv_op'](features_skip, num_classes, 1, 1, 0, 1, 1, False)

        self.tus = nn.ModuleList(self.tus)
        self.stages = nn.ModuleList(self.stages)
        self.deep_supervision_outputs = nn.ModuleList(self.deep_supervision_outputs)

    def forward(self, skips):
        # skips come from the encoder. They are sorted so that the bottleneck is last in the list
        # what is maybe not perfect is that the TUs and stages here are sorted the other way around
        # so let's just reverse the order of skips
        skips = skips[::-1]
        seg_outputs = []

        x = skips[0]  # this is the bottleneck

        for i in range(len(self.tus)):
            x = self.tus[i](x)
            x = torch.cat((x, skips[i + 1]), dim=1)
            x = self.stages[i](x)
            if self.deep_supervision and (i != len(self.tus) - 1):
                seg_outputs.append(self.deep_supervision_outputs[i](x))

        segmentation = self.segmentation_output(x)

        if self.deep_supervision:
            seg_outputs.append(segmentation)
            return seg_outputs[
                   ::-1]  # seg_outputs are ordered so that the seg from the highest layer is first, the seg from
            # the bottleneck of the UNet last
        else:
            return segmentation

    @staticmethod
    def compute_approx_vram_consumption(patch_size, base_num_features, max_num_features,
                                        num_classes, pool_op_kernel_sizes, num_blocks_per_stage_decoder,
                                        feat_map_mul_on_downscale, batch_size):
        """
        This only applies for num_conv_per_stage and convolutional_upsampling=True
        not real vram consumption. just a constant term to which the vram consumption will be approx proportional
        (+ offset for parameter storage)
        :param patch_size:
        :param num_pool_per_axis:
        :param base_num_features:
        :param max_num_features:
        :return:
        """
        npool = len(pool_op_kernel_sizes) - 1

        current_shape = np.array(patch_size)
        tmp = (num_blocks_per_stage_decoder[-1] * 2 + 1) * np.prod(
            current_shape) * base_num_features + num_classes * np.prod(current_shape)

        num_feat = base_num_features

        for p in range(1, npool):
            current_shape = current_shape / np.array(pool_op_kernel_sizes[p])
            num_feat = min(num_feat * feat_map_mul_on_downscale, max_num_features)
            num_convs = num_blocks_per_stage_decoder[-(p + 1)] * 2 + 1 + 1  # +1 for transpconv and +1 for conv in skip
            print(p, num_feat, num_convs, current_shape)
            tmp += num_convs * np.prod(current_shape) * num_feat

        return tmp * batch_size


class VAE(nn.Module):
    '''
    Variational Auto-Encoder : to group the features extracted by Encoder
    '''
    def __init__(self, inChans=256, outChans=4,activation="relu", normalization="group_normalization", mode="trilinear"):
        super(VAE, self).__init__()

        self.vd_resample = VDResampling(inChans=inChans, outChans=inChans)
        self.vd_block3 = VDecoderBlock(inChans, inChans//2, num_groups=10)
        self.vd_block2 = VDecoderBlock(inChans//2, inChans//4, num_groups=10)
        self.vd_block1 = VDecoderBlock(inChans//4, inChans//8, num_groups=10)
        self.vd_block0 = VDecoderBlock(inChans//8, inChans//16, num_groups=10)
        self.vd_end = nn.Conv3d(inChans//16, outChans, kernel_size=1)
        
    def forward(self, x):

        # out: (B, 320, 5, 14, 8)
        out, distr = self.vd_resample(x)
        
        # out: (B, 160, 10, 28, 16)
        out = self.vd_block3(out)
        # print('OUT2 size',out.size())
        # out: (B, 80, 20, 56, 32)
        out = self.vd_block2(out)
        # print('OUT3 size',out.size())
        # out: (B, 40, 40, 112, 64)
        out = self.vd_block1(out)
        # print('OUT4 size',out.size())
        # out: (B, 20, 80, 224, 128)
        out = self.vd_block0(out)
        # print('OUT5 size',out.size())
        # out: (B, 1, 80, 224, 128)
        out = self.vd_end(out)

        return out, distr

class Generic_UNet_VAE(Generic_UNet):

    def __init__(self, input_channels, base_num_features, num_classes, num_pool, num_conv_per_stage=2,
                 feat_map_mul_on_downscale=2, conv_op=nn.Conv2d,
                 norm_op=nn.BatchNorm2d, norm_op_kwargs=None,
                 dropout_op=nn.Dropout2d, dropout_op_kwargs=None,
                 nonlin=nn.LeakyReLU, nonlin_kwargs=None, deep_supervision=True, dropout_in_localization=False,
                 final_nonlin=softmax_helper, weightInitializer=InitWeights_He(1e-2), pool_op_kernel_sizes=None,
                 conv_kernel_sizes=None,
                 upscale_logits=False, convolutional_pooling=False, convolutional_upsampling=False,
                 max_num_features=None, basic_block=ConvDropoutNormNonlin,
                 seg_output_use_bias=False):
        """
        basically more flexible than v1, architecture is the same

        Does this look complicated? Nah bro. Functionality > usability

        This does everything you need, including world peace.

        Questions? -> <EMAIL>
        """
        super(Generic_UNet_VAE, self).__init__(input_channels, base_num_features, num_classes, num_pool, num_conv_per_stage,
                 feat_map_mul_on_downscale, conv_op,
                 norm_op, norm_op_kwargs,
                 dropout_op, dropout_op_kwargs,
                 nonlin, nonlin_kwargs, deep_supervision, dropout_in_localization,
                 final_nonlin, weightInitializer, pool_op_kernel_sizes,
                 conv_kernel_sizes,
                 upscale_logits, convolutional_pooling, convolutional_upsampling,
                 max_num_features, basic_block,
                 seg_output_use_bias)

        # Variational Auto-Encoder
        self.vae = VAE(320, outChans=input_channels)
    
    def forward(self, x):

        # DEBUG: input size torch.Size([1, 1, 80, 224, 128])
        # DEBUG: encoding size torch.Size([1, 320, 5, 7, 4])

        # DEBUG: input size torch.Size([2, 1, 80, 224, 128])
        # DEBUG: encoding size torch.Size([2, 320, 5, 7, 4])

        skips = []
        seg_outputs = []
        for d in range(len(self.conv_blocks_context) - 1):
            x = self.conv_blocks_context[d](x)
            # print('encoder', x.size())
            skips.append(x)
            if not self.convolutional_pooling:
                x = self.td[d](x)

        x = self.conv_blocks_context[-1](x)

        # VAE
        vae_outputs, vae_distr = self.vae(x)

        for u in range(len(self.tu)):
            x = self.tu[u](x)
            x = torch.cat((x, skips[-(u + 1)]), dim=1)
            x = self.conv_blocks_localization[u](x)
            seg_outputs.append(self.final_nonlin(self.seg_outputs[u](x)))
        
        if self._deep_supervision and self.do_ds:
            ds_outputs = tuple([seg_outputs[-1]] + [i(j) for i, j in
                                              zip(list(self.upscale_logits_ops)[::-1], seg_outputs[:-1][::-1])])
            return tuple([(ds_out, vae_outputs, vae_distr) for ds_out in ds_outputs])
        else:
            return seg_outputs[-1], vae_outputs, vae_distr

#    Copyright 2020 Division of Medical Image Computing, German Cancer Research Center (DKFZ), Heidelberg, Germany
#
#    Licensed under the Apache License, Version 2.0 (the "License");
#    you may not use this file except in compliance with the License.
#    You may obtain a copy of the License at
#
#        http://www.apache.org/licenses/LICENSE-2.0
#
#    Unless required by applicable law or agreed to in writing, software
#    distributed under the License is distributed on an "AS IS" BASIS,
#    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#    See the License for the specific language governing permissions and
#    limitations under the License.


from copy import deepcopy
from utilities.nd_softmax import softmax_helper
from torch import nn
import torch
import numpy as np
from network_architecture.initialization import InitWeights_He
from network_architecture.neural_network import SegmentationNetwork
from network_architecture.generic_UNet import Generic_UNet, ConvDropoutNormNonlin
from network_architecture.custom_modules.conv_blocks import BasicResidualBlock, ResidualLayer
from network_architecture.generic_UNet import Upsample
from network_architecture.generic_modular_UNet import PlainConvUNetDecoder, get_default_network_config

import torch.nn.functional

class LinearUpSampling(nn.Module):
    '''
    Trilinear interpolate to upsampling
    '''
    def __init__(self, inChans, outChans, scale_factor=2, mode="trilinear", align_corners=True):
        super(LinearUpSampling, self).__init__()
        self.scale_factor = scale_factor
        self.mode = mode
        self.align_corners = align_corners
        self.conv1 = nn.Conv3d(in_channels=inChans, out_channels=outChans, kernel_size=1)
        self.conv2 = nn.Conv3d(in_channels=inChans, out_channels=outChans, kernel_size=1)
    
    def forward(self, x, skipx=None):
        out = self.conv1(x)
        # out = self.up1(out)
        out = nn.functional.interpolate(out, scale_factor=self.scale_factor, mode=self.mode, align_corners=self.align_corners)

        if skipx is not None:
            out = torch.cat((out, skipx), 1)
            out = self.conv2(out)
        
        return out

class VDResampling(nn.Module):
    '''
    Variational Auto-Encoder Resampling block
    '''
    def __init__(self, inChans=256, outChans=256, in_dense_features=(3,4,2), out_dense_features=(5,14,8), stride=2, kernel_size=3, padding=1, activation="LeakyReLU", normalization="group_normalization"):
        super(VDResampling, self).__init__()
        
        self.midChans = int(inChans / 2)
        self.inChans = inChans
        self.in_dense_features = in_dense_features
        self.out_dense_features = out_dense_features
        if normalization == "group_normalization":
            self.gn1 = nn.GroupNorm(num_groups=8,num_channels=inChans)
        if activation == "relu":
            self.actv1 = nn.ReLU(inplace=True)
            self.actv2 = nn.ReLU(inplace=True)
        elif activation == "LeakyReLU":
            self.actv1 = nn.LeakyReLU(negative_slope=1e-2, inplace=True)
            self.actv2 = nn.LeakyReLU(negative_slope=1e-2, inplace=True)
        self.actv_vd = nn.Sigmoid()
        self.conv1 = nn.Conv3d(in_channels=inChans, out_channels=16, kernel_size=kernel_size, stride=stride, padding=padding)
        self.dense1 = nn.Linear(in_features=16*in_dense_features[0]*in_dense_features[1]*in_dense_features[2], out_features=inChans)
        self.dense2 = nn.Linear(in_features=self.midChans, out_features=inChans*out_dense_features[0]*out_dense_features[1]*out_dense_features[2])
        
    def forward(self, x):
        out = self.gn1(x)
        out = self.actv1(out)
        # print("ACTIVATE SIZE:", out.size())
        out = self.conv1(out)
        # print("COV SIZE:", out.size())
        out = out.view(-1, self.num_flat_features(out))
        # print("FLAT SIZE:", out.size())
        out_vd = self.dense1(out)
        # print("DENSE1 SIZE:", out_vd.size())
        out_vd = self.actv_vd(out_vd)
        # ACTIVATE SIZE: torch.Size([2, 320, 5, 7, 4])
        # COV SIZE: torch.Size([2, 16, 3, 4, 2])
        # FLAT SIZE: torch.Size([2, 384])
        # DENSE1 SIZE: torch.Size([2, 320])
        distr = out_vd 
        out = VDraw(out_vd)
        # print("VDraw SIZE:", out.size())

    #     VDraw SIZE: torch.Size([2, 160])
    # DENSE2 SIZE: torch.Size([2, 3840])  match [2, 5,14,8]
    # ACTI2 SIZE: torch.Size([2, 3840])
    # VIEW SIZE: torch.Size([2, 160, 3, 4, 2])
    # UPSAMPL SIZE: torch.Size([2, 320, 6, 8, 4])
        out = self.dense2(out)
        # print("DENSE2 SIZE:", out.size())
        out = self.actv2(out)
        # print("ACTI2 SIZE:", out.size())
        out = out.view((-1, self.inChans, self.out_dense_features[0],self.out_dense_features[1],self.out_dense_features[2]))
        # print("VIEW SIZE:", out.size())
        
        return out, distr
        
    def num_flat_features(self, x):
        size = x.size()[1:]
        num_features = 1
        for s in size:
            num_features *= s
            
        return num_features

def VDraw(x):
    # Generate a Gaussian distribution with the given mean(128-d) and std(128-d)
    d = x.size(1) // 2
    return torch.distributions.Normal(x[:,:d], x[:,d:]).sample()

class VDecoderBlock(nn.Module):
    '''
    Variational Decoder block
    '''
    def __init__(self, inChans, outChans, num_groups=8, activation="LeakyReLU", normalization="group_normalization", mode="trilinear"):
        super(VDecoderBlock, self).__init__()

        self.up0 = LinearUpSampling(inChans, outChans, mode=mode)
        self.block = DecoderBlock(outChans, outChans, num_groups=num_groups, activation=activation, normalization=normalization)
    
    def forward(self, x):
        out = self.up0(x)
        out = self.block(out)

        return out

class DecoderBlock(nn.Module):
    '''
    Decoder block
    '''
    def __init__(self, inChans, outChans, stride=1, padding=1, num_groups=8, activation="LeakyReLU", normalization="group_normalization"):
        super(DecoderBlock, self).__init__()
        
        if normalization == "group_normalization":
            self.norm1 = nn.GroupNorm(num_groups=num_groups, num_channels=outChans)
            self.norm2 = nn.GroupNorm(num_groups=num_groups, num_channels=outChans)
        else:
            self.norm1 = None
            self.norm2 = None
        if activation == "relu":
            self.actv1 = nn.ReLU(inplace=True)
            self.actv2 = nn.ReLU(inplace=True)
        elif activation == "LeakyReLU":
            self.actv1 = nn.LeakyReLU(negative_slope=1e-2, inplace=True)
            self.actv2 = nn.LeakyReLU(negative_slope=1e-2, inplace=True)           
        self.conv1 = nn.Conv3d(in_channels=inChans, out_channels=outChans, kernel_size=3, stride=stride, padding=padding)
        self.conv2 = nn.Conv3d(in_channels=outChans, out_channels=outChans, kernel_size=3, stride=stride, padding=padding)
        
        
    def forward(self, x):
        residual = x
        
        out = x
        
        if self.norm1:
            out = self.norm1(out)
        out = self.actv1(out)
        out = self.conv1(out)
        if self.norm2:
            out = self.norm2(out)
        out = self.actv2(out)
        out = self.conv2(out)
        
        out += residual
        
        return out

class VAE_ResidualUNetEncoder(nn.Module):
    def __init__(self, input_channels, base_num_features, num_blocks_per_stage, feat_map_mul_on_downscale,
                 pool_op_kernel_sizes, conv_kernel_sizes, props, default_return_skips=True,
                 max_num_features=480, block=BasicResidualBlock):
        """
        Following UNet building blocks can be added by utilizing the properties this class exposes (TODO)

        this one includes the bottleneck layer!

        :param input_channels:
        :param base_num_features:
        :param num_blocks_per_stage:
        :param feat_map_mul_on_downscale:
        :param pool_op_kernel_sizes:
        :param conv_kernel_sizes:
        :param props:
        """
        super(VAE_ResidualUNetEncoder, self).__init__()

        self.default_return_skips = default_return_skips
        self.props = props

        self.stages = []
        self.stage_output_features = []
        self.stage_pool_kernel_size = []
        self.stage_conv_op_kernel_size = []

        assert len(pool_op_kernel_sizes) == len(conv_kernel_sizes)

        num_stages = len(conv_kernel_sizes)

        if not isinstance(num_blocks_per_stage, (list, tuple)):
            num_blocks_per_stage = [num_blocks_per_stage] * num_stages
        else:
            assert len(num_blocks_per_stage) == num_stages

        self.num_blocks_per_stage = num_blocks_per_stage  # decoder may need this

        

        # self.initial_conv = props['conv_op'](input_channels, base_num_features, 3, padding=1, **props['conv_op_kwargs'])
        # self.initial_norm = props['norm_op'](base_num_features, **props['norm_op_kwargs'])
        # self.initial_nonlin = props['nonlin'](**props['nonlin_kwargs'])

        self.conv1 = props['conv_op'](input_channels, base_num_features, 3, padding=1, **props['conv_op_kwargs'])
        self.norm1 = props['norm_op'](base_num_features, **props['norm_op_kwargs'])
        self.nonlin = props['nonlin'](**props['nonlin_kwargs'])
        

        current_input_features = base_num_features

        
        for stage in range(num_stages):
            current_output_features = min(base_num_features * feat_map_mul_on_downscale ** stage, max_num_features)
            current_kernel_size = conv_kernel_sizes[stage]
            current_pool_kernel_size = pool_op_kernel_sizes[stage]

            current_stage = ResidualLayer(current_input_features, current_output_features, current_kernel_size, props,
                                          self.num_blocks_per_stage[stage], current_pool_kernel_size, block)

            self.stages.append(current_stage)
            self.stage_output_features.append(current_output_features)
            self.stage_conv_op_kernel_size.append(current_kernel_size)
            self.stage_pool_kernel_size.append(current_pool_kernel_size)

            # update current_input_features
            current_input_features = current_output_features
        
        # cum_upsample = np.cumprod(np.vstack(self.stage_pool_kernel_size), axis=0).astype(int)
        

        self.stages = nn.ModuleList(self.stages)

    def forward(self, x, return_skips=None):
        """

        :param x:
        :param return_skips: if none then self.default_return_skips is used
        :return:
        """
        skips = []

        # x = self.initial_nonlin(self.initial_norm(self.initial_conv(x)))
        x = self.nonlin(self.norm1(self.conv1(x)))
    


        for s in self.stages:
            x = s(x)
            if self.default_return_skips:
                skips.append(x)

        if return_skips is None:
            return_skips = self.default_return_skips

        if return_skips:
            return skips
        else:
            return x

    @staticmethod
    def compute_approx_vram_consumption(patch_size, base_num_features, max_num_features,
                                        num_modalities, pool_op_kernel_sizes, num_conv_per_stage_encoder,
                                        feat_map_mul_on_downscale, batch_size):
        npool = len(pool_op_kernel_sizes) - 1

        current_shape = np.array(patch_size)

        tmp = (num_conv_per_stage_encoder[0] * 2 + 1) * np.prod(current_shape) * base_num_features \
              + num_modalities * np.prod(current_shape)

        num_feat = base_num_features

        for p in range(1, npool + 1):
            current_shape = current_shape / np.array(pool_op_kernel_sizes[p])
            num_feat = min(num_feat * feat_map_mul_on_downscale, max_num_features)
            num_convs = num_conv_per_stage_encoder[p] * 2 + 1  # + 1 for conv in skip in first block
            print(p, num_feat, num_convs, current_shape)
            tmp += num_convs * np.prod(current_shape) * num_feat
        return tmp * batch_size



class ResidualUNetDecoder(nn.Module):
    def __init__(self, previous, num_classes, num_blocks_per_stage=None, network_props=None, deep_supervision=False,
                 upscale_logits=False, block=BasicResidualBlock):
        super(ResidualUNetDecoder, self).__init__()
        self.num_classes = num_classes
        self.deep_supervision = deep_supervision
        """
        We assume the bottleneck is part of the encoder, so we can start with upsample -> concat here
        """
        previous_stages = previous.stages
        previous_stage_output_features = previous.stage_output_features
        previous_stage_pool_kernel_size = previous.stage_pool_kernel_size
        previous_stage_conv_op_kernel_size = previous.stage_conv_op_kernel_size

        if network_props is None:
            self.props = previous.props
        else:
            self.props = network_props

        if self.props['conv_op'] == nn.Conv2d:
            transpconv = nn.ConvTranspose2d
            upsample_mode = "bilinear"
        elif self.props['conv_op'] == nn.Conv3d:
            transpconv = nn.ConvTranspose3d
            upsample_mode = "trilinear"
        else:
            raise ValueError("unknown convolution dimensionality, conv op: %s" % str(self.props['conv_op']))

        if num_blocks_per_stage is None:
            num_blocks_per_stage = previous.num_blocks_per_stage[:-1][::-1]

        assert len(num_blocks_per_stage) == len(previous.num_blocks_per_stage) - 1

        self.stage_pool_kernel_size = previous_stage_pool_kernel_size
        self.stage_output_features = previous_stage_output_features
        self.stage_conv_op_kernel_size = previous_stage_conv_op_kernel_size

        num_stages = len(previous_stages) - 1  # we have one less as the first stage here is what comes after the
        # bottleneck

        self.tus = []
        self.stages = []
        self.deep_supervision_outputs = []

        # only used for upsample_logits
        cum_upsample = np.cumprod(np.vstack(self.stage_pool_kernel_size), axis=0).astype(int)

        for i, s in enumerate(np.arange(num_stages)[::-1]):
            features_below = previous_stage_output_features[s + 1]
            features_skip = previous_stage_output_features[s]

            self.tus.append(transpconv(features_below, features_skip, previous_stage_pool_kernel_size[s + 1],
                                       previous_stage_pool_kernel_size[s + 1], bias=False))
            # after we tu we concat features so now we have 2xfeatures_skip
            self.stages.append(ResidualLayer(2 * features_skip, features_skip, previous_stage_conv_op_kernel_size[s],
                                             self.props, num_blocks_per_stage[i], None, block))

            if deep_supervision and s != 0:
                seg_layer = self.props['conv_op'](features_skip, num_classes, 1, 1, 0, 1, 1, False)
                if upscale_logits:
                    upsample = Upsample(scale_factor=cum_upsample[s], mode=upsample_mode)
                    self.deep_supervision_outputs.append(nn.Sequential(seg_layer, upsample))
                else:
                    self.deep_supervision_outputs.append(seg_layer)

        self.segmentation_output = self.props['conv_op'](features_skip, num_classes, 1, 1, 0, 1, 1, False)

        self.tus = nn.ModuleList(self.tus)
        self.stages = nn.ModuleList(self.stages)
        self.deep_supervision_outputs = nn.ModuleList(self.deep_supervision_outputs)

    def forward(self, skips):
        # skips come from the encoder. They are sorted so that the bottleneck is last in the list
        # what is maybe not perfect is that the TUs and stages here are sorted the other way around
        # so let's just reverse the order of skips
        skips = skips[::-1]
        seg_outputs = []

        x = skips[0]  # this is the bottleneck

        for i in range(len(self.tus)):
            x = self.tus[i](x)
            x = torch.cat((x, skips[i + 1]), dim=1)
            x = self.stages[i](x)
            if self.deep_supervision and (i != len(self.tus) - 1):
                seg_outputs.append(self.deep_supervision_outputs[i](x))

        segmentation = self.segmentation_output(x)

        if self.deep_supervision:
            seg_outputs.append(segmentation)
            return seg_outputs[
                   ::-1]  # seg_outputs are ordered so that the seg from the highest layer is first, the seg from
            # the bottleneck of the UNet last
        else:
            return segmentation

    @staticmethod
    def compute_approx_vram_consumption(patch_size, base_num_features, max_num_features,
                                        num_classes, pool_op_kernel_sizes, num_blocks_per_stage_decoder,
                                        feat_map_mul_on_downscale, batch_size):
        """
        This only applies for num_conv_per_stage and convolutional_upsampling=True
        not real vram consumption. just a constant term to which the vram consumption will be approx proportional
        (+ offset for parameter storage)
        :param patch_size:
        :param num_pool_per_axis:
        :param base_num_features:
        :param max_num_features:
        :return:
        """
        npool = len(pool_op_kernel_sizes) - 1

        current_shape = np.array(patch_size)
        tmp = (num_blocks_per_stage_decoder[-1] * 2 + 1) * np.prod(
            current_shape) * base_num_features + num_classes * np.prod(current_shape)

        num_feat = base_num_features

        for p in range(1, npool):
            current_shape = current_shape / np.array(pool_op_kernel_sizes[p])
            num_feat = min(num_feat * feat_map_mul_on_downscale, max_num_features)
            num_convs = num_blocks_per_stage_decoder[-(p + 1)] * 2 + 1 + 1  # +1 for transpconv and +1 for conv in skip
            print(p, num_feat, num_convs, current_shape)
            tmp += num_convs * np.prod(current_shape) * num_feat

        return tmp * batch_size


class VAE(nn.Module):
    '''
    Variational Auto-Encoder : to group the features extracted by Encoder
    '''
    def __init__(self, inChans=256, outChans=4,activation="relu", normalization="group_normalization", mode="trilinear"):
        super(VAE, self).__init__()

        self.vd_resample = VDResampling(inChans=inChans, outChans=inChans)
        self.vd_block3 = VDecoderBlock(inChans, inChans//2, num_groups=10)
        self.vd_block2 = VDecoderBlock(inChans//2, inChans//4, num_groups=10)
        self.vd_block1 = VDecoderBlock(inChans//4, inChans//8, num_groups=10)
        self.vd_block0 = VDecoderBlock(inChans//8, inChans//16, num_groups=10)
        self.vd_end = nn.Conv3d(inChans//16, outChans, kernel_size=1)
        
    def forward(self, x):

        # out: (B, 320, 5, 14, 8)
        out, distr = self.vd_resample(x)
        
        # out: (B, 160, 10, 28, 16)
        out = self.vd_block3(out)
        # print('OUT2 size',out.size())
        # out: (B, 80, 20, 56, 32)
        out = self.vd_block2(out)
        # print('OUT3 size',out.size())
        # out: (B, 40, 40, 112, 64)
        out = self.vd_block1(out)
        # print('OUT4 size',out.size())
        # out: (B, 20, 80, 224, 128)
        out = self.vd_block0(out)
        # print('OUT5 size',out.size())
        # out: (B, 1, 80, 224, 128)
        out = self.vd_end(out)

        return out, distr

class Generic_UNet_VAE(Generic_UNet):

    def __init__(self, input_channels, base_num_features, num_classes, num_pool, num_conv_per_stage=2,
                 feat_map_mul_on_downscale=2, conv_op=nn.Conv2d,
                 norm_op=nn.BatchNorm2d, norm_op_kwargs=None,
                 dropout_op=nn.Dropout2d, dropout_op_kwargs=None,
                 nonlin=nn.LeakyReLU, nonlin_kwargs=None, deep_supervision=True, dropout_in_localization=False,
                 final_nonlin=softmax_helper, weightInitializer=InitWeights_He(1e-2), pool_op_kernel_sizes=None,
                 conv_kernel_sizes=None,
                 upscale_logits=False, convolutional_pooling=False, convolutional_upsampling=False,
                 max_num_features=None, basic_block=ConvDropoutNormNonlin,
                 seg_output_use_bias=False):
        """
        basically more flexible than v1, architecture is the same

        Does this look complicated? Nah bro. Functionality > usability

        This does everything you need, including world peace.

        Questions? -> <EMAIL>
        """
        super(Generic_UNet_VAE, self).__init__(input_channels, base_num_features, num_classes, num_pool, num_conv_per_stage,
                 feat_map_mul_on_downscale, conv_op,
                 norm_op, norm_op_kwargs,
                 dropout_op, dropout_op_kwargs,
                 nonlin, nonlin_kwargs, deep_supervision, dropout_in_localization,
                 final_nonlin, weightInitializer, pool_op_kernel_sizes,
                 conv_kernel_sizes,
                 upscale_logits, convolutional_pooling, convolutional_upsampling,
                 max_num_features, basic_block,
                 seg_output_use_bias)

        # Variational Auto-Encoder
        self.vae = VAE(320, outChans=input_channels)
    
    def forward(self, x):

        # DEBUG: input size torch.Size([1, 1, 80, 224, 128])
        # DEBUG: encoding size torch.Size([1, 320, 5, 7, 4])

        # DEBUG: input size torch.Size([2, 1, 80, 224, 128])
        # DEBUG: encoding size torch.Size([2, 320, 5, 7, 4])

        skips = []
        seg_outputs = []
        for d in range(len(self.conv_blocks_context) - 1):
            x = self.conv_blocks_context[d](x)
            # print('encoder', x.size())
            skips.append(x)
            if not self.convolutional_pooling:
                x = self.td[d](x)

        x = self.conv_blocks_context[-1](x)

        # VAE
        vae_outputs, vae_distr = self.vae(x)

        for u in range(len(self.tu)):
            x = self.tu[u](x)
            x = torch.cat((x, skips[-(u + 1)]), dim=1)
            x = self.conv_blocks_localization[u](x)
            seg_outputs.append(self.final_nonlin(self.seg_outputs[u](x)))
        
        if self._deep_supervision and self.do_ds:
            ds_outputs = tuple([seg_outputs[-1]] + [i(j) for i, j in
                                              zip(list(self.upscale_logits_ops)[::-1], seg_outputs[:-1][::-1])])
            return tuple([(ds_out, vae_outputs, vae_distr) for ds_out in ds_outputs])
        else:
            return seg_outputs[-1], vae_outputs, vae_distr

#    Copyright 2020 Division of Medical Image Computing, German Cancer Research Center (DKFZ), Heidelberg, Germany # #    Licensed under the Apache License, Version 2.0 (the "License"); #    you may not use this file except in compliance with the License. #    You may obtain a copy of the License at # #        http://www.apache.org/licenses/LICENSE-2.0 # #    Unless required by applicable law or agreed to in writing, software #    distributed under the License is distributed on an "AS IS" BASIS, #    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. #    See the License for the specific language governing permissions and #    limitations under the License. Trilinear interpolate to upsampling # out = self.up1(out) Variational Auto-Encoder Resampling block # print("ACTIVATE SIZE:", out.size()) # print("COV SIZE:", out.size()) # print("FLAT SIZE:", out.size()) # print("DENSE1 SIZE:", out_vd.size()) # ACTIVATE SIZE: torch.Size([2, 320, 5, 7, 4]) # COV SIZE: torch.Size([2, 16, 3, 4, 2]) # FLAT SIZE: torch.Size([2, 384]) # DENSE1 SIZE: torch.Size([2, 320]) # print("VDraw SIZE:", out.size()) #     VDraw SIZE: torch.Size([2, 160]) # DENSE2 SIZE: torch.Size([2, 3840])  match [2, 5,14,8] # ACTI2 SIZE: torch.Size([2, 3840]) # VIEW SIZE: torch.Size([2, 160, 3, 4, 2]) # UPSAMPL SIZE: torch.Size([2, 320, 6, 8, 4]) # print("DENSE2 SIZE:", out.size()) # print("ACTI2 SIZE:", out.size()) # print("VIEW SIZE:", out.size()) # Generate a Gaussian distribution with the given mean(128-d) and std(128-d) Variational Decoder block Decoder block Following UNet building blocks can be added by utilizing the properties this class exposes (TODO)          this one includes the bottleneck layer!          :param input_channels:         :param base_num_features:         :param num_blocks_per_stage:         :param feat_map_mul_on_downscale:         :param pool_op_kernel_sizes:         :param conv_kernel_sizes:         :param props: # decoder may need this # self.initial_conv = props['conv_op'](input_channels, base_num_features, 3, padding=1, **props['conv_op_kwargs']) # self.initial_norm = props['norm_op'](base_num_features, **props['norm_op_kwargs']) # self.initial_nonlin = props['nonlin'](**props['nonlin_kwargs']) # update current_input_features # cum_upsample = np.cumprod(np.vstack(self.stage_pool_kernel_size), axis=0).astype(int) :param x:         :param return_skips: if none then self.default_return_skips is used         :return: # x = self.initial_nonlin(self.initial_norm(self.initial_conv(x))) # + 1 for conv in skip in first block We assume the bottleneck is part of the encoder, so we can start with upsample -> concat here # we have one less as the first stage here is what comes after the # bottleneck # only used for upsample_logits # after we tu we concat features so now we have 2xfeatures_skip # skips come from the encoder. They are sorted so that the bottleneck is last in the list # what is maybe not perfect is that the TUs and stages here are sorted the other way around # so let's just reverse the order of skips # this is the bottleneck # seg_outputs are ordered so that the seg from the highest layer is first, the seg from # the bottleneck of the UNet last This only applies for num_conv_per_stage and convolutional_upsampling=True         not real vram consumption. just a constant term to which the vram consumption will be approx proportional         (+ offset for parameter storage)         :param patch_size:         :param num_pool_per_axis:         :param base_num_features:         :param max_num_features:         :return: # +1 for transpconv and +1 for conv in skip Variational Auto-Encoder : to group the features extracted by Encoder # out: (B, 320, 5, 14, 8) # out: (B, 160, 10, 28, 16) # print('OUT2 size',out.size()) # out: (B, 80, 20, 56, 32) # print('OUT3 size',out.size()) # out: (B, 40, 40, 112, 64) # print('OUT4 size',out.size()) # out: (B, 20, 80, 224, 128) # print('OUT5 size',out.size()) # out: (B, 1, 80, 224, 128) basically more flexible than v1, architecture is the same          Does this look complicated? Nah bro. Functionality > usability          This does everything you need, including world peace.          Questions? -> <EMAIL> # Variational Auto-Encoder # DEBUG: input size torch.Size([1, 1, 80, 224, 128]) # DEBUG: encoding size torch.Size([1, 320, 5, 7, 4]) # DEBUG: input size torch.Size([2, 1, 80, 224, 128]) # DEBUG: encoding size torch.Size([2, 320, 5, 7, 4]) # print('encoder', x.size()) # VAE