import tensorflow as tf
import os

batch_size = 32
img_size=64
num_channels = 3;
num_classes = 41;

## labels
img_data = tf.placeholder(tf.float32, shape=[None, img_size, img_size, num_channels], name='img_data')
img_labels = tf.placeholder(tf.float32, shape=[None, num_classes], name='img_labels')
img_labels_cls = tf.argmax(img_labels, dimension=1)

##Network graph params
filter_size_conv1 = 3
num_filters_conv1 = 32

filter_size_conv2 = 3
num_filters_conv2 = 32

filter_size_conv3 = 3
num_filters_conv3 = 64

fc_layer_size = 128

def create_weights(shape):
    return tf.Variable(tf.truncated_normal(shape,stddev=0.05))

def create_biases(size):
    return tf.Variable(tf.constant(0.05, shape=[size]))

def create_convolutional_layer(input,
               num_input_channels,
               conv_filter_size,
               num_filters):

    weights = create_weights(shape=[conv_filter_size, conv_filter_size, num_input_channels, num_filters])
    biases = create_biases(num_filters)

    ## Add conv layer
    layer = tf.nn.conv2d(input=input,
                     filter=weights,
                     strides=[1, 1, 1, 1],
                     padding='SAME')

    layer += biases

    ## We shall be using max-pooling.
    layer = tf.nn.max_pool(value=layer,
                            ksize=[1, 2, 2, 1],
                            strides=[1, 2, 2, 1],
                            padding='SAME')
    ## Output of pooling is fed to Relu which is the activation function for us.
    layer = tf.nn.relu(layer)
    return layer

def create_flatten_layer(layer):
    layer_shape = layer.get_shape()
    num_features = layer_shape[1:4].num_elements()
    layer = tf.reshape(layer, [-1, num_features])
    return layer

def create_fc_layer(input,
             num_inputs,
             num_outputs,
             use_relu=True):

    #Let's define trainable weights and biases.
    weights = create_weights(shape=[num_inputs, num_outputs])
    biases = create_biases(num_outputs)

    layer = tf.matmul(input, weights) + biases
    if use_relu:
        layer = tf.nn.relu(layer)
    return layer

    layer_conv1 = create_convolutional_layer(input=img_data,
               num_input_channels=num_channels,
               conv_filter_size=filter_size_conv1,
               num_filters=num_filters_conv1)

layer_conv2 = create_convolutional_layer(input=layer_conv1,
               num_input_channels=num_filters_conv1,
               conv_filter_size=filter_size_conv2,
               num_filters=num_filters_conv2)

layer_conv3= create_convolutional_layer(input=layer_conv2,
               num_input_channels=num_filters_conv2,
               conv_filter_size=filter_size_conv3,
               num_filters=num_filters_conv3)

layer_flat = create_flatten_layer(layer_conv3)

layer_fc1 = create_fc_layer(input=layer_flat,
                     num_inputs=layer_flat.get_shape()[1:4].num_elements(),
                     num_outputs=fc_layer_size,
                     use_relu=True)

layer_fc2 = create_fc_layer(input=layer_fc1,
                     num_inputs=fc_layer_size,
                     num_outputs=num_classes,
                     use_relu=False)



cross_entropy = tf.nn.softmax_cross_entropy_with_logits(logits=layer_fc2,
                                                    labels=img_labels)
cost = tf.reduce_mean(cross_entropy)
pred_lab = tf.nn.softmax(layer_fc2,name='pred_lab')
pred_lab_cls = tf.argmax(pred_lab, dimension=1)

tr_op = tf.train.AdamOptimizer(learning_rate=1e-4).minimize(cost)

correct_prediction = tf.equal(pred_lab_cls, img_labels_cls)
accuracy = tf.reduce_mean(tf.cast(correct_prediction, tf.float32))



def read_and_decode_single_example(filename):
    # first construct a queue containing a list of filenames.
    # this lets a user split up there dataset in multiple files to keep
    # size down
    filename_queue = tf.train.string_input_producer([filename],
                                                    num_epochs=None)
    # Unlike the TFRecordWriter, the TFRecordReader is symbolic
    reader = tf.TFRecordReader()
    # One can read a single serialized example from a filename
    # serialized_example is a Tensor of type string.
    _, serialized_example = reader.read(filename_queue)
    # The serialized example is converted back to actual values.
    # One needs to describe the format of the objects to be returned
    features = tf.parse_single_example(
        serialized_example,
        features={
            'label':tf.FixedLenFeature([41], tf.float32),
            'image':tf.FixedLenFeature([12288], tf.float32)
        })
    # now return the converted data
    label = features['label']
    img = features['image']
    image = tf.reshape(img, [64,64,3])
    return label, image


# returns symbolic label and image
path = os.path.join(r'C:\Users\CalcBox\Documents\NeuralNet\Kaggle\Freesound',
                    'TfRecords\FreeSoundTrain64.tfrecords')
# get single examples
label, image = read_and_decode_single_example(path)

# groups examples into batches randomly
images_batch, labels_batch = tf.train.shuffle_batch(
    [image, label], batch_size=128,
    capacity=2000,
    min_after_dequeue=1000)

sess = tf.Session()
init = tf.initialize_all_variables()
sess.run(init)
tf.train.start_queue_runners(sess=sess)


labels, images= sess.run([labels_batch, images_batch])
feed_dict = {img_data: images,img_labels: labels,}
print(type(tr_op))
print(type(feed_dict))
try:
    some_object_iterator = iter(feed_dict)
except TypeError as te:
    print (feed_dict, 'is not iterable')
print(feed_dict)
_ , loss_val = sess.run(tr_op,
                        feed_dict=feed_dict)
print('Loss of the current batch is {}'.format(loss_val))