import os
from abc import ABC, abstractmethod

import torch
import torch.nn as nn
from torchsummary import summary

from base_model import BaseModel

def passthrough(x, **kwargs):
  return x

def ELUCons(elu, nchan):
  if elu:
    return nn.ELU(inplace=True)
  else:
    return nn.PReLU(nchan)

class LUConv(nn.Module):
  def __init__(self, nchan, elu):
    super(LUConv, self).__init__()
    self.relu1 = ELUCons(elu, nchan)
    self.conv1 = nn.Conv3d(nchan, nchan, kernel_size=5, padding=2)
    self.bn1 = torch.nn.BatchNorm3d(nchan)

  def forward(self, x):
    out = self.relu1(self.bn1(self.conv1(x)))
    return out

def _make_nConv(nchan, depth, elu):
    layers = []
    for _ in range(depth):
        layers.append(LUConv(nchan, elu))
    return nn.Sequential(*layers)

class InputTransition(nn.Module):
    def __init__(self, in_channels, elu):
        super(InputTransition, self).__init__()
        self.num_features = 16
        self.in_channels = in_channels
        self.conv1 = nn.Conv3d(self.in_channels, self.num_features, kernel_size=5, padding=2)
        self.bn1 = torch.nn.BatchNorm3d(self.num_features)
        self.relu1 = ELUCons(elu, self.num_features)

    def forward(self, x):
        out = self.conv1(x)
        repeat_rate = int(self.num_features / self.in_channels)
        out = self.bn1(out)
        x16 = x.repeat(1, repeat_rate, 1, 1, 1)
        return self.relu1(torch.add(out, x16))

class DownTransition(nn.Module):
  def __init__(self, inChans, nConvs, elu, dropout=False):
    super(DownTransition, self).__init__()
    outChans = 2 * inChans
    self.down_conv = nn.Conv3d(inChans, outChans, kernel_size=2, stride=2)
    self.bn1 = torch.nn.BatchNorm3d(outChans)

    self.do1 = passthrough
    self.relu1 = ELUCons(elu, outChans)
    self.relu2 = ELUCons(elu, outChans)
    if dropout:
        self.do1 = nn.Dropout3d()
    self.ops = _make_nConv(outChans, nConvs, elu)

  def forward(self, x):
    down = self.relu1(self.bn1(self.down_conv(x)))
    out = self.do1(down)
    out = self.ops(out)
    out = self.relu2(torch.add(out, down))
    return out

class UpTransition(nn.Module):
  def __init__(self, inChans, outChans, nConvs, elu, dropout=False):
    super(UpTransition, self).__init__()
    self.up_conv = nn.ConvTranspose3d(inChans, outChans // 2, kernel_size=2, stride=2)

    self.bn1 = torch.nn.BatchNorm3d(outChans // 2)
    self.do1 = passthrough
    self.do2 = nn.Dropout3d()
    self.relu1 = ELUCons(elu, outChans // 2)
    self.relu2 = ELUCons(elu, outChans)
    if dropout:
        self.do1 = nn.Dropout3d()
    self.ops = _make_nConv(outChans, nConvs, elu)

  def forward(self, x, skipx):
    out = self.do1(x)
    skipxdo = self.do2(skipx)
    out = self.relu1(self.bn1(self.up_conv(out)))
    xcat = torch.cat((out, skipxdo), 1)
    out = self.ops(xcat)
    out = self.relu2(torch.add(out, xcat))
    return out

class OutputTransition(nn.Module):
  def __init__(self, in_channels, classes, elu):
    super(OutputTransition, self).__init__()
    self.classes = classes
    self.conv1 = nn.Conv3d(in_channels, classes, kernel_size=5, padding=2)
    self.bn1 = torch.nn.BatchNorm3d(classes)

    self.conv2 = nn.Conv3d(classes, classes, kernel_size=1)
    self.relu1 = ELUCons(elu, classes)

  def forward(self, x):
    # convolve 32 down to channels as the desired classes
    out = self.relu1(self.bn1(self.conv1(x)))
    out = self.conv2(out)
    return out