import numpy as np
import tensorflow as tf


class NeuralNetwork(object):
  def conv2d(self, input, filter_shape, scope_name, strides=[1, 1, 1, 1]):
    xavier_initializer = tf.contrib.layers.xavier_initializer
    with tf.variable_scope(scope_name):
      W = tf.get_variable(initializer=xavier_initializer(uniform=False),
                          name="W",
                          shape=filter_shape)
      b = tf.get_variable(initializer=xavier_initializer(uniform=False),
                          name="b",
                          shape=[filter_shape[3]])
      out = tf.nn.conv2d(input, W, padding="SAME", strides=strides)
      out = tf.nn.bias_add(out, b)
      return out

  def conv2d_relu(self, input, filter_shape, scope_name, strides=[1, 1, 1, 1]):
    xavier_initializer = tf.contrib.layers.xavier_initializer
    with tf.variable_scope(scope_name):
      W = tf.get_variable(initializer=xavier_initializer(uniform=False),
                          name="W",
                          shape=filter_shape)
      b = tf.get_variable(initializer=xavier_initializer(uniform=False),
                          name="b",
                          shape=[filter_shape[3]])
      out = tf.nn.conv2d(input, W, padding="SAME", strides=strides)
      out = tf.nn.bias_add(out, b)
      out = tf.nn.relu(out, name="out")
      return out

  def maxpool2d(self, input, scope_name, pool_shape=[1, 2, 2, 1], strides=[1, 2, 2, 1]):
    with tf.variable_scope(scope_name):
      out = tf.nn.max_pool(input,
                           ksize=pool_shape,
                           name="out",
                           padding="SAME",
                           strides=strides)
      return out

  def avgpool2d(self, input, scope_name, pool_shape=[1, 2, 2, 1], strides=[1, 2, 2, 1]):
    with tf.variable_scope(scope_name):
      out = tf.nn.avg_pool(input,
                           ksize=pool_shape,
                           name="out",
                           padding="SAME",
                           strides=strides)
      return out


  def dropout(self, input, keep_prob, scope_name):
    with tf.variable_scope(scope_name):
      out = tf.nn.dropout(input, keep_prob, name="out")
      return out

  def fc(self, input, output_shape, scope_name):
    xavier_initializer = tf.contrib.layers.xavier_initializer
    with tf.variable_scope(scope_name):
      input_shape = input.shape[1:]
      input_shape = int(np.prod(input_shape))
      W = tf.get_variable(name="W",
                          shape=[input_shape, output_shape],
                          initializer=xavier_initializer(uniform=False))
      b = tf.get_variable(name="b",
                          shape=[output_shape],
                          initializer=xavier_initializer(uniform=False))
      input = tf.reshape(input, [-1, input_shape])
      # out = tf.nn.relu(tf.add(tf.matmul(input, W), b), name="out")
      out = tf.add(tf.matmul(input, W), b, name="out")
      return out

  def deconv2d(self, input, filter_shape, num_outputs, scope_name, strides=[1, 1]):
    xavier_initializer = tf.contrib.layers.xavier_initializer
    xavier_initializer_conv2d = tf.contrib.layers.xavier_initializer_conv2d
    with tf.variable_scope(scope_name):
      out = tf.contrib.layers.conv2d_transpose(input,
                                               # activation_fn=tf.nn.relu,
                                               activation_fn=None,
                                               biases_initializer=xavier_initializer(uniform=False),
                                               kernel_size=filter_shape,
                                               num_outputs=num_outputs,
                                               padding="SAME",
                                               stride=strides,
                                               weights_initializer=xavier_initializer_conv2d(uniform=False))
      out = tf.identity(out, name="out")
      return out

  def upsample(self, input, scope_name, factor=[2, 2]):
    size = [int(input.shape[1] * factor[0]), int(input.shape[2] * factor[1])]
    with tf.variable_scope(scope_name):
      out = tf.image.resize_bilinear(input, size=size, align_corners=None, name="out")
      return out


class ConvEncoder(NeuralNetwork):
  def __init__(self, input_shape, keep_prob, scope_name="encoder"):
    self.input_shape = input_shape
    self.keep_prob = keep_prob
    self.scope_name = scope_name

  def build_graph(self, input):
    with tf.variable_scope(self.scope_name):
      conv1 = self.conv2d_relu(input, filter_shape=[3, 3, 1, 8], scope_name="conv1")  # (232, 196, 8)
      pool1 = self.maxpool2d(conv1, scope_name="pool1")  # (116, 98, 8)
      drop1 = self.dropout(pool1, keep_prob=self.keep_prob, scope_name="drop1")
      conv2 = self.conv2d_relu(drop1, filter_shape=[5, 5, 8, 16], scope_name="conv2")  # (116, 98, 16)
      pool2 = self.maxpool2d(conv2, scope_name="pool2")  # (58, 49, 16)
      drop2 = self.dropout(pool2, keep_prob=self.keep_prob, scope_name="drop2")
      drop2 = tf.reshape(drop2, shape=[-1, 58*49*16])  # (45472,)
      fc1 = self.fc(drop2, output_shape=1024, scope_name="fc1")
      drop3 = self.dropout(fc1, keep_prob=self.keep_prob, scope_name="drop3")
      fc2 = self.fc(drop3, output_shape=256, scope_name="fc2")
      out = tf.identity(fc2, name="out")

    return out


class DeconvDecoder(NeuralNetwork):
  def __init__(self, keep_prob, output_shape, scope_name="decoder"):
    self.keep_prob = keep_prob
    self.output_shape = output_shape
    self.scope_name = scope_name

  def build_graph(self, input):
    with tf.variable_scope(self.scope_name):
      fc1 = self.fc(input, output_shape=1024, scope_name="fc1")
      drop1 = self.dropout(fc1, keep_prob=self.keep_prob, scope_name="drop1")
      fc2 = self.fc(drop1, output_shape=58*49*16, scope_name="fc2")
      drop2 = self.dropout(fc2, keep_prob=self.keep_prob, scope_name="drop2")
      drop2 = tf.reshape(drop2, shape=[-1, 58, 49, 16])
      up1 = self.upsample(drop2, scope_name="up1", factor=[2, 2])  # (116, 98, 16)
      deconv1 = self.deconv2d(up1, filter_shape=[5, 5], num_outputs=8, scope_name="deconv1")  # (116, 98, 8)
      up2 = self.upsample(deconv1, scope_name="up2", factor=[2, 2])
      deconv2 = self.deconv2d(up2, filter_shape=[3, 3], num_outputs=1, scope_name="deconv2")  # (232, 196, 1)
      out = tf.identity(deconv2, name="out")

    return out

class ConvEncoderNew(NeuralNetwork):
  def __init__(self, input_shape, keep_prob, scope_name="encoder_new"):
    self.input_shape = input_shape
    self.keep_prob = keep_prob
    self.scope_name = scope_name

  def build_graph(self, input):
    #David's ATLAS MODEL (changed maxpool2d to avepool2d)
    with tf.variable_scope(self.scope_name):
      conv1 = self.conv2d_relu(input, filter_shape=[3, 3, 1, 8], scope_name="conv1")  # (232, 196, 8)
      pool1 = self.avgpool2d(conv1, scope_name="pool1")  # (116, 98, 8)
      drop1 = self.dropout(pool1, keep_prob=self.keep_prob, scope_name="drop1")
      conv2 = self.conv2d_relu(drop1, filter_shape=[5, 5, 8, 16], scope_name="conv2")  # (116, 98, 16)
      pool2 = self.avgpool2d(conv2, scope_name="pool2")  # (58, 49, 16)
      drop2 = self.dropout(pool2, keep_prob=self.keep_prob, scope_name="drop2")
      drop2 = tf.reshape(drop2, shape=[-1, 58*49*16])  # (45472,)
      fc1 = self.fc(drop2, output_shape=1024, scope_name="fc1")
      drop3 = self.dropout(fc1, keep_prob=self.keep_prob, scope_name="drop3")
      fc2 = self.fc(drop3, output_shape=256, scope_name="fc2")
      out = tf.identity(fc2, name="out")

    return out

class DeconvDecoderNew(NeuralNetwork):
  def __init__(self, keep_prob, output_shape, scope_name="decoder_new"):
    self.keep_prob = keep_prob
    self.output_shape = output_shape
    self.scope_name = scope_name

  def build_graph(self, input):
    #David's ATLAS MODEL
    with tf.variable_scope(self.scope_name):
      fc1 = self.fc(input, output_shape=1024, scope_name="fc1")
      drop1 = self.dropout(fc1, keep_prob=self.keep_prob, scope_name="drop1")
      fc2 = self.fc(drop1, output_shape=58*49*16, scope_name="fc2")
      drop2 = self.dropout(fc2, keep_prob=self.keep_prob, scope_name="drop2")
      drop2 = tf.reshape(drop2, shape=[-1, 58, 49, 16])
      up1 = self.upsample(drop2, scope_name="up1", factor=[2, 2])  # (116, 98, 16)
      deconv1 = self.deconv2d(up1, filter_shape=[5, 5], num_outputs=8, scope_name="deconv1")  # (116, 98, 8)
      up2 = self.upsample(deconv1, scope_name="up2", factor=[2, 2])
      deconv2 = self.deconv2d(up2, filter_shape=[3, 3], num_outputs=1, scope_name="deconv2")  # (232, 196, 1)
      out = tf.identity(deconv2, name="out")

    return out


class ConvEncoderUNet(NeuralNetwork):
  def __init__(self, input_shape, keep_prob, scope_name="encoder_unet"):
    self.input_shape = input_shape
    self.keep_prob = keep_prob
    self.scope_name = scope_name

  def build_graph(self, input):
    with tf.variable_scope(self.scope_name):
      # 233 x 196
      paddings = tf.constant([[0, 0], [0, 0], [2, 2], [0, 0]])
      # 233 x 200
      padded_input = tf.pad(input, paddings, "CONSTANT")  
      conv1 = self.conv2d_relu(padded_input, filter_shape=[3, 3, 1, 64], scope_name="conv1")  # (230, 194, 64)
      conv2 = self.conv2d_relu(conv1, filter_shape=[3, 3, 64, 64], scope_name="conv2")  # (228, 192, 64)
      pool1 = self.maxpool2d(conv2, scope_name="pool1") # (114, 96, 64)

      conv3 = self.conv2d_relu(pool1, filter_shape=[3, 3, 64, 128], scope_name="conv3")  # (b, 112, 94, 128)
      conv4 = self.conv2d_relu(conv3, filter_shape=[3, 3, 128, 128], scope_name="conv4")  # (b, 110, 92, 128)
      pool2 = self.maxpool2d(conv4, scope_name="pool2") # (55, 46, 128)

      conv5 = self.conv2d_relu(pool2, filter_shape=[3, 3, 128, 256], scope_name="conv5")  # (b, 58, 49, 256)
      conv6 = self.conv2d_relu(conv5, filter_shape=[3, 3, 256, 256], scope_name="conv6")  # (b, 58, 49, 256)
      pool3 = self.maxpool2d(conv6, scope_name="pool3")

      conv7 = self.conv2d_relu(pool3, filter_shape=[3, 3, 256, 512], scope_name="conv7")  # (b, 116, 98, 128)
      conv8 = self.conv2d_relu(conv7, filter_shape=[3, 3, 512, 512], scope_name="conv8")  # (b, 116, 98, 128)
      cache = (conv6, conv4, conv2)
      out = conv8
    return (out, cache)
    
class DeconvDecoderUNet(NeuralNetwork):
  def __init__(self, keep_prob, output_shape, scope_name="decoder_unet"):
    self.keep_prob = keep_prob
    self.output_shape = output_shape
    self.scope_name = scope_name

  def build_graph(self, input, cache):
    conv6, conv4, conv2 = cache
    with tf.variable_scope(self.scope_name):
      up1 = self.upsample(input, scope_name="up1", factor=[2, 2])  # (b, 116, 98, 256)
      deconv1 = self.deconv2d(up1, filter_shape=[2, 2], num_outputs=256, scope_name="deconv1")  # (b, 116, 98, 128)

      concat1 = tf.concat([conv6, deconv1], axis=3)  # (b, 116, 98, 256)
      conv9 = self.conv2d_relu(concat1, filter_shape=[3, 3, 512, 512], scope_name="conv9")  # (b, 116, 98, 128)
      conv10 = self.conv2d_relu(conv9, filter_shape=[3, 3, 512, 256], scope_name="conv10")  # (b, 116, 98, 128)

      up2 = self.upsample(conv10, scope_name="up2", factor=[2, 2])  # (b, 116, 98, 256)
      deconv2 = self.deconv2d(up2, filter_shape=[2, 2], num_outputs=128, scope_name="deconv2")  # (b, 116, 98, 128)
      concat2 = tf.concat([conv4, deconv2], axis=3)  # (b, 116, 98, 256)
      conv11 = self.conv2d_relu(concat2, filter_shape=[3, 3, 256, 256], scope_name="conv11")
      conv12 = self.conv2d_relu(conv11, filter_shape=[3, 3, 256, 128], scope_name="conv12")

      up3 = self.upsample(conv12, scope_name="up3", factor=[2, 2])  # (b, 116, 98, 256)
      deconv3 = self.deconv2d(up3, filter_shape=[2, 2], num_outputs=64, scope_name="deconv3")  # (b, 116, 98, 128)
      concat3 = tf.concat([conv2, deconv3], axis=3)  # (b, 116, 98, 256)
      conv13 = self.conv2d_relu(concat3, filter_shape=[3, 3, 128, 128], scope_name="conv13")
      conv14 = self.conv2d_relu(conv13, filter_shape=[3, 3, 128, 64], scope_name="conv14")
      conv15 = self.conv2d_relu(conv14, filter_shape=[1, 1, 64, 1], scope_name="conv15")

      unpadded_out = tf.identity(conv15, name="out")
      out = unpadded_out[:, :, 2:-2, :]
      print(np.shape(out))
    return out

class ConvEncoderDeepUNet(NeuralNetwork):
  def __init__(self, input_shape, keep_prob, scope_name="encoder_unet_deep"):
    self.input_shape = input_shape
    self.keep_prob = keep_prob
    self.scope_name = scope_name

  def build_graph(self, input):
    with tf.variable_scope(self.scope_name):
      # 233 x 196
      paddings = tf.constant([[0, 0], [4, 4], [6, 6], [0, 0]])
      # 233 x 200
      padded_input = tf.pad(input, paddings, "CONSTANT")  
      conv1 = self.conv2d_relu(padded_input, filter_shape=[3, 3, 1, 64], scope_name="conv1")  # (230, 194, 64)
      conv2 = self.conv2d_relu(conv1, filter_shape=[3, 3, 64, 64], scope_name="conv2")  # (228, 192, 64)
      pool1 = self.maxpool2d(conv2, scope_name="pool1") # (114, 96, 64)

      conv3 = self.conv2d_relu(pool1, filter_shape=[3, 3, 64, 128], scope_name="conv3")  # (b, 112, 94, 128)
      conv4 = self.conv2d_relu(conv3, filter_shape=[3, 3, 128, 128], scope_name="conv4")  # (b, 110, 92, 128)
      pool2 = self.maxpool2d(conv4, scope_name="pool2") # (55, 46, 128)

      conv5 = self.conv2d_relu(pool2, filter_shape=[3, 3, 128, 256], scope_name="conv5")  # (b, 58, 49, 256)
      conv6 = self.conv2d_relu(conv5, filter_shape=[3, 3, 256, 256], scope_name="conv6")  # (b, 58, 49, 256)
      pool3 = self.maxpool2d(conv6, scope_name="pool3")

      conv7 = self.conv2d_relu(pool3, filter_shape=[3, 3, 256, 512], scope_name="conv7")  # (b, 116, 98, 128)
      conv8 = self.conv2d_relu(conv7, filter_shape=[3, 3, 512, 512], scope_name="conv8")  # (b, 116, 98, 128)
      pool4 = self.maxpool2d(conv8, scope_name="pool4")

      conv9 = self.conv2d_relu(pool4, filter_shape=[3, 3, 512, 1024], scope_name="conv9")  # (b, 116, 98, 128)
      conv10 = self.conv2d_relu(conv9, filter_shape=[3, 3, 1024, 1024], scope_name="conv10")  # (b, 116, 98, 128)

      cache = (conv8, conv6, conv4, conv2)
      out = conv10
    return (out, cache)

class DeconvDecoderDeepUNet(NeuralNetwork):
  def __init__(self, keep_prob, output_shape, scope_name="decoder_unet_deep"):
    self.keep_prob = keep_prob
    self.output_shape = output_shape
    self.scope_name = scope_name

  def build_graph(self, input, cache):
    conv8, conv6, conv4, conv2 = cache
    with tf.variable_scope(self.scope_name):

      up1 = self.upsample(input, scope_name="up1", factor=[2, 2])  # (b, 116, 98, 256)
      deconv1 = self.deconv2d(up1, filter_shape=[2, 2], num_outputs=512, scope_name="deconv1")  # (b, 116, 98, 128)

      concat1 = tf.concat([conv8, deconv1], axis=3)  # (b, 116, 98, 256)
      conv11 = self.conv2d_relu(concat1, filter_shape=[3, 3, 1024, 1024], scope_name="conv11")  # (b, 116, 98, 128)
      conv12 = self.conv2d_relu(conv11, filter_shape=[3, 3, 1024, 512], scope_name="conv12")  # (b, 116, 98, 128)

      up2 = self.upsample(conv12, scope_name="up2", factor=[2, 2])  # (b, 116, 98, 256)
      deconv2 = self.deconv2d(up2, filter_shape=[2, 2], num_outputs=256, scope_name="deconv2")  # (b, 116, 98, 128)
      concat2 = tf.concat([conv6, deconv2], axis=3)  # (b, 116, 98, 256)
      conv13 = self.conv2d_relu(concat2, filter_shape=[3, 3, 512, 512], scope_name="conv13")  # (b, 116, 98, 128)
      conv14 = self.conv2d_relu(conv13, filter_shape=[3, 3, 512, 256], scope_name="conv14")  # (b, 116, 98, 128)

      up3 = self.upsample(conv14, scope_name="up3", factor=[2, 2])  # (b, 116, 98, 256)
      deconv3 = self.deconv2d(up3, filter_shape=[2, 2], num_outputs=128, scope_name="deconv3")  # (b, 116, 98, 128)
      concat3 = tf.concat([conv4, deconv3], axis=3)  # (b, 116, 98, 256)
      conv15 = self.conv2d_relu(concat3, filter_shape=[3, 3, 256, 256], scope_name="conv15")
      conv16 = self.conv2d_relu(conv15, filter_shape=[3, 3, 256, 128], scope_name="conv16")

      up4 = self.upsample(conv16, scope_name="up4", factor=[2, 2])  # (b, 116, 98, 256)
      deconv4 = self.deconv2d(up4, filter_shape=[2, 2], num_outputs=64, scope_name="deconv4")  # (b, 116, 98, 128)
      concat4 = tf.concat([conv2, deconv4], axis=3)  # (b, 116, 98, 256)
      conv17 = self.conv2d_relu(concat4, filter_shape=[3, 3, 128, 128], scope_name="conv17")
      conv18 = self.conv2d_relu(conv17, filter_shape=[3, 3, 128, 64], scope_name="conv18")

      conv19 = self.conv2d_relu(conv18, filter_shape=[1, 1, 64, 1], scope_name="conv19")

      unpadded_out = tf.identity(conv19, name="out")
      out = unpadded_out[:, 4:-4, 6:-6, :]

    return out

class ConvEncoderShallowUNet(NeuralNetwork):
  def __init__(self, input_shape, keep_prob, scope_name="encoder_unet_shallow"):
    self.input_shape = input_shape
    self.keep_prob = keep_prob
    self.scope_name = scope_name

  def build_graph(self, input):
    with tf.variable_scope(self.scope_name):
      # 233 x 196
      paddings = tf.constant([[0, 0], [0, 0], [2, 2], [0, 0]])
      # 233 x 200
      padded_input = tf.pad(input, paddings, "CONSTANT")  
      conv1 = self.conv2d_relu(padded_input, filter_shape=[3, 3, 1, 64], scope_name="conv1")  # (230, 194, 64)
      conv2 = self.conv2d_relu(conv1, filter_shape=[3, 3, 64, 64], scope_name="conv2")  # (228, 192, 64)
      pool1 = self.maxpool2d(conv2, scope_name="pool1") # (114, 96, 64)

      conv3 = self.conv2d_relu(pool1, filter_shape=[3, 3, 64, 128], scope_name="conv3")  # (b, 112, 94, 128)
      conv4 = self.conv2d_relu(conv3, filter_shape=[3, 3, 128, 128], scope_name="conv4")  # (b, 110, 92, 128)
      pool2 = self.maxpool2d(conv4, scope_name="pool2") # (55, 46, 128)

      conv5 = self.conv2d_relu(pool2, filter_shape=[3, 3, 128, 256], scope_name="conv5")  # (b, 58, 49, 256)
      conv6 = self.conv2d_relu(conv5, filter_shape=[3, 3, 256, 256], scope_name="conv6")  # (b, 58, 49, 256)
      pool2 = self.maxpool2d(conv4, scope_name="pool2") # (55, 46, 128)

      cache = (conv6, conv4, conv2)
      out = conv6

    return (out, cache)

class DeconvDecoderShallowUNet(NeuralNetwork):
  def __init__(self, keep_prob, output_shape, scope_name="decoder_unet_shallow"):
    self.keep_prob = keep_prob
    self.output_shape = output_shape
    self.scope_name = scope_name

  def build_graph(self, input, cache):
    conv6, conv4, conv2 = cache
    with tf.variable_scope(self.scope_name):

      up2 = self.upsample(input, scope_name="up2", factor=[2, 2])  # (b, 116, 98, 256)
      deconv2 = self.deconv2d(up2, filter_shape=[2, 2], num_outputs=128, scope_name="deconv2")  # (b, 116, 98, 128)
      concat2 = tf.concat([conv4, deconv2], axis=3)  # (b, 116, 98, 256)
      conv11 = self.conv2d_relu(concat2, filter_shape=[3, 3, 256, 256], scope_name="conv11")
      conv12 = self.conv2d_relu(conv11, filter_shape=[3, 3, 256, 128], scope_name="conv12")

      up3 = self.upsample(conv12, scope_name="up3", factor=[2, 2])  # (b, 116, 98, 256)
      deconv3 = self.deconv2d(up3, filter_shape=[2, 2], num_outputs=64, scope_name="deconv3")  # (b, 116, 98, 128)
      concat3 = tf.concat([conv2, deconv3], axis=3)  # (b, 116, 98, 256)
      conv13 = self.conv2d_relu(concat3, filter_shape=[3, 3, 128, 128], scope_name="conv13")
      conv14 = self.conv2d_relu(conv13, filter_shape=[3, 3, 128, 64], scope_name="conv14")
      
      conv15 = self.conv2d_relu(conv14, filter_shape=[1, 1, 64, 1], scope_name="conv15")

      unpadded_out = tf.identity(conv15, name="out")
      out = unpadded_out[:, :, 2:-2, :]

    return out
