from datetime import datetime
import math
import time
import tensorflow as tf
import numpy as np
import scipy as sci
import cv2
import input_data_conv
import skimage.transform
from skimage import color

import tensorflow.python.platform
import tensorflow as tf

FLAGS = tf.app.flags.FLAGS
# Parameters
learning_rate = 0.001
training_iters = 200000
batch_size = 64
display_step = 20
n_classes=101 # number of classes

#Input data and classes
global train_data,train_class,test_data,test_classs,train_i,test_i
test_i, train_i = 0,0
train_data=input_data_conv.train_single_frames
train_class=input_data_conv.train_single_classes
test_data=input_data_conv.test_single_frames
test_classs=input_data_conv.test_single_classes


# Network Parameters
n_input = [227, 227, 3 ]# MNIST data input (img shape: 227*227*3)
dropout = 0.5 # Dropout, probability to keep units



# TODO: why is batch size 64 going OOM?
tf.app.flags.DEFINE_integer('batch_size', 64,
                            """Batch size.""")
tf.app.flags.DEFINE_integer('num_batches', 100,
                            """Number of batches to run.""")
tf.app.flags.DEFINE_boolean('forward_only', False,
                            """Only run the forward pass.""")
tf.app.flags.DEFINE_boolean('forward_backward_only', False,
                            """Only run the forward-forward pass.""")
tf.app.flags.DEFINE_string('data_format', 'NHWC',
                           """The data format for Convnet operations.
                           Can be either NHWC or NCHW.
                           """)

parameters = []

conv_counter = 1
pool_counter = 1
affine_counter = 1

def _conv(inpOp, nIn, nOut, kH, kW, dH, dW, padType):
    global conv_counter
    global parameters
    name = 'conv' + str(conv_counter)
    conv_counter += 1
    with tf.name_scope(name) as scope:
        kernel = tf.Variable(tf.truncated_normal([kH, kW, nIn, nOut],
                                                 dtype=tf.float32,
                                                 stddev=1e-1), name='weights')
        if FLAGS.data_format == 'NCHW':
          strides = [1, 1, dH, dW]
        else:
          strides = [1, dH, dW, 1]
        conv = tf.nn.conv2d(inpOp, kernel, strides, padding=padType)
        biases = tf.Variable(tf.constant(0.0, shape=[nOut], dtype=tf.float32),
                             trainable=True, name='biases')
        #bias = tf.reshape(tf.nn.bias_add(conv, biases, conv.get_shape()))
        bias = tf.nn.bias_add(conv, biases)
        conv1 = tf.nn.relu(bias, name=scope)
        parameters += [kernel, biases]# saved to global var
        return conv1

def _affine(inpOp, nIn, nOut): #Dense Layers defined
    global affine_counter
    global parameters
    name = 'affine' + str(affine_counter)
    affine_counter += 1
    with tf.name_scope(name) as scope:
        kernel = tf.Variable(tf.truncated_normal([nIn, nOut],
                                                 dtype=tf.float32,
                                                 stddev=1e-1), name='weights')
        biases = tf.Variable(tf.constant(0.0, shape=[nOut], dtype=tf.float32),
                             trainable=True, name='biases')
        affine1 = tf.nn.relu_layer(inpOp, kernel, biases, name=name)
        parameters += [kernel, biases]
        return affine1

def norm(name, l_input, lsize):
    return tf.nn.lrn(l_input, lsize, bias=1.0, alpha=0.0001 / 9.0, beta=0.75, name=name)

def _mpool(inpOp, kH, kW, dH, dW): # Max pooling
    global pool_counter
    global parameters
    name = 'pool' + str(pool_counter)
    pool_counter += 1
    if FLAGS.data_format == 'NCHW':
      ksize = [1, 1, kH, kW]
      strides = [1, 1, dH, dW]
    else:
      ksize = [1, kH, kW, 1]
      strides = [1, dH, dW, 1]
    return tf.nn.max_pool(inpOp,
                          ksize=ksize,
                          strides=strides,
                          padding='VALID',
                          name=name)

def loss(logits, labels): # loss function
    batch_size = tf.size(labels)
    labels = tf.expand_dims(labels, 1)
    indices = tf.expand_dims(tf.range(0, batch_size, 1), 1)
    concated = tf.concat(1, [indices, labels])
    onehot_labels = tf.sparse_to_dense(
        concated, tf.pack([batch_size, 1000]), 1.0, 0.0)
    cross_entropy = tf.nn.softmax_cross_entropy_with_logits(logits,
                                                            onehot_labels,
                                                            name='xentropy')
    loss = tf.reduce_mean(cross_entropy, name='xentropy_mean')

def resize_im(im, nh, nw):
    im=np.copy(im)
    h, w, _ = im.shape
    im = skimage.transform.resize(im, (nh, nw), preserve_range=True)
    return im

def create_class_vec(val,nuoclasses):
    x=np.zeros(nuoclasses)
    x[val]=1
    return x


def Vgg_single_frame(_X,_dropout): # Network definiton
    #_X = tf.reshape(_X, shape=[-1, 227, 227, 3])
    conv1 = _conv (_X, 3, 96, 7, 7, 2, 2, 'SAME')# _conv(inpOp, nIn, nOut, kH, kW, dH, dW, padType)
    pool1 = _mpool(conv1,  3, 3, 2, 2)#_mpool(inpOp, kH, kW, dH, dW)
    norm1 = norm('norm1', pool1, lsize=5)
    conv2 = _conv (norm1,  96, 384, 5, 5, 2, 2, 'SAME')
    pool2 = _mpool(conv2,  3, 3, 2, 2)
    norm2 = norm('norm2', pool2, lsize=5)
    conv3 = _conv (norm2,  384, 512, 3, 3, 1, 1, 'VALID')
    conv4 = _conv (conv3,  512, 512, 3, 3, 1, 1, 'VALID')
    conv5 = _conv (conv4,  512, 384, 3, 3, 1, 1, 'VALID')
    pool5 = _mpool(conv5,  3, 3, 2, 2)
    resh1 = tf.reshape(pool5, [-1, 3*3*384])
    dense1 = _affine(resh1, 3*3*384, 4096)
    drop1= tf.nn.dropout(dense1, _dropout)
    dense2= _affine(drop1, 4096, 4096)
    drop2= tf.nn.dropout(dense2, _dropout)
    affn3 = _affine(drop2, 4096, 101)

    return affn3

def train_next_batch(batch_size):
    temp_data=np.ndarray(shape=(batch_size,227,227,3),dtype=float)
    temp_class=np.ndarray(shape=(batch_size,n_classes),dtype=float)
    for idx,x in enumerate(train_data[train_i:train_i+batch_size]):
        temp_data[idx,:,:,:]=resize_im(cv2.imread(x,1),227,227)
        temp_class[idx,:]=create_class_vec(train_class[train_i+idx],101)
    return temp_data,temp_class

def test_next_batch(batch_size):
    temp_data=np.ndarray(shape=(batch_size,227,227,3),dtype=float)
    temp_class=np.ndarray(shape=(batch_size,n_classes),dtype=float)
    for idx,x in enumerate(test_data[test_i:test_i+batch_size]):
        temp_data[idx,:,:,:]=resize_im(cv2.imread(x,1),227,227)
        temp_class[idx,:]=create_class_vec(test_classs[train_i+idx],101)
    return temp_data,temp_class



# tf Graph input
x = tf.placeholder(tf.float32, [None, 227,227,3])
y = tf.placeholder(tf.float32, [None, n_classes])
keep_prob = tf.placeholder(tf.float32) # dropout (keep probability)
pred = Vgg_single_frame(x, keep_prob)
cost = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(pred, y))
optimizer = tf.train.GradientDescentOptimizer(learning_rate=learning_rate).minimize(cost)
correct_pred = tf.equal(tf.argmax(pred,1), tf.argmax(y,1))
accuracy = tf.reduce_mean(tf.cast(correct_pred, tf.float32))

init = tf.initialize_all_variables()

with tf.Session() as sess:
    sess.run(init)
    step = 1
    # Keep training until reach max iterations
    while step * batch_size < training_iters:
        batch_xs, batch_ys = train_next_batch(batch_size)
        train_i=train_i+batch_size
        # Fit training using batch data
        sess.run(optimizer, feed_dict={x: batch_xs, y: batch_ys, keep_prob: 0.5})

        if step % display_step == 0:
            # Calculate batch accuracy
            acc = sess.run(accuracy, feed_dict={x: batch_xs, y: batch_ys, keep_prob: 0.5})
            # Calculate batch loss
            loss = sess.run(cost, feed_dict={x: batch_xs, y: batch_ys, keep_prob: 0.5})
            print "Iter " + str(step*batch_size) + ", Minibatch Loss= " + "{:.6f}".format(loss) + ", Training Accuracy= " + "{:.5f}".format(acc)
            test_xs,test_ys = test_next_batch(batch_size)
            test_i=test_i+batch_size
            print "Testing Accuracy:", sess.run(accuracy, feed_dict={x: test_xs, y:test_ys, keep_prob: 0.0})
        step += 1
    print "Optimization Finished!"
    # Calculate accuracy for 256 mnist test images