#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Model exaplaination https://www.coursera.org/lecture/advanced-computer-vision-with-tensorflow/fcn-architecture-details-uwVDj
@author: anelmusic
"""

import tensorflow as tf
import config

vgg_weights_path = config.VGG_WEIGHTS_PATH

def block(x, n_convs, filters, kernel_size, activation, pool_size, pool_stride, block_name):
  for i in range(n_convs):
      x = tf.keras.layers.Conv2D(filters=filters, kernel_size=kernel_size, activation=activation, padding='same', name="{}_conv{}".format(block_name, i + 1))(x)
    
  x = tf.keras.layers.MaxPooling2D(pool_size=pool_size, strides=pool_stride, name="{}_pool{}".format(block_name, i+1 ))(x)
  return x

def VGG_16(image_input):
  x = block(image_input,n_convs=2, filters=64, kernel_size=(3,3), activation='relu',pool_size=(2,2), pool_stride=(2,2), block_name='block1')
  p1= x
  
  x = block(x,n_convs=2, filters=128, kernel_size=(3,3), activation='relu',pool_size=(2,2), pool_stride=(2,2), block_name='block2')
  p2 = x
 
  x = block(x,n_convs=3, filters=256, kernel_size=(3,3), activation='relu',pool_size=(2,2), pool_stride=(2,2), block_name='block3')
  p3 = x
  
  x = block(x,n_convs=3, filters=512, kernel_size=(3,3), activation='relu',pool_size=(2,2), pool_stride=(2,2), block_name='block4')
  p4 = x
 
  x = block(x,n_convs=3, filters=512, kernel_size=(3,3), activation='relu',pool_size=(2,2), pool_stride=(2,2), block_name='block5')
  p5 = x

  vgg  = tf.keras.Model(image_input , p5)
  vgg.load_weights(vgg_weights_path) 

  # number of filters for the output convolutional layers
  n = 4096

  c6 = tf.keras.layers.Conv2D( n , ( 7 , 7 ) , activation='relu' , padding='same', name="conv6")(p5)
  c7 = tf.keras.layers.Conv2D( n , ( 1 , 1 ) , activation='relu' , padding='same', name="conv7")(c6)
  return (p1, p2, p3, p4, c7)


def fcn8_decoder(convs, n_classes):
  f1, f2, f3, f4, f5 = convs
  o = tf.keras.layers.Conv2DTranspose(n_classes , kernel_size=(4,4) ,  strides=(2,2) , use_bias=False )(f5)
  o = tf.keras.layers.Cropping2D(cropping=(1,1))(o)
  
  o2 = f4
  o2 = ( tf.keras.layers.Conv2D(n_classes , ( 1 , 1 ) , activation='relu' , padding='same'))(o2)
  
  o = tf.keras.layers.Add()([o, o2])
  o = (tf.keras.layers.Conv2DTranspose( n_classes , kernel_size=(4,4) ,  strides=(2,2) , use_bias=False ))(o)
  o = tf.keras.layers.Cropping2D(cropping=(1, 1))(o)
  
  o2 = f3
  o2 = ( tf.keras.layers.Conv2D(n_classes , ( 1 , 1 ) , activation='relu' , padding='same'))(o2)
  
  o = tf.keras.layers.Add()([o, o2])
  o = tf.keras.layers.Conv2DTranspose(n_classes , kernel_size=(8,8) ,  strides=(8,8) , use_bias=False )(o)

  o = (tf.keras.layers.Activation('softmax'))(o)
  return o


def segmentation_model(): 
  inputs = tf.keras.layers.Input(shape=(config.NUM_PIXELS_SQRT,config.NUM_PIXELS_SQRT,3,))
  convs = VGG_16(image_input=inputs)
  outputs = fcn8_decoder(convs, config.OUTPUT_CHANNELS)
  model = tf.keras.Model(inputs=inputs, outputs=outputs)
  return model
tensorflow.keras.layers.Activation tensorflow.keras.layers.Conv2DTranspose tensorflow.keras.layers.Conv2D tensorflow.keras.Model tensorflow.keras.layers.Cropping2D tensorflow.keras.layers.Add tensorflow.keras.layers.Input
[(36, 'tensorflow.keras.Model', 'tf.keras.Model', (['image_input', 'p5'], {}), True, 'import tensorflow as tf\n'), (70, 'tensorflow.keras.layers.Input', 'tf.keras.layers.Input', ([], {'shape': '(config.NUM_PIXELS_SQRT, config.NUM_PIXELS_SQRT, 3)'}), True, 'import tensorflow as tf\n'), (73, 'tensorflow.keras.Model', 'tf.keras.Model', ([], {'inputs': 'inputs', 'outputs': 'outputs'}), True, 'import tensorflow as tf\n'), (42, 'tensorflow.keras.layers.Conv2D', 'tf.keras.layers.Conv2D', (['n', '(7, 7)'], {'activation': '"""relu"""', 'padding': '"""same"""', 'name': '"""conv6"""'}), True, 'import tensorflow as tf\n'), (43, 'tensorflow.keras.layers.Conv2D', 'tf.keras.layers.Conv2D', (['n', '(1, 1)'], {'activation': '"""relu"""', 'padding': '"""same"""', 'name': '"""conv7"""'}), True, 'import tensorflow as tf\n'), (49, 'tensorflow.keras.layers.Conv2DTranspose', 'tf.keras.layers.Conv2DTranspose', (['n_classes'], {'kernel_size': '(4, 4)', 'strides': '(2, 2)', 'use_bias': '(False)'}), True, 'import tensorflow as tf\n'), (50, 'tensorflow.keras.layers.Cropping2D', 'tf.keras.layers.Cropping2D', ([], {'cropping': '(1, 1)'}), True, 'import tensorflow as tf\n'), (53, 'tensorflow.keras.layers.Conv2D', 'tf.keras.layers.Conv2D', (['n_classes', '(1, 1)'], {'activation': '"""relu"""', 'padding': '"""same"""'}), True, 'import tensorflow as tf\n'), (55, 'tensorflow.keras.layers.Add', 'tf.keras.layers.Add', ([], {}), True, 'import tensorflow as tf\n'), (56, 'tensorflow.keras.layers.Conv2DTranspose', 'tf.keras.layers.Conv2DTranspose', (['n_classes'], {'kernel_size': '(4, 4)', 'strides': '(2, 2)', 'use_bias': '(False)'}), True, 'import tensorflow as tf\n'), (57, 'tensorflow.keras.layers.Cropping2D', 'tf.keras.layers.Cropping2D', ([], {'cropping': '(1, 1)'}), True, 'import tensorflow as tf\n'), (60, 'tensorflow.keras.layers.Conv2D', 'tf.keras.layers.Conv2D', (['n_classes', '(1, 1)'], {'activation': '"""relu"""', 'padding': '"""same"""'}), True, 'import tensorflow as tf\n'), (62, 'tensorflow.keras.layers.Add', 'tf.keras.layers.Add', ([], {}), True, 'import tensorflow as tf\n'), (63, 'tensorflow.keras.layers.Conv2DTranspose', 'tf.keras.layers.Conv2DTranspose', (['n_classes'], {'kernel_size': '(8, 8)', 'strides': '(8, 8)', 'use_bias': '(False)'}), True, 'import tensorflow as tf\n'), (65, 'tensorflow.keras.layers.Activation', 'tf.keras.layers.Activation', (['"""softmax"""'], {}), True, 'import tensorflow as tf\n')]