text: #
  # Copyright 2021 xxx. All Rights Reserved.
  #
  # 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.
  #
  # /*
  #  * @Author: dorming 
  #  * @Date: 2021-01-12 15:59:33 
  #  * @Last Modified by:   dorming 
  #  * @Last Modified time: 2021-01-12 15:59:33 
  #  */
  
  
  import numpy as np
  import torch
  import torch.nn.functional as F
  from torch import nn, optim
  from torch.nn import functional as F
  from torchvision import datasets, transforms
  from torchvision.utils import save_image
  import torch.nn as nn
  from torchsummary import summary
  import sys
  
  # sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
  sys.path.append("/home/zongdaoming/cv/multi-organ/multi-organ-ijcai/")
  from .BaseModelClass import BaseModel
  
  class UNet3D(BaseModel):
      """
      Implementations based on the Unet3D paper: https://arxiv.org/abs/1606.06650
      """
      def __init__(self, in_channels, n_classes, base_n_filter=8):
          super(UNet3D, self).__init__()
          self.in_channels = in_channels
          self.n_classes = n_classes
          self.base_n_filter = base_n_filter
  
          self.lrelu = nn.LeakyReLU()
          self.dropout3d = nn.Dropout3d(p=0.6)
          self.upsacle = nn.Upsample(scale_factor=2, mode='nearest')
          # typical first dim is the batch_size
          self.softmax = nn.Softmax(dim=1)
          self.conv3d_c1_1 = nn.Conv3d(self.in_channels, self.base_n_filter, kernel_size=3, stride=1, padding=1,
                                       bias=False)
          self.conv3d_c1_2 = nn.Conv3d(self.base_n_filter, self.base_n_filter, kernel_size=3, stride=1, padding=1,
                                       bias=False)
          self.lrelu_conv_c1 = self.lrelu_conv(self.base_n_filter, self.base_n_filter)                                     
          
          self.inorm3d_c1 = nn.InstanceNorm3d(self.base_n_filter)
  
          self.conv3d_c2 = nn.Conv3d(self.base_n_filter, self.base_n_filter * 2, kernel_size=3, stride=2, padding=1,
                                     bias=False)
          self.norm_lrelu_conv_c2 = self.norm_lrelu_conv(self.base_n_filter * 2, self.base_n_filter * 2)
          self.inorm3d_c2 = nn.InstanceNorm3d(self.base_n_filter * 2)
  
          self.conv3d_c3 = nn.Conv3d(self.base_n_filter * 2, self.base_n_filter * 4, kernel_size=3, stride=2, padding=1,
                                     bias=False)
          self.norm_lrelu_conv_c3 = self.norm_lrelu_conv(self.base_n_filter * 4, self.base_n_filter * 4)
          self.inorm3d_c3 = nn.InstanceNorm3d(self.base_n_filter * 4)
  
          self.conv3d_c4 = nn.Conv3d(self.base_n_filter * 4, self.base_n_filter * 8, kernel_size=3, stride=2, padding=1,
                                     bias=False)
          self.norm_lrelu_conv_c4 = self.norm_lrelu_conv(self.base_n_filter * 8, self.base_n_filter * 8)
          self.inorm3d_c4 = nn.InstanceNorm3d(self.base_n_filter * 8)
  
          self.conv3d_c5 = nn.Conv3d(self.base_n_filter * 8, self.base_n_filter * 16, kernel_size=3, stride=2, padding=1,
                                     bias=False)
          self.norm_lrelu_conv_c5 = self.norm_lrelu_conv(self.base_n_filter * 16, self.base_n_filter * 16)
          self.norm_lrelu_upscale_conv_norm_lrelu_l0 = self.norm_lrelu_upscale_conv_norm_lrelu(self.base_n_filter * 16,
                                                                                               self.base_n_filter * 8)
  
          self.conv3d_l0 = nn.Conv3d(self.base_n_filter * 8, self.base_n_filter * 8, kernel_size=1, stride=1, padding=0,
                                     bias=False)
          self.inorm3d_l0 = nn.InstanceNorm3d(self.base_n_filter * 8)
  
          self.conv_norm_lrelu_l1 = self.conv_norm_lrelu(self.base_n_filter * 16, self.base_n_filter * 16)
          self.conv3d_l1 = nn.Conv3d(self.base_n_filter * 16, self.base_n_filter * 8, kernel_size=1, stride=1, padding=0,
                                     bias=False)
          self.norm_lrelu_upscale_conv_norm_lrelu_l1 = self.norm_lrelu_upscale_conv_norm_lrelu(self.base_n_filter * 8,
                                                                                               self.base_n_filter * 4)
  
          self.conv_norm_lrelu_l2 = self.conv_norm_lrelu(self.base_n_filter * 8, self.base_n_filter * 8)
          self.conv3d_l2 = nn.Conv3d(self.base_n_filter * 8, self.base_n_filter * 4, kernel_size=1, stride=1, padding=0,
                                     bias=False)
          self.norm_lrelu_upscale_conv_norm_lrelu_l2 = self.norm_lrelu_upscale_conv_norm_lrelu(self.base_n_filter * 4,
                                                                                               self.base_n_filter * 2)
  
          self.conv_norm_lrelu_l3 = self.conv_norm_lrelu(self.base_n_filter * 4, self.base_n_filter * 4)
          self.conv3d_l3 = nn.Conv3d(self.base_n_filter * 4, self.base_n_filter * 2, kernel_size=1, stride=1, padding=0,
                                     bias=False)
          self.norm_lrelu_upscale_conv_norm_lrelu_l3 = self.norm_lrelu_upscale_conv_norm_lrelu(self.base_n_filter * 2,
                                                                                               self.base_n_filter)
  
          self.conv_norm_lrelu_l4 = self.conv_norm_lrelu(self.base_n_filter * 2, self.base_n_filter * 2)
          self.conv3d_l4 = nn.Conv3d(self.base_n_filter * 2, self.n_classes, kernel_size=1, stride=1, padding=0,
                                     bias=False)
  
          self.ds2_1x1_conv3d = nn.Conv3d(self.base_n_filter * 8, self.n_classes, kernel_size=1, stride=1, padding=0,
                                          bias=False)
          self.ds3_1x1_conv3d = nn.Conv3d(self.base_n_filter * 4, self.n_classes, kernel_size=1, stride=1, padding=0,
                                          bias=False)
  
          self.sigmod = nn.Sigmoid()
  
      def conv_norm_lrelu(self, feat_in, feat_out):
          return nn.Sequential(
              nn.Conv3d(feat_in, feat_out, kernel_size=3, stride=1, padding=1, bias=False),
              nn.InstanceNorm3d(feat_out),
              nn.LeakyReLU())
  
      def norm_lrelu_conv(self, feat_in, feat_out):
          return nn.Sequential(
              nn.InstanceNorm3d(feat_in),
              nn.LeakyReLU(),
              nn.Conv3d(feat_in, feat_out, kernel_size=3, stride=1, padding=1, bias=False))
  
      def lrelu_conv(self, feat_in, feat_out):
          return nn.Sequential(
              nn.LeakyReLU(),
              nn.Conv3d(feat_in, feat_out, kernel_size=3, stride=1, padding=1, bias=False))
  
      def norm_lrelu_upscale_conv_norm_lrelu(self, feat_in, feat_out):
          return nn.Sequential(
              nn.InstanceNorm3d(feat_in),
              nn.LeakyReLU(),
              nn.Upsample(scale_factor=2, mode='nearest'),
              # should be feat_in*2 or feat_in
              nn.Conv3d(feat_in, feat_out, kernel_size=3, stride=1, padding=1, bias=False),
              nn.InstanceNorm3d(feat_out),
              nn.LeakyReLU())
  
      def sample_gumbel(shape, eps=1e-20):
          U = torch.rand(shape)
          if is_cuda:
              U = U.cuda()
          return -torch.log(-torch.log(U + eps) + eps)
  
  
      def gumbel_softmax_sample(logits, temperature):
          y = logits + sample_gumbel(logits.size())
          return F.softmax(y / temperature, dim=-1)
  
  
      def gumbel_softmax(logits, temperature, hard=False):
          """
          ST-gumple-softmax
          input: [*, n_class]
          return: flatten --> [*, n_class] an one-hot vector
          """
          y = gumbel_softmax_sample(logits, temperature)
          if not hard:
              return y.view(-1, latent_dim * categorical_dim)
  
          shape = y.size()
          _, ind = y.max(dim=-1)
          y_hard = torch.zeros_like(y).view(-1, shape[-1])
          y_hard.scatter_(1, ind.view(-1, 1), 1)
          y_hard = y_hard.view(*shape)
          # Set gradients w.r.t. y_hard gradients w.r.t. y
          y_hard = (y_hard - y).detach() + y
          return y_hard.view(-1, latent_dim * categorical_dim)
          
  
      def forward(self, x):
          # Level 1 context pathway 
          out = self.conv3d_c1_1(x)
          residual_1 = out
          out = self.lrelu(out)
          out = self.conv3d_c1_2(out)
          out = self.dropout3d(out)
          out = self.lrelu_conv_c1(out)
          # Element Wise Summation
          out += residual_1
          context_1 = self.lrelu(out)
          out = self.inorm3d_c1(out)
          out = self.lrelu(out)
  
          # Level 2 context pathway
          out = self.conv3d_c2(out)
          residual_2 = out
          out = self.norm_lrelu_conv_c2(out)
          out = self.dropout3d(out)
          out = self.norm_lrelu_conv_c2(out)
          out += residual_2
          out = self.inorm3d_c2(out)
          out = self.lrelu(out)
          context_2 = out
  
          # Level 3 context pathway
          out = self.conv3d_c3(out)
          residual_3 = out
          out = self.norm_lrelu_conv_c3(out)
          out = self.dropout3d(out)
          out = self.norm_lrelu_conv_c3(out)
          out += residual_3
          out = self.inorm3d_c3(out)
          out = self.lrelu(out)
          context_3 = out
  
          # Level 4 context pathway
          out = self.conv3d_c4(out)
          residual_4 = out
          out = self.norm_lrelu_conv_c4(out)
          out = self.dropout3d(out)
          out = self.norm_lrelu_conv_c4(out)
          out += residual_4
          out = self.inorm3d_c4(out)
          out = self.lrelu(out)
          context_4 = out
  
  
          # Level 5 context pathway
          out = self.conv3d_c5(out)
          residual_5 = out
          out = self.norm_lrelu_conv_c5(out)
          out = self.dropout3d(out)
          out = self.norm_lrelu_conv_c5(out)
          out += residual_5
          out = self.norm_lrelu_upscale_conv_norm_lrelu_l0(out)
  
          out = self.conv3d_l0(out)
          out = self.inorm3d_l0(out)
          out = self.lrelu(out)
  
          # Level 1 localization pathway
          out = torch.cat([out, context_4], dim=1)
          out = self.conv_norm_lrelu_l1(out)
          out = self.conv3d_l1(out)
          out = self.norm_lrelu_upscale_conv_norm_lrelu_l1(out)
  
          # Level 2 localization pathway
          # print(out.shape)
          # print(context_3.shape)
          out = torch.cat([out, context_3], dim=1)
          out = self.conv_norm_lrelu_l2(out)
          ds2 = out
          out = self.conv3d_l2(out)
          out = self.norm_lrelu_upscale_conv_norm_lrelu_l2(out)
  
  
          # Level 3 localization pathway
          out = torch.cat([out, context_2], dim=1)
          out = self.conv_norm_lrelu_l3(out)
          ds3 = out
          out = self.conv3d_l3(out)
          out = self.norm_lrelu_upscale_conv_norm_lrelu_l3(out)
  
          # Level 4 localization pathway
          out = torch.cat([out, context_1], dim=1)
          out = self.conv_norm_lrelu_l4(out)
          out_pred = self.conv3d_l4(out)
  
  
          ds2_1x1_conv = self.ds2_1x1_conv3d(ds2)
          ds1_ds2_sum_upscale = self.upsacle(ds2_1x1_conv)
          ds3_1x1_conv = self.ds3_1x1_conv3d(ds3)
          ds1_ds2_sum_upscale_ds3_sum = ds1_ds2_sum_upscale + ds3_1x1_conv
          ds1_ds2_sum_upscale_ds3_sum_upscale = self.upsacle(ds1_ds2_sum_upscale_ds3_sum)
  
          out = out_pred + ds1_ds2_sum_upscale_ds3_sum_upscale
          seg_layer = out
          return seg_layer     
       
      def test(self,device='cpu'):
          input_tensor = torch.rand(1, 2, 32, 32, 32)
          ideal_out = torch.rand(1, self.n_classes, 32, 32, 32)
          out = self.forward(input_tensor)
          assert ideal_out.shape == out.shape
          device = torch.device(device)
          summary(self.to(device), input_size=(2,32,32,32), device='cpu')
          print("*UNet3D test is complete")
  
  
  
  # if __name__ == "__main__":
  #     unet3d = UNet3D(in_channels=2,n_classes=10 , base_n_filter=8)
  #     unet3d.test()
  
  
  
  
  
  
  
  
  
  
  
  
  
  
  
  
  
  
  
  
  