import torch

import torch.nn as nn

import torch.nn.functional as F

class ConvUnit(nn.Module):

  """

    Convolution Unit -

    for  now : (Conv3D -> BatchNorm -> ReLu) * 2

    Try modifying to Residual convolutions

  """

  def __init__(self, in_channels, out_channels):

    super(ConvUnit, self).__init__()

    self.double_conv = nn.Sequential(

        nn.Conv3d(in_channels, out_channels, kernel_size = 3, padding = 1),

        nn.BatchNorm3d(out_channels),

        nn.ReLU(inplace=True), # inplace=True means it changes the input directly, input is lost

        nn.Conv3d(out_channels, out_channels, kernel_size = 3, padding = 1),

        nn.BatchNorm3d(out_channels),

        nn.ReLU(inplace=True)

      )

  def forward(self,x):

    return self.double_conv(x)

class EncoderUnit(nn.Module):

  """

    An Encoder Unit with the ConvUnit and MaxPool

  """

  def __init__(self, in_channels, out_channels):

    super(EncoderUnit, self).__init__()

    self.encoder = nn.Sequential(

        nn.MaxPool3d(2),

        ConvUnit(in_channels, out_channels)

    )

  def forward(self, x):

    return self.encoder(x)

class DecoderUnit(nn.Module):

  """

    ConvUnit and upsample with Upsample or convTranspose

  """

  def __init__(self, in_channels, out_channels, bilinear=False):

    super().__init__()

    if bilinear:

      # Only for 2D model

      self.up = nn.Upsample(scale_factor=2, mode="bilinear", align_corners=True)

    else:

      self.up = nn.ConvTranspose3d(in_channels // 2, in_channels // 2, kernel_size=2, stride=2)

    self.conv = ConvUnit(in_channels, out_channels)

  def forward(self, x1, x2):

      x1 = self.up(x1)

      diffZ = x2.size()[2] - x1.size()[2]

      diffY = x2.size()[3] - x1.size()[3]

      diffX = x2.size()[4] - x1.size()[4]

      x1 = F.pad(x1, [diffX // 2, diffX - diffX // 2, diffY // 2, diffY - diffY // 2, diffZ // 2, diffZ - diffZ // 2])

      x = torch.cat([x2, x1], dim=1)

      return self.conv(x)

class OutConv(nn.Module):

  def __init__(self, in_channels, out_channels):

    super(OutConv, self).__init__()

    self.conv = nn.Conv3d(in_channels, out_channels, kernel_size = 1)

  def forward(self, x):

    return self.conv(x)

###########   Model :

class UNet(nn.Module):

  def __init__(self, in_channels, n_classes, s_channels, bilinear = False):

    super(UNet, self).__init__()

    self.in_channels = in_channels

    self.n_classes = n_classes

    self.s_channels = s_channels

    self.bilinear = bilinear

    self.conv = ConvUnit(in_channels, s_channels)

    self.enc1 = EncoderUnit(s_channels, 2 * s_channels)

    self.enc2 = EncoderUnit(2 * s_channels, 4 * s_channels)

    self.enc3 = EncoderUnit(4 * s_channels, 8 * s_channels)

    self.enc4 = EncoderUnit(8 * s_channels, 8 * s_channels)

    self.dec1 = DecoderUnit(16 * s_channels, 4 * s_channels, self.bilinear)

    self.dec2 = DecoderUnit(8 * s_channels, 2 * s_channels, self.bilinear)

    self.dec3 = DecoderUnit(4 * s_channels, s_channels, self.bilinear)

    self.dec4 = DecoderUnit(2 * s_channels, s_channels, self.bilinear)

    self.out = OutConv(s_channels, n_classes)

  def forward(self, x):

    x1 = self.conv(x)

    x2 = self.enc1(x1)

    x3 = self.enc2(x2)

    x4 = self.enc3(x3)

    x5 = self.enc4(x4)

    mask = self.dec1(x5, x4)

    mask = self.dec2(mask, x3)

    mask = self.dec3(mask, x2)

    mask = self.dec4(mask, x1)

    mask = self.out(mask)

    return mask