﻿import tensorflow as tf
import time
import numpy as np

def dropout(src, keep_prop, isDrop):
    src = tf.cond(isDrop, lambda:tf.nn.dropout(src, keep_prop), lambda:identity(src))
    return src

def identity(src):
    return src

def resize(src, dstH, dstW,interpol = 1):
    
    if interpol == 0: return tf.image.resize_nearest_neighbor(src, [dstH, dstW])
    elif interpol == 1: return tf.image.resize_bilinear(src, [dstH, dstW])
    else: return tf.image.resize_bicubic(src, [dstH, dstW])

def conv2d(src, weights, bias):
    
    conv = tf.nn.conv2d(src,weights,strides=[1, 1, 1, 1],padding='SAME')
    return tf.nn.bias_add(conv,bias)


def conv2dRelu(src, weights, bias):
    conv = conv2d(src,weights,bias)
    return tf.nn.relu(conv) 

def conv2dReluStride2(src, weights, bias):
    
    conv = tf.nn.conv2d(src,weights,strides=[1, 2, 2, 1],padding='SAME')
    conv = tf.nn.bias_add(conv,bias)
    return tf.nn.relu(conv) 

def Gaussian_noise_layer(input_layer, std):
    noise = tf.random_normal(tf.shape(input_layer), mean = 0.0, stddev = std, dtype = tf.float32) 
    return input_layer + noise

def Gaussian_noise_volumn_layer(input_layer, std):
    input_shape = tf.shape(input_layer)
    shape = [input_shape[0],1,1,input_shape[3]]
    noise = tf.random_normal(shape , mean = 1.0, stddev = std, dtype = tf.float32) 
    return input_layer * noise

def Gaussian_noise(input_layer, stdAlpha, stdBeta):
    src = Gaussian_noise_volumn_layer(input_layer,stdAlpha)
    return Gaussian_noise_layer(src,stdBeta)   

def Gaussian_noise_Add(input_layer, stdAlpha, stdBeta,isTrain):
    return tf.cond(isTrain, lambda:Gaussian_noise(input_layer,stdAlpha,stdBeta), lambda:identity(input_layer))

def upConv(src, weights, bias,up_shape):
    output = tf.nn.conv2d_transpose(src,weights,output_shape=up_shape, strides=[1, 2, 2, 1],padding='SAME')
    output = tf.nn.bias_add(output,bias)  
    return output

def upConvRelu(src, weights, bias,up_shape):
    output = tf.nn.conv2d_transpose(src,weights,output_shape=up_shape, strides=[1, 2, 2, 1],padding='SAME')
    output = tf.nn.bias_add(output,bias)  
    return tf.nn.relu(output) 

def depthwiseConv2dRelu(src, weights, bias):
    
    output = tf.nn.depthwise_conv2d(src,weights,strides=[1,1,1,1],padding='SAME')
    output = tf.nn.bias_add(output, bias)
    output = tf.nn.relu(output)
    return output 

def LayerNormalize(src4d):
    src_shape = tf.shape(src4d)
    src_2d = tf.reshape(src4d, [src_shape[0],-1])
    m, v = tf.nn.moments(x=src_2d, axes=[1], keep_dims=True)
    src_normal = (src_2d - m)/tf.sqrt(v + 1e-5)
    src_4d = tf.reshape(src_normal, src_shape)
    return src_4d

def StdNormalize(src, beta,gamma):
    batch_mean, batch_var = tf.nn.moments(x=src,axes=[0,1,2])
    out = tf.nn.batch_normalization(x=src, mean=batch_mean, variance=batch_var,offset=beta,scale=gamma,variance_epsilon=1e-3)    
    return tf.nn.relu(out) 

def batchNormal(src, beta,gamma):
    batch_mean, batch_var = tf.nn.moments(x=src,axes=[0,1,2])
    out = tf.nn.batch_normalization(x=src, mean=batch_mean, variance=batch_var,offset=beta,scale=gamma,variance_epsilon=1e-3)    
    return out

def batch_norm(src, beta, gamma, isTrain):
    batch_mean, batch_var = tf.nn.moments(x=src,axes=[0,1,2])
    ema = tf.train.ExponentialMovingAverage(decay=0.5)

    def mean_var_with_update():
        ema_apply_op = ema.apply([batch_mean, batch_var])

        with tf.control_dependencies([ema_apply_op]):

            return tf.identity(batch_mean), tf.identity(batch_var)
    mean, var = tf.cond(isTrain, mean_var_with_update, lambda:(ema.average(batch_mean),ema.average(batch_var)))    
    return tf.nn.batch_normalization(src, mean,var, beta, gamma,1e-3)

def conv2dBN(src, weights, beta,gamma):
    
    conv = tf.nn.conv2d(src,weights,strides=[1, 1, 1, 1],padding='SAME')
    return batchNormal(conv,beta,gamma)

def conv2dBN_Relu(src, weights, beta,gamma):
    src_s = dropout(src)
    conv = tf.nn.conv2d(src_s,weights,strides=[1, 1, 1, 1],padding='SAME')
    bn = batchNormal(conv,beta,gamma)
    return tf.nn.relu(bn)

def max_pool_2(src):
    return tf.nn.max_pool(src,[1, 2, 2, 1],strides=[1, 2, 2, 1],padding='SAME')    

def max_pool_k3_s2(src):
    return tf.nn.max_pool(src,[1, 3, 3, 1],strides=[1, 2, 2, 1],padding='SAME')    

def bilinear_resize(src, step):
    #input_shape = tf.shape(src)
    input_shape = src.get_shape().as_list()
    srcH = 128
    srcW = 256
    scale = step % 3 + 1
    dstH = np.round(srcH/scale)
    dstW = np.round(srcW/scale)
    src= resize(src,dstH,dstW)
    return src

def reshape_4d_to_2d(src):
    shape = src.get_shape().as_list()
    dim = 1
    for d in shape[1:]:
        dim *= d
    return tf.reshape(src, [-1, dim])