import tensorflow as tf

from tensorflow.contrib import slim

from nets import resnet_v1

from nets.convlstm_cell import ConvLSTMCell

from utils.training import get_valid_logits_and_labels

FLAGS = tf.app.flags.FLAGS

def unpool(inputs,scale):

    return tf.image.resize_bilinear(inputs, size=[tf.shape(inputs)[1]*scale,  tf.shape(inputs)[2]*scale])

def ResidualConvUnit(inputs,features=256,kernel_size=3):

    net=tf.nn.relu(inputs)

    net=slim.conv2d(net, features, kernel_size)

    net=tf.nn.relu(net)

    net=slim.conv2d(net,features,kernel_size)

    net=tf.add(net,inputs)

    return net

def ChainedResidualPooling(inputs,features=256):

    net_relu=tf.nn.relu(inputs)

    net=slim.max_pool2d(net_relu, [5, 5],stride=1,padding='SAME')

    net=slim.conv2d(net,features,3)

    net_sum_1=tf.add(net,net_relu)

    net = slim.max_pool2d(net_relu, [5, 5], stride=1, padding='SAME')

    net = slim.conv2d(net, features, 3)

    net_sum_2=tf.add(net,net_sum_1)

    return net_sum_2

def MultiResolutionFusion(high_inputs=None,low_inputs=None,features=256):

    if high_inputs is None:#refineNet block 4

        rcu_low_1 = low_inputs[0]

        rcu_low_2 = low_inputs[1]

        rcu_low_1 = slim.conv2d(rcu_low_1, features, 3)

        rcu_low_2 = slim.conv2d(rcu_low_2, features, 3)

        return tf.add(rcu_low_1,rcu_low_2)

    else:

        rcu_low_1 = low_inputs[0]

        rcu_low_2 = low_inputs[1]

        rcu_low_1 = slim.conv2d(rcu_low_1, features, 3)

        rcu_low_2 = slim.conv2d(rcu_low_2, features, 3)

        rcu_low = tf.add(rcu_low_1,rcu_low_2)

        rcu_high_1 = high_inputs[0]

        rcu_high_2 = high_inputs[1]

        rcu_high_1 = unpool(slim.conv2d(rcu_high_1, features, 3),2)

        rcu_high_2 = unpool(slim.conv2d(rcu_high_2, features, 3),2)

        rcu_high = tf.add(rcu_high_1,rcu_high_2)

        return tf.add(rcu_low, rcu_high)

def RefineBlock(high_inputs=None,low_inputs=None):

    if high_inputs is None: # block 4

        rcu_low_1= ResidualConvUnit(low_inputs, features=256)

        rcu_low_2 = ResidualConvUnit(low_inputs, features=256)

        rcu_low = [rcu_low_1, rcu_low_2]

        fuse = MultiResolutionFusion(high_inputs=None, low_inputs=rcu_low, features=256)

        fuse_pooling = ChainedResidualPooling(fuse, features=256)

        output = ResidualConvUnit(fuse_pooling, features=256)

        return output

    else:

        rcu_low_1 = ResidualConvUnit(low_inputs, features=256)

        rcu_low_2 = ResidualConvUnit(low_inputs, features=256)

        rcu_low = [rcu_low_1, rcu_low_2]

        rcu_high_1 = ResidualConvUnit(high_inputs, features=256)

        rcu_high_2 = ResidualConvUnit(high_inputs, features=256)

        rcu_high = [rcu_high_1, rcu_high_2]

        fuse = MultiResolutionFusion(rcu_high, rcu_low,features=256)

        fuse_pooling = ChainedResidualPooling(fuse, features=256)

        output = ResidualConvUnit(fuse_pooling, features=256)

        return output

def model(model_type, images, weight_decay=1e-5, is_training=True):

    images = mean_image_subtraction(images)

    with slim.arg_scope(resnet_v1.resnet_arg_scope(weight_decay=weight_decay)):

        logits, end_points = resnet_v1.resnet_v1_101(images, is_training=is_training, scope='resnet_v1_101')

    with tf.variable_scope('feature_fusion', values=[end_points.values]):

        batch_norm_params = {'decay': 0.997,'epsilon': 1e-5,'scale': True,'is_training': is_training}

        with slim.arg_scope([slim.conv2d],

                            activation_fn=tf.nn.relu,

                            normalizer_fn=slim.batch_norm,

                            normalizer_params=batch_norm_params,

                            weights_regularizer=slim.l2_regularizer(weight_decay)):

            f = [end_points['pool5'], end_points['pool4'],

                 end_points['pool3'], end_points['pool2']]

            for i in range(4):

                print('Shape of f_{} {}'.format(i, f[i].shape))

            g = [None, None, None, None]

            h = [None, None, None, None]

            for i in range(4):

                h[i]=slim.conv2d(f[i], 256, 1)

            for i in range(4):

                print('Shape of h_{} {}'.format(i, h[i].shape))

            g[0]=RefineBlock(high_inputs=None,low_inputs=h[0])

            g[1]=RefineBlock(g[0],h[1])

            g[2]=RefineBlock(g[1],h[2])

            g[3]=RefineBlock(g[2],h[3])

            #g[3]=unpool(g[3],scale=4)

            output = g[3]

            if model_type == 'sesnet':

                in_shape = g[3].shape

                output = tf.expand_dims(g[3], axis=0)

                lstm_cell_1 = ConvLSTMCell([in_shape[1], in_shape[2]], in_shape[3] // 2, [3, 3])

                lstm_cell_2 = ConvLSTMCell([in_shape[1], in_shape[2]], in_shape[3] // 4, [3, 3])

                with tf.variable_scope('rnn_scope_0', reuse=tf.AUTO_REUSE):

                    output0, _ = tf.nn.dynamic_rnn(lstm_cell_1, output, dtype=output.dtype)

                with tf.variable_scope('rnn_scope_1', reuse=tf.AUTO_REUSE):

                    output1, _ = tf.nn.dynamic_rnn(lstm_cell_2, output0, dtype=output0.dtype)

                output = tf.squeeze(output1, axis=0)

            F_score = slim.conv2d(output, 2, 1, activation_fn=tf.nn.relu, normalizer_fn=None)

    return F_score

def mean_image_subtraction(images, means=[123.68, 116.78, 103.94]):

    images=tf.to_float(images)

    num_channels = images.get_shape().as_list()[-1]

    if len(means) != num_channels:

        raise ValueError('len(means) must match the number of channels')

    channels = tf.split(axis=3, num_or_size_splits=num_channels, value=images)

    for i in range(num_channels):

        channels[i] -= means[i]

    return tf.concat(axis=3, values=channels)

def loss(annotation_batch,upsampled_logits_batch,class_labels):

    valid_labels_batch_tensor, valid_logits_batch_tensor = get_valid_logits_and_labels(

        annotation_batch_tensor=annotation_batch,

        logits_batch_tensor=upsampled_logits_batch,

        class_labels=class_labels)

    valid_labels_batch_tensor = tf.argmax(valid_labels_batch_tensor, dimension=3)

    cross_entropies = tf.nn.sparse_softmax_cross_entropy_with_logits(logits=valid_logits_batch_tensor,

                                                              labels=valid_labels_batch_tensor)

    cross_entropy_sum = tf.reduce_mean(cross_entropies)

    return cross_entropy_sum