"""
.............................................................
Számjegyfelismerés konvolúciós hálózattal (MNIST, TensorFlow)
.............................................................
"""

print(__doc__)

import tensorflow as tf
import numpy as np

from tensorflow.examples.tutorials.mnist import input_data              # adatok betöltése
mnist = input_data.read_data_sets('MNIST_data', one_hot=True)

x = tf.placeholder(tf.float32, shape=[None,784])            # input images placeholder
d = tf.placeholder(tf.float32, shape=[None,10])             # labels placeholder

def weight(shape):                              # normális eloszlásúra inicializálja 
    return tf.Variable(tf.random_normal(shape, stddev=0.1))

def bias(shape):                                # kis pozitív torzítás
    return tf.Variable(tf.constant(0.1, shape=shape))

def conv2d(x, W):                                        # zero-padded: ugyanakkora az eredmény, output size = input size (0-kkal kipótolja a kép szélét)
    return tf.nn.conv2d(x, W, strides=[1, 1, 1, 1], padding='SAME')

def max_pool(x):                                         # 2x2-es max pooling, stride: 2 -> nincs átfedés a kis négyzetek között
    return tf.nn.max_pool(x, ksize=[1, 2, 2, 1], strides=[1, 2, 2, 1], padding='SAME')

# first conv layer

W_conv1 = weight([5, 5, 1, 32])               # 32 features
b_conv1 = bias([32])

x_image = tf.reshape(x, [-1, 28, 28, 1])

conv1 = tf.nn.relu(conv2d(x_image, W_conv1) + b_conv1)
pool1 = max_pool(conv1)

# second conv layer

W_conv2 = weight([5, 5, 32, 64])              # 64 features
b_conv2 = bias([64])

conv2 = tf.nn.relu(conv2d(pool1, W_conv2) + b_conv2)
pool2 = max_pool(conv2)

# fully-connected (vagy dense) layer

W_fc1 = weight([7 * 7 * 64, 1024])            # fully-connected layerben 1024 neuron
b_fc1 = bias([1024])

pool2_flat = tf.reshape(pool2, [-1, 7*7*64])
fc1 = tf.nn.relu(tf.matmul(pool2_flat, W_fc1) + b_fc1)

# dropout layer

keep_prob = tf.placeholder(tf.float32)
fc1_drop = tf.nn.dropout(fc1, keep_prob)

# output layer

W_fc2 = weight([1024, 10])                   # 1024 neuron, output: 10-elemű vektor (one-hot vectors)
b_fc2 = bias([10])

y = tf.matmul(fc1_drop, W_fc2) + b_fc2

loss = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(labels=d, logits=y))
optimizer = tf.train.AdamOptimizer(0.0001).minimize(loss)
correct_prediction = tf.equal(tf.argmax(y, 1), tf.argmax(d, 1))
accuracy = tf.reduce_mean(tf.cast(correct_prediction, tf.float32))

with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())
    for _ in range(500):                              # 500 iteráció
        batch = mnist.train.next_batch(50)            # batchsize: 50
        optimizer.run(feed_dict={x: batch[0], d: batch[1], keep_prob: 0.5})        # tanítás futtatása: loss és gradients kiszámolása, súlyok update-elése
    print('test accuracy %g' % accuracy.eval(feed_dict={x: mnist.test.images, d: mnist.test.labels, keep_prob: 1.0}))   # pontosság a teszthalmazon

##################################################

# Batch training: vagy a fenti Session kell, vagy ez:

with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())
    for i in range(500):                          
        batch = mnist.train.next_batch(50)
        if i % 100 == 0:                          # ha i osztható 100-zal..., azaz minden 100. iterációnál kiírja, hogy áll a tanítás
            train_accuracy = accuracy.eval(feed_dict={x: batch[0], d: batch[1], keep_prob: 0.5})
            print('step %d, training accuracy %g' % (i, train_accuracy))
        optimizer.run(feed_dict={x: batch[0], d: batch[1], keep_prob: 0.5})
    print('test accuracy %g' % accuracy.eval(feed_dict={x: mnist.test.images, d: mnist.test.labels, keep_prob: 1.0}))