text: """
  Implementations of different CNNs
  
  by Christopher M. Sandino (sandino@stanford.edu), 2019.
  
  """
  import sys
  import torch
  from torch import nn
  from utils.transforms import center_crop
  import torch.nn.functional as F
  
  
  class SeparableConv3d(nn.Module):
      """
      A separable 3D convolutional operator.
      """
      def __init__(self, in_chans, out_chans, kernel_size, spatial_chans=None, act_type='relu'):
          """
          Args:
              in_chans (int): Number of channels in the input.
              out_chans (int): Number of channels in the output.
              kernel_size (int): Size of kernel (repeated for all three dimensions).
          """
          super().__init__()
  
          sp_kernel_size = (1, kernel_size, kernel_size)
          sp_pad_size = (0, 1, 1)
          t_kernel_size  = (kernel_size, 1, 1) 
          t_pad_size = (1, 0, 0)
  
          if spatial_chans is None:
              # Force number of spatial features, such that the total number of
              # parameters is the same as a nn.Conv3D(in_chans, out_chans)
              spatial_chans = (kernel_size**3)*in_chans*out_chans 
              spatial_chans /= (kernel_size**2)*in_chans + kernel_size*out_chans
              spatial_chans = int(spatial_chans)
  
          # Define each layer in SeparableConv3d block
          spatial_conv = nn.Conv3d(in_chans, spatial_chans, kernel_size=sp_kernel_size, padding=sp_pad_size)
          temporal_conv = nn.Conv3d(spatial_chans, out_chans, kernel_size=t_kernel_size, padding=t_pad_size)
  
          # Define choices for intermediate activation layer
          activations = nn.ModuleDict([
                  ['none',  nn.Identity()],
                  ['relu',  nn.ReLU()]
          ])
  
          # Define the forward pass
          self.layers = nn.Sequential(spatial_conv, activations[act_type], temporal_conv)
  
  
      def forward(self, input):
          """
          Args:
              input (torch.Tensor): Input tensor of shape [batch_size, self.in_chans, depth, width, height]
  
          Returns:
              (torch.Tensor): Output tensor of shape [batch_size, self.out_chans, depth, width, height]
          """
          return self.layers(input)
  
  
  class ConvBlock(nn.Module):
      """
      A 3D Convolutional Block that consists of Norm -> ReLU -> Dropout -> Conv
  
      Based on implementation described by: 
          K He, et al. "Identity Mappings in Deep Residual Networks" arXiv:1603.05027
      """
      def __init__(self, in_chans, out_chans, kernel_size, drop_prob, 
                   conv_type='conv3d', act_type='relu', norm_type='none'):
          """
          Args:
              in_chans (int): Number of channels in the input.
              out_chans (int): Number of channels in the output.
              drop_prob (float): Dropout probability.
          """
          super().__init__()
  
          self.in_chans = in_chans
          self.out_chans = out_chans
          self.drop_prob = drop_prob
  
          # Define choices for each layer in ConvBlock
          normalizations = nn.ModuleDict([
                  ['none',     nn.Identity()],
                  ['instance', nn.InstanceNorm3d(in_chans, affine=False)],
                  ['batch',    nn.BatchNorm3d(in_chans, affine=False)]
          ])
          activations = nn.ModuleDict([
                  ['relu',  nn.ReLU()],
                  ['leaky_relu', nn.LeakyReLU()]
          ])
          dropout = nn.Dropout3d(p=drop_prob, inplace=True)
  
          # Note: don't use ModuleDict here. Otherwise, the parameters for the un-selected
          # convolution type will still be initialized and added to model.parameters()
  #         print(conv_type=='conv3d')
          if conv_type == 'conv3d':
              convolution = nn.Conv3d(in_chans, out_chans, kernel_size=kernel_size, padding=1)
  #             print("is conv3d")
  #             sys.exit()
          else:
              convolution = SeparableConv3d(in_chans, out_chans, kernel_size=kernel_size)
  #             print("is conv2p1d")
              
  #             print("Separable conv")
  #             sys.exit()
              
  
          # Define forward pass
          self.layers = nn.Sequential(
              normalizations[norm_type], 
              activations[act_type],
              dropout, convolution
          )
  
      def forward(self, input):
          """
          Args:
              input (torch.Tensor): Input tensor of shape [batch_size, self.in_chans, depth, width, height]
  
          Returns:
              (torch.Tensor): Output tensor of shape [batch_size, self.out_chans, depth, width, height]
          """
          return self.layers(input)
  
      def __repr__(self):
          return f'ConvBlock3D(in_chans={self.in_chans}, out_chans={self.out_chans}, ' \
              f'drop_prob={self.drop_prob})'
  
  
  class ResBlock(nn.Module):
      """
      A ResNet block that consists of two convolutional layers followed by a residual connection.
      """
  
      def __init__(self, in_chans, out_chans, kernel_size, drop_prob, conv_type='conv3d'):
          """
          Args:
              in_chans (int): Number of channels in the input.
              out_chans (int): Number of channels in the output.
              drop_prob (float): Dropout probability.
          """
          super().__init__()
  
          self.layers = nn.Sequential(
              ConvBlock(in_chans, out_chans, kernel_size, drop_prob,conv_type),
              ConvBlock(out_chans, out_chans, kernel_size, drop_prob,conv_type)
          )
  
          if in_chans != out_chans:
              self.resample = nn.Conv3d(in_chans, out_chans, kernel_size=1)
          else:
              self.resample = nn.Identity()
  
      def forward(self, input):
          """
          Args:
              input (torch.Tensor): Input tensor of shape [batch_size, self.in_chans, depth, width, height]
  
          Returns:
              (torch.Tensor): Output tensor of shape [batch_size, self.out_chans, depth, width, height]
          """
  
          # To have a residual connection, number of inputs must be equal to outputs
          shortcut = self.resample(input)
  
          return self.layers(input) + shortcut
  
  
  class ResNet(nn.Module):
      """
      Prototype for 3D ResNet architecture
      """
  
      def __init__(self, num_resblocks, in_chans, chans, kernel_size, drop_prob, circular_pad=True,conv_type='conv3d'):
          """
  
          """
          super().__init__()
  
          self.circular_pad = circular_pad
          self.pad_size = 2*num_resblocks + 1
  
          # Declare ResBlock layers
          self.res_blocks = nn.ModuleList([ResBlock(in_chans, chans, kernel_size, drop_prob,conv_type)])
          for _ in range(num_resblocks-1):
              self.res_blocks += [ResBlock(chans, chans, kernel_size, drop_prob,conv_type)]
  
          # Declare final conv layer (down-sample to original in_chans)
          self.final_layer = nn.Conv3d(chans, in_chans, kernel_size=kernel_size, padding=1)
  
  
      def forward(self, input):
          """
          Args:
              input (torch.Tensor): Input tensor of shape [batch_size, self.in_chans, depth, width, height]
  
          Returns:
              (torch.Tensor): Output tensor of shape [batch_size, self.in_chans, depth, width, height]
          """
  
          orig_shape = input.shape
          if self.circular_pad:
              input = nn.functional.pad(input, (0,0,0,0) + 2*(self.pad_size,), mode='circular')
              # input = nn.functional.pad(input, 4*(self.pad_size,) + (0,0), mode='replicate')
  
          # Perform forward pass through the network
          output = input
          for res_block in self.res_blocks:
              output = res_block(output)
          output = self.final_layer(output) + input
  
          return center_crop(output, orig_shape)
  
  