import numpy as np

import tensorflow as tf

import pickle

from tensorflow.examples.tutorials.mnist import input_data

import matplotlib.pyplot as plt

# 读入mnist输入

mnist = input_data.read_data_sets('/data')

# 构造输入样本的数据,real_size 和 noise_size, 生成初始化的real_image 和 noise_image

def get_inputs(real_size, noise_size):

    # 真实图片的维度为[None, 784]

    real_img = tf.placeholder('float32', [None, real_size])

    # 噪音数据的维度为[None , 100]

    noise_img = tf.placeholder('float32', [None, noise_size])

    return real_img, noise_img

# 构造generator网络，将[None, 100]的噪音数据, 生成[None, 28, 28, 1]的图片数据

def get_generator(noise_img, output_dim, is_train, alpha=0.01):

    with tf.variable_scope('generator', reuse=not is_train):

        # 100 -> 4*4*512 - > 4, 4, 512

        # 进行全连接操作，将100的数据通过全连接，转换为[None ,4*4*512]

        layer1 = tf.layers.dense(noise_img, 4*4*512)

        # 进行维度变化，将其转换为[-1, 4, 4, 512]

        layer1 = tf.reshape(layer1, [-1, 4, 4, 512])

        # 使用tf.layers.batch_normalization进行归一化操作

        layer1 = tf.layers.batch_normalization(layer1, training=is_train)

        # 使用tf.maximum进行lrelu操作

        layer1 = tf.maximum(alpha*layer1, layer1)

        # 使用tf.nn.dropout进行dropout操作

        layer1 = tf.nn.dropout(layer1, keep_prob=0.8)

        # 4*4*512 -> 7*7*256

        # 进行反卷积操作，将4*4*256 -> 7*7*256

        layer2 = tf.layers.conv2d_transpose(layer1, 256, 4, 1, padding='valid')

        # 进行归一化操作

        layer2 = tf.layers.batch_normalization(layer2, training=is_train)

        # 进行lrelu操作

        layer2 = tf.maximum(alpha*layer2, layer2)

        # 进行dropout操作

        layer2 = tf.nn.dropout(layer2, keep_prob=0.8)

        #  7*7*256 - 14*14*128

        # 进行反卷积操作，stride=2，维度扩大为原来的2倍

        layer3 = tf.layers.conv2d_transpose(layer2, 128, 3, 2, padding='SAME')

        # 归一化操作

        layer3 = tf.layers.batch_normalization(layer3, training=is_train)

        # 进行lrelu操作

        layer3 = tf.maximum(alpha*layer3, layer3)

        # 进行dropout的操作

        layer3 = tf.nn.dropout(layer3, keep_prob=0.8)

        # 进行反卷积操作将 14*14*128 -> 28*28*1

        logits = tf.layers.conv2d_transpose(layer3, output_dim, 3, 2, padding='SAME')

        # 将输出结果转换为-1, 1之间

        outputs = tf.tanh(logits)

    return outputs

# 生成判别网络，输入样本的大小为[28, 28, 1], reuse表示是否进行参数复用

def get_discrimator(input_image, reuse, alpha=0.01):

    # 指定参数范围

    with tf.variable_scope('discrimator', reuse=reuse):

        # 28*28*1 -> 14*14*128

        # 进行卷积操作，将28*28*1变化为14*14*128，步长为2

        layer1 = tf.layers.conv2d(input_image, 128, 2, padding='SAME')

        # 进行lrelu操作

        layer1 = tf.maximum(layer1*alpha, layer1)

        # 进行dropout操作

        layer1 = tf.nn.dropout(layer1, keep_prob=0.8)

        # 14*14*128 -> 7*7*256

        # 进行卷积操作，将14*14*128变换为7*7*256, 步长为2

        layer2 = tf.layers.conv2d(layer1, 256, 2, padding='SAME')

        # 使用tf.layers.batch_normalization进行归一化操作

        layer2 = tf.layers.batch_normalization(layer2, training=True)

        # 进行lrelu操作

        layer2 = tf.maximum(layer2*alpha, layer2)

        # 进行dropout操作

        layer2 = tf.nn.dropout(layer2, keep_prob=0.8)

        # 7*7*64 -> 4*4*512

        # 进行卷积操作

        layer3 = tf.layers.conv2d(layer2, 512, 2, padding='SAME')

        # 进行归一化操作

        layer3 = tf.layers.batch_normalization(layer3, training=True)

        # 进行lrelu操作

        layer3 = tf.maximum(layer3*alpha, layer3)

        # 进行dropout操作

        layer3 = tf.nn.dropout(layer3, keep_prob=0.8)\

        # 进行维度变换

        flatten_fc = tf.reshape(layer3, [-1, 4*4*512])

        # 进行全连接操作，将4*4*512 变为1

        logits = tf.layers.dense(flatten_fc, 1)

        # 将结果变化为0, 1之间，输出的是概率值

        outputs = tf.sigmoid(logits)

    return logits, outputs

# 用于生成损失值,输入为真实图片，噪音图片, 生成图片的维度， smooth

def get_loss(real_image, noise_image, output_dim, smooth=0.1):

    # 生成图片和判别网络

    # 使用生成网络，输入噪音图片，生成图片

    g_outputs = get_generator(noise_image, output_dim, True)

    # 使用判别网络，对真实图片进行判别，返回Logits和inputs

    d_real_logits, d_real_outputs = get_discrimator(real_image, False)

    # 使用判别网络，对生成图片进行判别，返回为logits和inputs

    d_fake_logits, d_real_outputs = get_discrimator(g_outputs, True)

    # 生成loss函数， 使用tf.reduce_mean构成d_loss

    d_loss_1 = tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(logits=d_real_logits, labels=tf.ones_like(d_real_logits) * (1-smooth)))

    d_loss_2 = tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(logits=d_fake_logits, labels=tf.zeros_like(d_fake_logits)))

    d_loss = d_loss_1 + d_loss_2

    # 使用tf.reduce_mean 构成g_loss, labels=tf.ones_like(d_fake_logits)

    g_loss = tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(logits=d_fake_logits, labels=tf.ones_like(d_fake_logits) * (1-smooth)))

    # 返回判别和生成的损失值

    return d_loss, g_loss

# 获得用于降低损失的g_opt 和 d_opt

def get_opt(d_loss, g_loss, learning_rate=0.001):

    # 获得训练过程中的所有训练参数

    varies = tf.trainable_variables()

    # 循环训练参数，根据参数名字的开头获得生成网络的参数

    g_var = [var for var in varies if var.name.startswith('generator')]

    # 根据参数名字的开头获得判别网络的参数

    d_var = [var for var in varies if var.name.startswith('discrimator')]

    # 对于d_loss的损失值，只训练d_var参数

    d_opt = tf.train.AdamOptimizer(learning_rate).minimize(d_loss, var_list=d_var)

    # 对于g_loss的损失值，只训练g_var参数

    g_opt = tf.train.AdamOptimizer(learning_rate).minimize(g_loss, var_list=g_var)

    return d_opt, g_opt

# 调用get_generator生成图片

def show_sample_image(sess, n_sample, input_noise, output_dim):

    # 获得输入图片的维度，这里值噪声图片,维度为100

    input_size = input_noise.get_shape().as_list()[-1]

    # 生成-1,1的[50, 100]的噪音数据

    noise_sample = np.random.uniform(-1, 1, (n_sample, input_size))

    # 执行生成函数，将噪音数据传入

    samples = sess.run(get_generator(input_noise, output_dim, False), feed_dict={input_noise:noise_sample})

    return samples

# 进行画图操作

def plot_samples(samples):

    # 表示图片以及图片的维度

    fig, axes = plt.subplots(nrows=1, ncols=25)

    # 循环每一个subplots进行画图操作

    for img, ax in zip(samples, axes):

        # 在每个subplot上进行画图

        ax.imshow(img, cmap='Greys_r')

        # 去除ax的边界

        ax.get_xaxis().set_visible(False)

        ax.get_yaxis().set_visible(False)

    # 图片没有边界

    fig.tight_layer(pad=0)

    # 图片展示

    plt.show()

# 主函数进行训练

def train():

    # 第一步：设置基本的参数

    # 图片的大小，这里是784

    input_size = mnist.train.images[0].shape[0]

    # 噪声数据的大小，100

    noise_size = 100

    # 每一个batch值

    batch_size = 64

    # 生成图片的维度

    output_dim = 1

    # 迭代的epoch值

    epoch = 30

    # 用于展示的图片大小

    n_sample = 25

    # 第二步：调用get_inputs,生成初始化的输入数据,[None, 784], [None, 100]

    real_img, noise_img = get_inputs(input_size, noise_size)

    # 对输入数据的真实数据进行维度变化，用来输入到判别网络进行判别

    real_img_reshape = tf.reshape(real_img, [-1, 28, 28, 1])

    # 第三步: 将真实数据和噪音数据输入到loss函数中，生成判别损失值d_loss和生成损失值g_loss

    d_loss, g_loss = get_loss(real_img_reshape, noise_img, output_dim)

    # 第四步：嗲用get_op根据损失值获得操作步骤opt

    d_opt, g_opt = get_opt(d_loss, g_loss)

    # 第五步：实际传入数据，进行模型的训练

    with tf.Session() as sess:

        # 参数初始化

        sess.run(tf.global_variables_initializer())

        # 循环epoch

        for i in range(epoch):

            # 迭代每一个batch_size

            for batch_i in range(mnist.train.num_examples // batch_size):

                # 获得一个batch_size的数据

                batch = mnist.train.next_batch(batch_size)

                # 获得一个batch_size的图片

                batch_img = batch[0]

                # 将其范围变为-1, 1

                batch_img = batch_img * 2 - 1

                # 获得一个batch_size的噪音数据

                batch_noise = np.random.uniform(-1, 1, (batch_size, noise_size))

                # 实际传入参数，执行g_opt 和 d_opt

                _ = sess.run(g_opt, feed_dict={real_img:batch_img, noise_img:batch_noise})

                _ = sess.run(d_opt, feed_dict={real_img:batch_img, noise_img:batch_noise})

            # 一个epoch后获的损失值

            train_d_loss = sess.run(d_loss, feed_dict={real_img:batch_img, noise_img:batch_noise})

            train_g_loss = sess.run(g_loss, feed_dict={real_img: batch_img, noise_img: batch_noise})

            # 迭代10次以后

            if i % 10 == 0:

                # 根据噪音数据生成samples数据

                samples = show_sample_image(sess, n_sample, noise_img, output_dim)

                # 对生成的samples数据进行图像展示

                plot_samples(samples)

train()