import os
os.environ['TF_CPP_MIN_LOG_LEVEL']='2'
import tensorflow as tf
import numpy as np
import tf_utils

def model_l5(X, params, is_training):
    regulizer = tf.contrib.layers.l2_regularizer(params["weight_decay"]) if params["weight_decay"] > 0. else None

    bn_params = None
    if "bn_decay" in params.keys() and "bn_epsilon" in params.keys():
        bn_params = {'decay': params["bn_decay"], 'epsilon': params["bn_epsilon"], 'is_training': is_training}

    conv1_1 = tf_utils.conv2d(X, "conv1_1", 3, 64,
                              bn_params = bn_params, use_bias = (bn_params == None),
                              weight_regularizer = regulizer,
                              bias_regularizer = regulizer)
    conv1_2 = tf_utils.conv2d(conv1_1, "conv1_2", 3, 64,
                              bn_params = bn_params, use_bias = (bn_params == None),
                              weight_regularizer = regulizer,
                              bias_regularizer = regulizer)
    pool1 = tf.nn.max_pool(conv1_2, ksize = [1, 2, 2, 1], strides = [1, 2, 2, 1], padding = "SAME", name = "pool1")
    #############################################################################
    conv2_1 = tf_utils.conv2d(pool1, "conv2_1", 3, 128,
                              bn_params = bn_params, use_bias = (bn_params == None),
                              weight_regularizer = regulizer,
                              bias_regularizer = regulizer)
    conv2_2 = tf_utils.conv2d(conv2_1, "conv2_2", 3, 128,
                              bn_params = bn_params, use_bias = (bn_params == None),
                              weight_regularizer = regulizer,
                              bias_regularizer = regulizer)
    pool2 = tf.nn.max_pool(conv2_2, ksize = [1, 2, 2, 1], strides = [1, 2, 2, 1], padding = "SAME", name = "pool2")
    #############################################################################

    conv3_1 = tf_utils.conv2d(pool2, "conv3_1", 3, 256,
                              bn_params = bn_params, use_bias = (bn_params == None),
                              weight_regularizer = regulizer,
                              bias_regularizer = regulizer)
    conv3_2 = tf_utils.conv2d(conv3_1, "conv3_2", 3, 256,
                              bn_params = bn_params, use_bias = (bn_params == None),
                              weight_regularizer = regulizer,
                              bias_regularizer = regulizer)
    pool3 = tf.nn.max_pool(conv3_2, ksize = [1, 2, 2, 1], strides = [1, 2, 2, 1], padding = "SAME", name = "pool3")
    #############################################################################
    conv4_1 = tf_utils.conv2d(pool3, "conv4_1", 3, 512,
                              bn_params = bn_params, use_bias = (bn_params == None),
                              weight_regularizer = regulizer,
                              bias_regularizer = regulizer)
    conv4_2 = tf_utils.conv2d(conv4_1, "conv4_2", 3, 512,
                              bn_params = bn_params, use_bias = (bn_params == None),
                              weight_regularizer = regulizer,
                              bias_regularizer = regulizer)
    pool4 = tf.nn.max_pool(conv4_2, ksize = [1, 2, 2, 1], strides = [1, 2, 2, 1], padding = "SAME", name = "pool4")
    #############################################################################
    conv5_1 = tf_utils.conv2d(pool4, "conv5_1", 3, 1024,
                              bn_params = bn_params, use_bias = (bn_params == None),
                              weight_regularizer = regulizer,
                              bias_regularizer = regulizer)
    conv5_2 = tf_utils.conv2d(conv5_1, "conv5_2", 3, 1024,
                              bn_params = bn_params, use_bias = (bn_params == None),
                              weight_regularizer = regulizer,
                              bias_regularizer = regulizer)
    #############################################################################
    upconv5 = tf_utils.conv2d_transpose(conv5_2, "upconv5",
                                        kernel_sz = 2, out_shape = tf.shape(conv4_2), strides = 2, out_chns = 512,
                                        bn_params = None,
                                        weight_regularizer = regulizer,
                                        bias_regularizer = regulizer)
    concat4 = tf.concat([conv4_2, upconv5], axis = -1)
    deconv4_2 = tf_utils.conv2d(concat4, "deconv4_2", 3, 512,
                                bn_params = bn_params, use_bias = (bn_params == None),
                                weight_regularizer = regulizer,
                                bias_regularizer = regulizer)
    deconv4_1 = tf_utils.conv2d(deconv4_2, "deconv4_1", 3, 512,
                                bn_params = bn_params, use_bias = (bn_params == None),
                                weight_regularizer = regulizer,
                                bias_regularizer = regulizer)
    #############################################################################
    upconv4 = tf_utils.conv2d_transpose(deconv4_1, "upconv4",
                                        kernel_sz = 2, out_shape = tf.shape(conv3_2), strides = 2, out_chns = 256,
                                        bn_params = None,
                                        weight_regularizer = regulizer,
                                        bias_regularizer = regulizer)
    concat3 = tf.concat([conv3_2, upconv4], axis = -1)
    deconv3_2 = tf_utils.conv2d(concat3, "deconv3_2", 3, 256,
                                bn_params = bn_params, use_bias = (bn_params == None),
                                weight_regularizer = regulizer,
                                bias_regularizer = regulizer)
    deconv3_1 = tf_utils.conv2d(deconv3_2, "deconv3_1", 3, 256,
                                bn_params = bn_params, use_bias = (bn_params == None),
                                weight_regularizer = regulizer,
                                bias_regularizer = regulizer)
    #############################################################################
    upconv3 = tf_utils.conv2d_transpose(deconv3_1, "upconv3",
                                        kernel_sz = 2, out_shape = tf.shape(conv2_2), strides = 2, out_chns = 128,
                                        bn_params = None,
                                        weight_regularizer = regulizer,
                                        bias_regularizer = regulizer)
    concat2 = tf.concat([conv2_2, upconv3], axis = -1)
    deconv2_2 = tf_utils.conv2d(concat2, "deconv2_2", 3, 128,
                                bn_params = bn_params, use_bias = (bn_params == None),
                                weight_regularizer = regulizer,
                                bias_regularizer = regulizer)
    deconv2_1 = tf_utils.conv2d(deconv2_2, "deconv2_1", 3, 128,
                                bn_params = bn_params, use_bias = (bn_params == None),
                                weight_regularizer = regulizer,
                                bias_regularizer = regulizer)
    #############################################################################
    upconv2 = tf_utils.conv2d_transpose(deconv2_1, "upconv2",
                                        kernel_sz = 2, out_shape = tf.shape(conv1_2), strides = 2, out_chns = 64,
                                        bn_params = None,
                                        weight_regularizer = regulizer,
                                        bias_regularizer = regulizer)
    concat2 = tf.concat([conv1_2, upconv2], axis = -1)
    deconv1_2 = tf_utils.conv2d(concat2, "deconv1_2", 3, 64,
                                bn_params = bn_params, use_bias = (bn_params == None),
                                weight_regularizer = regulizer,
                                bias_regularizer = regulizer)
    deconv1_1 = tf_utils.conv2d(deconv1_2, "deconv1_1", 3, 64,
                                bn_params = bn_params, use_bias = (bn_params == None),
                                weight_regularizer = regulizer,
                                bias_regularizer = regulizer)
    #############################################################################
    result = tf_utils.conv2d(deconv1_1, "result_deconv", 1, params["classes_cnt"],
                             bn_params = None, use_bias = False, non_linear_func = None,
                             weight_regularizer = regulizer,
                             bias_regularizer = regulizer)
    return result

