"""
With this script you can finetune AlexNet as provided in the alexnet.py
class on any given dataset. 
Specify the configuration settings at the beginning according to your 
problem.
This script was written for TensorFlow 1.0 and come with a blog post 
you can find here:
  
https://kratzert.github.io/2017/02/24/finetuning-alexnet-with-tensorflow.html

Author: Frederik Kratzert 
contact: f.kratzert(at)gmail.com
"""
import os
import numpy as np
import tensorflow as tf
from datetime import datetime
from alexnet import AlexNet
from dataLoader import ImageDataGenerator
from testDataGen import getTestImage
from tensorflow.contrib.data import Iterator


"""
Configuration settings
"""
TRAINING=1;

# Path to the textfiles for the trainings and validation set
base_path = '/Users/apurvnigam/study_ucl/term1/MScThesis/hinterstoisser/test/05/dummy/'

train_file = './train.txt'
val_file = base_path+ 'validationData/validate.txt'

# Learning params
learning_rate = 0.01
num_epochs = 10
batch_size = 10

npArrChunkSize =4000
# Network params
dropout_rate = 0.5
num_classes = 2
train_layers = ['fc8', 'fc7']

# How often we want to write the tf.summary data to disk
display_step = 1

# Path for tf.summary.FileWriter and to store model checkpoints
filewriter_path = "./finetune_alexnet/duckPose/"
checkpoint_path = "./finetune_alexnet/"

# Create parent path if it doesn't exist
if not os.path.isdir(checkpoint_path): os.mkdir(checkpoint_path)


if(TRAINING==1):
    # TF placeholder for graph input and output
    x = tf.placeholder(tf.float32, [batch_size, 227, 227, 4])
    y = tf.placeholder(tf.float32, [batch_size, num_classes])
    keep_prob = tf.placeholder(tf.float32)

    # Initialize model
    model = AlexNet(x, keep_prob, num_classes)
    #
    # Link variable to model output
    score = model.fc8
    softmax = tf.nn.softmax(score)
    with tf.name_scope("cross_ent"):
        loss = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(logits=score, labels=y))

    # Op for calculating the loss
    # Train op
    with tf.name_scope("train"):
        # Create optimizer and apply gradient descent to the trainable variables
        train_op = tf.train.GradientDescentOptimizer(learning_rate).minimize(loss)

    # Add the loss to summary
    tf.summary.scalar('cross_entropy', loss)

    # Evaluation op: Accuracy of the model
    with tf.name_scope("accuracy"):
        correct_pred = tf.equal(tf.argmax(score, 1), tf.argmax(y, 1))
        accuracy = tf.reduce_mean(tf.cast(correct_pred, tf.float32))

    # Add the accuracy to the summary
    tf.summary.scalar('accuracy', accuracy)
    # Merge all summaries together
    merged_summary = tf.summary.merge_all()

    # Initialize the FileWriter
    # writer = tf.summary.FileWriter(filewriter_path)

    # Initialize an saver for store model checkpoints
    saver = tf.train.Saver()

    # Initalize the data generator seperately for the training and validation set
    # train_generator = ImageDataGenerator(train_file, 1, npArrChunkSize=100,
    #                                      horizontal_flip=False, shuffle=True)
    # val_generator = ImageDataGenerator(val_file, 2, npArrChunkSize=100, shuffle=False)

    tr_data = ImageDataGenerator(train_file,
                                 mode='training',
                                 batch_size=batch_size,
                                 num_classes=num_classes,
                                 shuffle=False)
    val_data = ImageDataGenerator(val_file,
                                  mode='inference',
                                  batch_size=batch_size,
                                  num_classes=num_classes,
                                  shuffle=False)

    # create an reinitializable iterator given the dataset structure
    iterator = Iterator.from_structure(tr_data.data.output_types,
                                       tr_data.data.output_shapes)
    next_batch = iterator.get_next()

# Ops for initializing the two different iterators
    training_init_op = iterator.make_initializer(tr_data.data)
    validation_init_op = iterator.make_initializer(val_data.data)
    # Get the number of training/validation steps per epoch
    # num_train_npy_files = np.floor(train_generator.data_size / npArrChunkSize).astype(np.int16)
    # num_validate_npy_files = np.floor(val_generator.data_size / npArrChunkSize).astype(np.int16)
    #
    # train_batches_per_epoch = (num_train_npy_files * (npArrChunkSize / batch_size)).astype(np.int16)
    # val_batches_per_epoch = (num_validate_npy_files * (npArrChunkSize / batch_size)).astype(np.int16)

    # Get the number of training/validation steps per epoch
    train_batches_per_epoch = int(np.floor(tr_data.data_size / batch_size))
    val_batches_per_epoch = int(np.floor(val_data.data_size / batch_size))

elif (TRAINING==0):
    # TF placeholder for graph input and output
    x = tf.placeholder(tf.float32, [1, 227, 227, 4])
    keep_prob = tf.placeholder(tf.float32)

    # Initialize model
    model = AlexNet(x, keep_prob, num_classes)
    #
    # Link variable to model output
    score = model.fc8
    softmax = tf.nn.softmax(score)
    testImg = getTestImage()
writer = tf.summary.FileWriter(filewriter_path)

# Start Tensorflow session
with tf.Session() as sess:

    if(TRAINING==1):
          # Initialize all variables
          sess.run(tf.global_variables_initializer())

          # Add the model graph to TensorBoard
          writer.add_graph(sess.graph)

          print("{} Start training...".format(datetime.now()))
          print("{} Open Tensorboard at --logdir {}".format(datetime.now(),
                                                            filewriter_path))

          # Loop over number of epochs
          for epoch in range(num_epochs):

                print("{} Epoch number: {}".format(datetime.now(), epoch+1))

                # Initialize iterator with the training dataset
                sess.run(training_init_op)

                for step in range(train_batches_per_epoch):

                    # get next batch of data
                    img_batch, label_batch = sess.run(next_batch)

                    image_summary_t = tf.summary.image('tn', img_batch, max_outputs=1)

                    # And run the training op
                    _, image_summary=sess.run([train_op,image_summary_t], feed_dict={x: img_batch,
                                                  y: label_batch,
                                                  keep_prob: dropout_rate})

                    # Generate summary with the current batch of data and write to file
                    if step%display_step == 0:
                        s = sess.run(merged_summary, feed_dict={x: img_batch,
                                                                y: label_batch,
                                                                keep_prob: 1.})
                        writer.add_summary(s, epoch*train_batches_per_epoch + step)
                        writer.add_summary(image_summary)

                    step += 1

                # Validate the model on the entire validation set
                print("{} Start validation".format(datetime.now()))
                test_acc = 0.
                test_count = 0
                for _ in range(val_batches_per_epoch):

                    img_batch, label_batch = sess.run(next_batch)
                    acc = sess.run(accuracy, feed_dict={x: img_batch,
                                                        y: label_batch,
                                                        keep_prob: 1.})
                    test_acc += acc
                    test_count += 1
                test_acc /= test_count
                print("{} Validation Accuracy = {:.4f}".format(datetime.now(), test_acc))

                # # Reset the file pointer of the image data generator
                # val_generator.reset_pointer()
                # train_generator.reset_pointer()

                print("{} Saving checkpoint of model...".format(datetime.now()))

                #save checkpoint of the model
                checkpoint_name = os.path.join(checkpoint_path, 'model_epoch'+str(epoch+1)+'.ckpt')
                save_path = saver.save(sess, checkpoint_name)

                print("{} Model checkpoint saved at {}".format(datetime.now(), checkpoint_name))

    elif(TRAINING==0): #Evaluation
        print 'inside Evaluation'
        # sess.run(tf.global_variables_initializer())

        # define activation of last layer as score
        saver = tf.train.import_meta_graph('./models/model_epoch100.ckpt.meta')
        saver.restore(sess, "./models/model_epoch100.ckpt")
        # image_summary_t = tf.summary.image('tn', testImg[:,:,:,3:4], max_outputs=1)
        all_vars = tf.get_collection('weight')
        for v in all_vars:
            v_ = sess.run(v)
            print(v_)
        # And run the training op

        pred,image_summary = sess.run([softmax], feed_dict={x: testImg,keep_prob: 1.})
        writer.add_summary(image_summary)

        print np.argmax(pred,1)





