# Authors:
# Akshay Chaudhari and Zhongnan Fang
# May 2018
# akshaysc@stanford.edu

from __future__ import print_function, division

import numpy as np
import pickle
import json
import math
import os

from keras.models import Model, Sequential
from keras.layers import Input, Conv2D, MaxPooling2D, Conv2DTranspose, concatenate, add, Lambda, Dropout, AlphaDropout
from keras.layers import BatchNormalization as BN
from keras.utils import plot_model
from keras import backend as K
import tensorflow as tf

def unet_2d_model(input_size):

    # input size is a tuple of the size of the image
    # assuming channel last
    # input_size = (dim1, dim2, dim3, ch)
    # unet begins

    nfeatures = [2**feat*32 for feat in np.arange(6)]
    depth = len(nfeatures)    

    conv_ptr = []

    # input layer
    inputs = Input(input_size)

    # step down convolutional layers 
    pool = inputs
    for depth_cnt in xrange(depth):

        conv = Conv2D(nfeatures[depth_cnt], (3,3), 
                      padding='same', 
                      activation='relu',
                      kernel_initializer='he_normal')(pool)
        conv = Conv2D(nfeatures[depth_cnt], (3,3), 
                      padding='same', 
                      activation='relu',
                      kernel_initializer='he_normal')(conv)

        conv = BN(axis=-1, momentum=0.95, epsilon=0.001)(conv)
        conv = Dropout(rate=0.0)(conv)

        conv_ptr.append(conv)

        # Only maxpool till penultimate depth
        if depth_cnt < depth-1:

            # If size of input is odd, only do a 3x3 max pool
            xres = conv.shape.as_list()[1]
            if (xres % 2 == 0):
                pooling_size = (2,2)
            elif (xres % 2 == 1):
                pooling_size = (3,3)

            pool = MaxPooling2D(pool_size=pooling_size)(conv)


    # step up convolutional layers
    for depth_cnt in xrange(depth-2,-1,-1):

        deconv_shape = conv_ptr[depth_cnt].shape.as_list()
        deconv_shape[0] = None

        # If size of input is odd, then do a 3x3 deconv  
        if (deconv_shape[1] % 2 == 0):
            unpooling_size = (2,2)
        elif (deconv_shape[1] % 2 == 1):
            unpooling_size = (3,3)

        up = concatenate([Conv2DTranspose(nfeatures[depth_cnt],(3,3),
                          padding='same',
                          strides=unpooling_size,
                          output_shape=deconv_shape)(conv),
                          conv_ptr[depth_cnt]], 
                          axis=3)

        conv = Conv2D(nfeatures[depth_cnt], (3,3), 
                      padding='same', 
                      activation='relu',
                      kernel_initializer='he_normal')(up)
        conv = Conv2D(nfeatures[depth_cnt], (3,3), 
                      padding='same', 
                      activation='relu',
                      kernel_initializer='he_normal')(conv)

        conv = BN(axis=-1, momentum=0.95, epsilon=0.001)(conv)
        conv = Dropout(rate=0.00)(conv)

    # combine features
    recon = Conv2D(1, (1,1), padding='same', activation='sigmoid')(conv)

    model = Model(inputs=[inputs], outputs=[recon])
    plot_model(model, to_file='unet2d.png',show_shapes=True)
    
    return model


if __name__ == '__main__':
  
    os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
    img_size = (288,288,1)
    model = unet_2d_model(img_size)
    print(model.summary())