# python code describing CNN architecture and training/validating/testing process

import tensorflow as tf
 
# dimension of input data
x = tf.placeholder(tf.float32, [None, training_image_size], name='x')

# number of categories
y = tf.placeholder(tf.float32, [None, num_labels], name='y')

# learning rate
learn_rate = tf.placeholder(tf.float32, name='learn_rate')

# drop out rate
keep_prob = tf.placeholder(tf.float32, name='keep_prob')

# input dimension to the CNN
images = tf.reshape(x, [-1, training_image_width, training_image_width, 1])

# now we set up the network architecture

# first layer of filters
conv1 = tf.layers.conv2d(inputs=images, filters=p.conv1_number_of_filters, kernel_size=[p.kernel_size, p.kernel_size], padding="same", activation=tf.nn.relu)
pool1 = tf.layers.max_pooling2d(conv1, pool_size=[p.pool_size, p.pool_size], strides=p.strides)

# second layer of filters
conv2 = tf.layers.conv2d(inputs=pool1, filters=p.conv2_number_of_filters, kernel_size=[p.kernel_size, p.kernel_size], padding="same", activation=tf.nn.relu)
pool2 = tf.layers.max_pooling2d(conv2, pool_size=[p.pool_size, p.pool_size], strides=p.strides)

# flatten and use fully connected layer
flat1 = tf.contrib.layers.flatten(pool2)
nn = tf.layers.dense(inputs=flat1, units=p.flat_layer_size, activation=tf.nn.relu)
nn = tf.layers.batch_normalization(nn)
nn = tf.nn.dropout(nn, keep_prob)

# output layer and optimisation processes
logits = tf.layers.dense(nn, p.number_of_categories, kernel_initializer=tf.contrib.layers.xavier_initializer(), name="logits")
cost = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(logits=logits, labels=y))
optimizer = tf.train.AdamOptimizer(learn_rate).minimize(cost)
correct_pred = tf.equal(tf.argmax(logits, 1), tf.argmax(y, 1))
accuracy = tf.reduce_mean(tf.cast(correct_pred, tf.float32), name='accuracy')
pred = tf.argmax(logits, 1)

# option to load a previous session, and/or test on new data
saver = tf.train.Saver(max_to_keep=p.number_to_save)

def run():
    
    with tf.Session() as sess:
        sess.run(tf.global_variables_initializer())
    
        # load checkpoint?
        if (not p.checkpoint == None):            
            saver.restore(sess, tf.train.latest_checkpoint(p.checkpoint))
        
        # only test then exit?
        if (p.test_only == True):                        
            predictions_test = pred.eval({x: x_test, keep_prob: 1.0})
            probabilities_test = logits.eval({x: x_test, keep_prob: 1.0})
            # now save x_test, y_test, predictions_test, and probabilities_test        
            return
        
        for e in range(epochs + 1):
    
            # get the next set of training images
            for x_batch, y_batch in get_batch(x_train, y_train, batch_size): # "get_batch()" returns training data
            
                # train the network on training images
                sess.run(optimizer, feed_dict={x: x_batch, y: y_batch, keep_prob: p.keep_probability, learn_rate:p.lr})
    
                # test the accuracy of the network on validation images, and display
                accuracy_val = accuracy.eval({x: x_val, y: y_val, keep_prob: 1.0})
                print(accuracy_val)
    
            # for example, after every iteration save network weighting
            # now save x_val, y_val, predictions_val, and probabilities_val             
            saver.save(sess, os.path.join(str(p.fname), "model"), global_step=e)
                
    return    
    
run()