import numpy as np
import tensorflow as tf

import awesome_gans.modules as t

tf.set_random_seed(777)  # reproducibility
np.random.seed(777)  # reproducibility

he_normal = tf.contrib.layers.variance_scaling_initializer(factor=1.0, mode='FAN_AVG', uniform=True)
l2_reg = tf.contrib.layers.l2_regularizer


def pixel_norm(x, eps=1e-8):
    return x * tf.rsqrt(tf.reduce_mean(tf.square(x), axis=1, keepdims=True) + eps)


def resize_nn(x, size):
    return tf.image.resize_nearest_neighbor(x, size=(int(size), int(size)))


def bacth_concat(x, eps=1e-8, averaging='all'):
    """
    ref : https://github.com/zhangqianhui/progressive_growing_of_gans_tensorflow/blob/master/ops.py#L145
    """
    adj_std = lambda x_, **kwargs: tf.sqrt(tf.reduce_mean((x_ - tf.reduce_mean(x_, **kwargs)) ** 2, **kwargs) + eps)

    val_ = adj_std(x, axis=0, keepdims=True)
    if averaging == 'all':
        val_ = tf.reduce_mean(val_, keepdims=True)
    val_ = tf.tile(val_, multiples=[tf.shape(x)[0], 4, 4, 1])
    return tf.concat([x, val_], axis=3)


class PGGAN:
    def __init__(
        self,
        s,
        batch_size=16,
        input_height=128,
        input_width=128,
        input_channel=3,