"""Inception-v3 dnnlib.Network port."""

import tensorflow as tf

# Initialize dict which maps parameter names to TF variables.
_inception_params_dict = dict()
scope_prefix = 'InceptionV3/import'

#----------------------------------------------------------------------------
# Get variables and add names to params dict.

def get_parameters(shape, name):
    w = tf.get_variable(name, shape=shape, initializer=tf.initializers.zeros())
    _inception_params_dict[scope_prefix + '/' + name] = w
    return w

# Softmax.

def softmax(x, fmaps, weights_name):
    w = get_parameters([x.shape[1].value,fmaps], weights_name)
    w = tf.cast(w, x.dtype)
    out = tf.matmul(x, w)
    out = tf.exp(out)
    out /= tf.reduce_sum(out, axis=1, keepdims=True)
    return out

# Max-pooling.

def maxpool2d(x, stride=2, padding='VALID'):
    return tf.nn.max_pool(x, ksize=[1,3,3,1], strides=[1,stride,stride,1], padding=padding, data_format='NHWC')

# Avg-pooling.

def avgpool2d(x, kernel_size=3, padding='SAME'):
    return tf.nn.avg_pool(x, ksize=[1,kernel_size,kernel_size,1], strides=[1,1,1,1], padding=padding, data_format='NHWC')

# Conv. layer without bias.

def conv2d(x, kernel, fmaps, weights_name, stride=1, padding='VALID'):
    w = get_parameters([kernel,kernel,x.shape[3].value,fmaps], weights_name)
    return tf.nn.conv2d(x, w, strides=[1,stride,stride,1], padding=padding, data_format='NHWC')

# Factorized conv. layer without bias.

def factorized_conv2d(x, kernel_height, kernel_width, fmaps, weights_name, stride=1, padding='VALID'):
    w = get_parameters([kernel_height,kernel_width,x.shape[3].value,fmaps], weights_name)

# Batch-normalization.

def batch_normalization(x, beta_name, avg_name, var_name):
    shape = [x.shape[3].value]
    beta = get_parameters(shape, beta_name)
    avg_mean = get_parameters(shape, avg_name)
    avg_var = get_parameters(shape, var_name)

    # Note: variance_epsilon must be 1e-3 to match legacy network.
    return tf.nn.batch_normalization(x, mean=avg_mean, variance=avg_var,
                                     offset=beta, scale=None, variance_epsilon=1e-3)

# Inception module, type A.

def inception_module_A(x, fmaps_after_pool=64):
    # Conv. tower 1: 1x1->3x3->3x3.
    with tf.variable_scope('tower_1'):
        # 1x1 conv.
        conv1_tower1 = conv2d(x, 1, 64, 'conv/conv2d_params', padding='SAME')
        bn1_tower1 = batch_normalization(conv1_tower1, 'conv/batchnorm/beta', 'conv/batchnorm/moving_mean', 'conv/batchnorm/moving_variance')
        relu1_tower1 = tf.nn.relu(bn1_tower1)

        # 3x3 conv.
        conv2_tower1 = conv2d(relu1_tower1, 3, 96, 'conv_1/conv2d_params', padding='SAME')
        bn2_tower1 = batch_normalization(conv2_tower1, 'conv_1/batchnorm/beta', 'conv_1/batchnorm/moving_mean', 'conv_1/batchnorm/moving_variance')
        relu2_tower1 = tf.nn.relu(bn2_tower1)

        # 3x3 conv.
        conv3_tower1 = conv2d(relu2_tower1, 3, 96, 'conv_2/conv2d_params', padding='SAME')
        bn3_tower1 = batch_normalization(conv3_tower1, 'conv_2/batchnorm/beta', 'conv_2/batchnorm/moving_mean', 'conv_2/batchnorm/moving_variance')
        relu3_tower1 = tf.nn.relu(bn3_tower1)

    # Conv. tower 2: 1x1->5x5.
    with tf.variable_scope('tower'):
        # 1x1 conv.
        conv1_tower2 = conv2d(x, 1, 48, 'conv/conv2d_params', padding='SAME')
        bn1_tower2 = batch_normalization(conv1_tower2, 'conv/batchnorm/beta', 'conv/batchnorm/moving_mean', 'conv/batchnorm/moving_variance')
        relu1_tower2 = tf.nn.relu(bn1_tower2)

        # 3x3 conv.
        conv2_tower2 = conv2d(relu1_tower2, 5, 64, 'conv_1/conv2d_params', padding='SAME')
        bn2_tower2 = batch_normalization(conv2_tower2, 'conv_1/batchnorm/beta', 'conv_1/batchnorm/moving_mean', 'conv_1/batchnorm/moving_variance')
        relu2_tower2 = tf.nn.relu(bn2_tower2)

    # Conv. tower 3: Avg-pool->1x1.
    with tf.variable_scope('tower_2'):
        # Pool.
        pool1_out = avgpool2d(x)
        conv1_tower3 = conv2d(pool1_out, 1, fmaps_after_pool, 'conv/conv2d_params', padding='SAME')
        bn1_tower3 = batch_normalization(conv1_tower3, 'conv/batchnorm/beta', 'conv/batchnorm/moving_mean', 'conv/batchnorm/moving_variance')
        relu1_tower3 = tf.nn.relu(bn1_tower3)

    # Conv. tower 4: 1x1.
    with tf.variable_scope('conv'):
        conv1_tower4 = conv2d(x, 1, 64, 'conv2d_params', padding='SAME')
        bn1_tower4 = batch_normalization(conv1_tower4, 'batchnorm/beta', 'batchnorm/moving_mean', 'batchnorm/moving_variance')
        relu1_tower4 = tf.nn.relu(bn1_tower4)

    # Concatenate channels.
    module_output = tf.concat([relu1_tower4,relu2_tower2,relu3_tower1,relu1_tower3], axis=3)

    return module_output

# Inception module, type B.

def inception_module_B(x):
        # 1x1 conv.

        # 3x3 conv. stride 1.

        # 3x3 conv. stride 2.
        conv3_tower1 = conv2d(relu2_tower1, 3, 96, 'conv_2/conv2d_params', padding='VALID', stride=2)

    # Conv. tower 2: 3x3.
        conv1_tower2 = conv2d(x, 3, 384, 'conv2d_params', padding='VALID', stride=2)
        bn1_tower2 = batch_normalization(conv1_tower2, 'batchnorm/beta', 'batchnorm/moving_mean', 'batchnorm/moving_variance')

    # Max pooling.
    with tf.variable_scope('pool'):
        pool_out = maxpool2d(x)

    module_output = tf.concat([relu1_tower2, relu3_tower1, pool_out], axis=3)

# Inception module, type C.

def inception_module_C(x, tower_fmaps=128, tower_1_fmaps=160):
    # Conv. tower 1: 1x1->7x1->1x7->7x1->1x7.
        # 1x1 conv.
        conv1_tower1 = conv2d(x, 1, tower_1_fmaps, 'conv/conv2d_params', padding='SAME')

        # 7x1 conv.
        conv2_tower1 = factorized_conv2d(relu1_tower1, 7, 1, tower_1_fmaps, 'conv_1/conv2d_params', stride=1, padding='SAME')

        # 1x7 conv.
        conv3_tower1 = factorized_conv2d(relu2_tower1, 1, 7, tower_1_fmaps, 'conv_2/conv2d_params', stride=1, padding='SAME')

        # 7x1 conv.
        conv4_tower1 = factorized_conv2d(relu3_tower1, 7, 1, tower_1_fmaps, 'conv_3/conv2d_params', stride=1, padding='SAME')
        bn4_tower1 = batch_normalization(conv4_tower1, 'conv_3/batchnorm/beta', 'conv_3/batchnorm/moving_mean', 'conv_3/batchnorm/moving_variance')
        relu4_tower1 = tf.nn.relu(bn4_tower1)

        # 1x7 conv.
        conv5_tower1 = factorized_conv2d(relu4_tower1, 1, 7, 192, 'conv_4/conv2d_params', stride=1, padding='SAME')
        bn5_tower1 = batch_normalization(conv5_tower1, 'conv_4/batchnorm/beta', 'conv_4/batchnorm/moving_mean', 'conv_4/batchnorm/moving_variance')
        relu5_tower1 = tf.nn.relu(bn5_tower1)
    
    # Conv. tower 2: 1x1->1x7->7x1.
        # 1x1 conv.
        conv1_tower2 = conv2d(x, 1, tower_fmaps, 'conv/conv2d_params', padding='SAME')

        # 1x7 conv.
        conv2_tower2 = factorized_conv2d(relu1_tower2, 1, 7, tower_fmaps, 'conv_1/conv2d_params', padding='SAME')

        # 7x1 conv.
        conv3_tower2 = factorized_conv2d(relu2_tower2, 7, 1, 192, 'conv_2/conv2d_params', padding='SAME')
        bn3_tower2 = batch_normalization(conv3_tower2, 'conv_2/batchnorm/beta', 'conv_2/batchnorm/moving_mean', 'conv_2/batchnorm/moving_variance')
        relu3_tower2 = tf.nn.relu(bn3_tower2)

    # Conv. tower 3: Avg.pool->1x1.
        # Avg. pool.
        pool_out = avgpool2d(x)

        # 1x1 conv.
        conv1_tower3 = conv2d(pool_out, 1, 192, 'conv/conv2d_params', padding='SAME')

    # Conv. tower4: 1x1.
        conv1_tower4 = conv2d(x, 1, 192, 'conv2d_params', padding='SAME')

    module_output = tf.concat([relu1_tower4, relu3_tower2, relu5_tower1, relu1_tower3], axis=3)

# Inception module, type D.

def inception_module_D(x):
    # Conv. tower 1: 1x1->1x7->7x1->3x3.
        # 1x1 conv.
        conv1_tower1 = conv2d(x, 1, 192, 'conv/conv2d_params', padding='SAME')

        # 1x7 conv.
        conv2_tower1 = factorized_conv2d(relu1_tower1, 1, 7, 192, 'conv_1/conv2d_params', padding='SAME')

        # 7x1 conv.
        conv3_tower1 = factorized_conv2d(relu2_tower1, 7, 1, 192, 'conv_2/conv2d_params', padding='SAME')

        # 3x3 conv.
        conv4_tower1 = conv2d(relu3_tower1, 3, 192, 'conv_3/conv2d_params', padding='VALID', stride=2)

    # Conv. tower 2: 1x1->3x3.
        # 1x1 conv.
        conv1_tower2 = conv2d(x, 1, 192, 'conv/conv2d_params', padding='SAME')

        # 3x3 conv.
        conv2_tower2 = conv2d(relu1_tower2, 3, 320, 'conv_1/conv2d_params', padding='VALID', stride=2)

    # Max-pool.

    module_output = tf.concat([relu2_tower2, relu4_tower1, pool_out], axis=3)

# Inception module, type E.

def inception_module_E(x, use_avg_pool=True):
    # Conv. tower 1: 1x1.
        conv1_tower1 = conv2d(x, 1, 320, 'conv2d_params', padding='SAME')
        bn1_tower1 = batch_normalization(conv1_tower1, 'batchnorm/beta', 'batchnorm/moving_mean', 'batchnorm/moving_variance')

    # Conv. tower 2: 1x1->1x3->3x1.
        # 1x1 conv.
        conv1_tower2 = conv2d(x, 1, 384, 'conv/conv2d_params', padding='SAME')

        # 1x3 conv.
        conv2_tower2 = factorized_conv2d(relu1_tower2, 1, 3, 384, 'mixed/conv/conv2d_params', padding='SAME')
        bn2_tower2 = batch_normalization(conv2_tower2, 'mixed/conv/batchnorm/beta', 'mixed/conv/batchnorm/moving_mean', 'mixed/conv/batchnorm/moving_variance')
        relu2_tower2 = tf.nn.relu(bn2_tower2)    

        # 3x1 conv.
        conv3_tower2 = factorized_conv2d(relu1_tower2, 3, 1, 384, 'mixed/conv_1/conv2d_params', padding='SAME')
        bn3_tower2 = batch_normalization(conv3_tower2, 'mixed/conv_1/batchnorm/beta', 'mixed/conv_1/batchnorm/moving_mean', 'mixed/conv_1/batchnorm/moving_variance')
        relu3_tower2 = tf.nn.relu(bn3_tower2) 

    # Conv. tower 3: 1x1->3x3->1x3->3x1.
        # 1x1 conv.
        conv1_tower3 = conv2d(x, 1, 448, 'conv/conv2d_params', padding='SAME')

        # 3x3 conv.
        conv2_tower3 = conv2d(relu1_tower3, 3, 384, 'conv_1/conv2d_params', padding='SAME')
        bn2_tower3 = batch_normalization(conv2_tower3, 'conv_1/batchnorm/beta', 'conv_1/batchnorm/moving_mean', 'conv_1/batchnorm/moving_variance')
        relu2_tower3 = tf.nn.relu(bn2_tower3)

        # 1x3 conv.
        conv3_tower3 = factorized_conv2d(relu2_tower3, 1, 3, 384, 'mixed/conv/conv2d_params', padding='SAME')
        bn3_tower3 = batch_normalization(conv3_tower3, 'mixed/conv/batchnorm/beta', 'mixed/conv/batchnorm/moving_mean', 'mixed/conv/batchnorm/moving_variance')
        relu3_tower3 = tf.nn.relu(bn3_tower3)    

        # 3x1 conv.
        conv4_tower3 = factorized_conv2d(relu2_tower3, 3, 1, 384, 'mixed/conv_1/conv2d_params', padding='SAME')
        bn4_tower3 = batch_normalization(conv4_tower3, 'mixed/conv_1/batchnorm/beta', 'mixed/conv_1/batchnorm/moving_mean', 'mixed/conv_1/batchnorm/moving_variance')
        relu4_tower3 = tf.nn.relu(bn4_tower3) 
    
    # Conv. tower 4: Avg.pool->1x1.
        if use_avg_pool:
            pool_out = avgpool2d(x)
        else:
            pool_out = maxpool2d(x, stride=1, padding='SAME')

        # 1x1 conv.
        conv1_tower4 = conv2d(pool_out, 1, 192, 'conv/conv2d_params', padding='SAME')
        bn1_tower4 = batch_normalization(conv1_tower4, 'conv/batchnorm/beta', 'conv/batchnorm/moving_mean', 'conv/batchnorm/moving_variance')
    
    module_output = tf.concat([relu1_tower1, relu2_tower2, relu3_tower2, relu3_tower3, relu4_tower3, relu1_tower4], axis=3)

def inception_v3(images_in, include_softmax=True, **_kwargs):
    """Inception-v3 builder function."""
    # Input NCHW, [0,255].
    images_in.set_shape([None, 3, None, None])
    x = tf.cast(images_in, tf.float32)

    # Convert to NHWC, 299x299, [-1,1].
    x = tf.transpose(x, [0, 2, 3, 1])
    x = tf.image.resize_bilinear(x, [299, 299])
    x -= 128.0
    x *= 1 / 128.0

    #----------------------------------------------------------------------------

    with tf.variable_scope(scope_prefix + '/conv'):
        conv0_out = conv2d(x, 3, 32, 'conv2d_params', stride=2)
        bn0_out = batch_normalization(conv0_out, 'batchnorm/beta', 'batchnorm/moving_mean', 'batchnorm/moving_variance')
        relu0_out = tf.nn.relu(bn0_out)

    with tf.variable_scope(scope_prefix + '/conv_1'):
        conv1_out = conv2d(relu0_out, 3, 32, 'conv2d_params')
        bn1_out = batch_normalization(conv1_out, 'batchnorm/beta', 'batchnorm/moving_mean', 'batchnorm/moving_variance')
        relu1_out = tf.nn.relu(bn1_out)

    with tf.variable_scope(scope_prefix + '/conv_2'):
        conv2_out = conv2d(relu1_out, 3, 64, 'conv2d_params', padding='SAME')
        bn2_out = batch_normalization(conv2_out, 'batchnorm/beta', 'batchnorm/moving_mean', 'batchnorm/moving_variance')
        relu2_out = tf.nn.relu(bn2_out)
   
    with tf.variable_scope(scope_prefix + '/pool0'):
        pool0_out = maxpool2d(relu2_out)

    with tf.variable_scope(scope_prefix + '/conv_3'):
        conv3_out = conv2d(pool0_out, 1, 80, 'conv2d_params')
        bn3_out = batch_normalization(conv3_out, 'batchnorm/beta', 'batchnorm/moving_mean', 'batchnorm/moving_variance')
        relu3_out = tf.nn.relu(bn3_out)

    with tf.variable_scope(scope_prefix + '/conv_4'):
        conv4_out = conv2d(relu3_out, 3, 192, 'conv2d_params')
        bn4_out = batch_normalization(conv4_out, 'batchnorm/beta', 'batchnorm/moving_mean', 'batchnorm/moving_variance')
        relu4_out = tf.nn.relu(bn4_out)

    with tf.variable_scope(scope_prefix + '/pool1'):
        pool1_out = maxpool2d(relu4_out)

    # 3x type A Inception modules.
    with tf.variable_scope(scope_prefix + '/mixed'):
        module_A1_out = inception_module_A(pool1_out, fmaps_after_pool=32)

    with tf.variable_scope(scope_prefix + '/mixed_1'):
        module_A2_out = inception_module_A(module_A1_out)

    with tf.variable_scope(scope_prefix + '/mixed_2'):
        module_A3_out = inception_module_A(module_A2_out)

    # 1x type B Inception module.
    with tf.variable_scope(scope_prefix + '/mixed_3'):
        module_B1_out = inception_module_B(module_A3_out)

    # 4x type C Inception modules.
    with tf.variable_scope(scope_prefix + '/mixed_4'):
        module_C1_out = inception_module_C(module_B1_out, tower_fmaps=128, tower_1_fmaps=128)

    with tf.variable_scope(scope_prefix + '/mixed_5'):
        module_C2_out = inception_module_C(module_C1_out, tower_fmaps=160, tower_1_fmaps=160)

    with tf.variable_scope(scope_prefix + '/mixed_6'):
        module_C3_out = inception_module_C(module_C2_out, tower_fmaps=160, tower_1_fmaps=160)

    with tf.variable_scope(scope_prefix + '/mixed_7'):
        module_C4_out = inception_module_C(module_C3_out, tower_fmaps=192, tower_1_fmaps=192)

    # 1x type D Inception module.
    with tf.variable_scope(scope_prefix + '/mixed_8'):
        module_D1_out = inception_module_D(module_C4_out)

    # 2x type E Inception modules.
    with tf.variable_scope(scope_prefix + '/mixed_9'):
        module_E1_out = inception_module_E(module_D1_out, use_avg_pool=True)

    with tf.variable_scope(scope_prefix + '/mixed_10'):
        module_E2_out = inception_module_E(module_E1_out, use_avg_pool=False)

    # Pool_3.
    with tf.variable_scope(scope_prefix + '/pool_3'):
        pool3_out = avgpool2d(module_E2_out, kernel_size=8, padding='VALID')

    logits = tf.reshape(pool3_out, [-1,pool3_out.shape[3].value])

    assert logits.dtype == tf.float32

    if include_softmax:
        class_probabilities = softmax(logits, 1008, scope_prefix + '/softmax/weights')
        return class_probabilities
    else:
        return logits