# -*- coding: utf-8 -*-
import tensorflow as tf
from tensorflow.examples.tutorials.mnist import input_data
from tensorflow.contrib.tensorboard.plugins import projector


old_v = tf.logging.get_verbosity()
tf.logging.set_verbosity(tf.logging.ERROR)
# 载入数据集
mnist = input_data.read_data_sets('MNIST_data', one_hot=True)
# 运行次数
max_steps = 3001
# 图片数量
image_num = 5000
# 文件路径
DIR = "D:/AIdata/tf_data/tf_test1/"

sess = tf.Session()

# 载入图片，
# tf.stack矩阵拼接函数，
embedding = tf.Variable(tf.stack(mnist.test.images[:image_num]),
                        trainable=False, name="embedding")


def variable_summaries(var):
    with tf.name_scope("summaries"):
        mean = tf.reduce_mean(var)
        with tf.name_scope("stddev"):
            # 计算标准差
            stddev = tf.sqrt(tf.reduce_mean(tf.square(var-mean)))
        # 绘制标准差信息
        tf.summary.scalar("stddev", stddev)
        # 绘制最大值
        tf.summary.scalar("max", tf.reduce_max(var))
        tf.summary.scalar("min", tf.reduce_min(var))
        # 绘制直方图信息
        tf.summary.histogram("histogram", var)


with tf.name_scope('Input'):
    # 这里的none表示第一个维度可以是任意的长度
    x = tf.placeholder(tf.float32, [None, 784], name="x_input")
    y = tf.placeholder(tf.float32, [None, 10], name="y_input")
    LR = tf.Variable(0.001, dtype=tf.float32)

# 显示图片
with tf.name_scope("input_reshape"):
    # 改变x的形状（28x28x1）
    image_shape_input = tf.reshape(x, [-1, 28, 28, 1])
    # 将图像写入summary，输出带图像的probuf
    tf.summary.image("Input", image_shape_input, 10)

with tf.name_scope('layer'):
    with tf.name_scope('weights'):
        W = tf.Variable(tf.zeros([784, 10]), name='W')
        variable_summaries(W)
    with tf.name_scope('biases'):
        b = tf.Variable(tf.zeros([10]), name='b')
        variable_summaries(b)
    with tf.name_scope('wxb'):
        wxb = tf.matmul(x, W) + b
    with tf.name_scope('softmax'):
        prediction = tf.nn.softmax(wxb)


with tf.name_scope("loss"):
    # 交叉熵函数
    loss = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits_v2(labels=y,
                                                                     logits=prediction))
    # 绘制loss值
    tf.summary.scalar("loss", loss)

with tf.name_scope("Train"):
    # AdamOptimizer优化器
    train_step = tf.train.AdamOptimizer(LR).minimize(loss)

init_op = tf.global_variables_initializer()
sess.run(init_op)   # 变量初始化


with tf.name_scope("Result"):
    with tf.name_scope("correct_prediction"):
        # 记录预测值和标签值对比结果
        correct_prediction = tf.equal(tf.argmax(y, 1), tf.argmax(prediction, 1))
    with tf.name_scope("Accuracy"):
        # 求准确率
        accuracy = tf.reduce_mean(tf.cast(correct_prediction, tf.float32))
        # 绘制准确率
        tf.summary.scalar("accuracy", accuracy)

# 判断是否已存在metadata.tsv文件，若存在则删除
if tf.gfile.Exists(DIR+"projector/projector/metadata.tsv"):
    tf.gfile.Remove(DIR+"projector/projector/metadata.tsv")

# 创建并写入metadata.tsv文件
with open(DIR+"projector/projector/metadata.tsv", 'w') as f:
    labels = sess.run(tf.argmax(mnist.test.labels[:], 1))
    for i in range(image_num):
        f.write(str(labels[i]) + '\n')

# 合并默认图表管理summary
merged = tf.summary.merge_all()

projector_writer = tf.summary.FileWriter(DIR+"/projector/projector", sess.graph)
# 定义saver对象，以保存和恢复模型变量
saver = tf.train.Saver()
# 定义配置
config = projector.ProjectorConfig()
embed = config.embeddings.add()
embed.tensor_name = embedding.name
# metadata_path文件路径
embed.metadata_path = DIR+"projector/projector/metadata.tsv"
# sprite image文件路径
embed.sprite.image_path = DIR+'projector/data/mnist_10k_sprite.png'
# sprite image中每一单个图像的大小
embed.sprite.single_image_dim.extend([28, 28])
# 写入可视化配置
projector.visualize_embeddings(projector_writer, config)

for i in range(max_steps):
    # 每个批次100个样本
    batch_xs, batch_ys = mnist.train.next_batch(100)
    run_options = tf.RunOptions(trace_level=tf.RunOptions.FULL_TRACE)
    run_metadata = tf.RunMetadata()
    summary, _ = sess.run([merged, train_step], feed_dict={x: batch_xs, y: batch_ys},
                          options=run_options, run_metadata=run_metadata)
    projector_writer.add_run_metadata(run_metadata, 'step%03d' % i)
    projector_writer.add_summary(summary, i)

    if i % 100 == 0:
        sess.run(tf.assign(LR, 0.001))
        acc = sess.run(accuracy, feed_dict={x: mnist.test.images, y: mnist.test.labels})
        print("Iter " + str(i) + ", Testing Accuracy= " + str(acc))
# 保存模型
saver.save(sess, DIR+'projector/projector/mnist_model.ckpt', global_step=max_steps)
projector_writer.close()
sess.close()