import os
import sys
import numpy as np
import tensorflow as tf
import random
import basic_cnn_model
import pickle
import cv2
from scipy import stats
import math


global g_train_files1,g_train_classes,g_test_files1,g_test_classes,g_val_files1,g_val_classes,g_val_files2,g_test_files2,g_train_files2

############################Parms ################################################
learning_rate=0.009
log_dir='/media/anilil/Data/charm/Cnn/Output_logs/xml_ucf_cnn/'
batch_size=64
training_iters=70000
n_classes=101
data_root='/media/anilil/Data/Datasets/UCf_scales/ucf_xmlparse/MV_UCF101_H264_CRF18_GOP30/'
display_step=100
############################Parms ################################################

def load_from_file(name):
    with open(name, "rb") as f:
        dictname = pickle.load(f)
    return dictname

def keep_png(path,listr):
    ll = list()
    for x in listr:
        if os.path.isfile(path+str(x)) & str(x).endswith('.jpg'):
            ll.append(x)
    return ll

def read_from_tex(filename):
    data1=list()
    classes=list()
    with open(filename,'r') as f:
        for each_line in f:
            x=str(each_line).replace('\n','').replace('\r','').split(',')
            data1.append(x[0])
            data2.append(x[1])
            classes.append(x[2])
    return data1,data2,classes

def readfiles_fromtext(filename):
    data1=list()
    data2=list()
    classes=list()
    with open(filename,'r') as f:
        for each_line in f:
            x=str(each_line).replace('\n','')
            data1.append(x.split(' ')[0])
            classes.append(x.split(' ')[1])
    return data1,classes

def train_loadbatch_from_lists(b_siz,pos):
    if pos+b_siz>len(g_train_classes):
        pos=pos%len(g_train_classes)
    files1,classes=g_train_files1[pos:pos+b_siz],g_train_classes[pos:pos+b_siz]
    temp_batch=np.zeros(shape=(b_siz,227,227,3),dtype=np.float32)
    temp_classes=np.zeros(shape=(b_siz,101),dtype=np.float32)
    for x in range(0,b_siz):
        temp_batch[x,:,:,:]=cv2.resize(cv2.imread(files1[x]),(227,227))
        temp_classes[x,:]=create_one_hot(int(classes[x]))
    temp_batch= (temp_batch-127)/127
    return temp_batch,temp_classes

def replace_root(listw):
    x= list()
    for each in listw:
        x.append(str(each).replace('/media/anilil/Data/Datasets/UCf_scales/ucf_xmlparse/MV_UCF101_H264_CRF18_GOP30/',data_root))
    return x

def test_loadbatch_from_lists(b_siz,pos):
    if pos+b_siz>len(g_test_classes):
        pos=pos%len(g_test_classes)
    files1,classes=g_test_files1[pos:pos+b_siz],g_test_classes[pos:pos+b_siz]
    temp_batch=np.zeros(shape=(b_siz,227,227,3),dtype=np.float32)
    temp_classes=np.zeros(shape=(b_siz,101),dtype=np.float32)
    for x in range(0,b_siz):
        temp_batch[x,:,:,:]=cv2.resize(cv2.imread(files1[x]),(227,227))
        temp_classes[x,:]=create_one_hot(int(classes[x]))
    temp_batch= (temp_batch-127)/127
    return temp_batch,temp_classes


def create_one_hot(num,len1=101):
    tem=np.zeros(shape=(len1))
    tem[int(num)]=1.0
    return tem

def shuffle_list(list1,list2,num=3):
    combined = zip(list1, list2)
    for x in range(0,num):
        random.shuffle(combined)
    list1,list2= zip(*combined)
    return list1,list2

def load_model(log_dir):
    ckpt = tf.train.get_checkpoint_state(log_dir)
    if ckpt and ckpt.model_checkpoint_path:
        saver.restore(sess, ckpt.model_checkpoint_path)


# Store layers weight & bias
weights = {
    'wc1': tf.Variable(tf.random_normal([7, 7, 3, 96],stddev=0.01)),
    'wc2': tf.Variable(tf.random_normal([5, 5, 96, 384],stddev=0.01)),
    'wc3': tf.Variable(tf.random_normal([3, 3, 384, 512],stddev=0.01)),
    'wc4': tf.Variable(tf.random_normal([3, 3, 512, 512],stddev=0.01)),
    'wc5': tf.Variable(tf.random_normal([3, 3, 512, 384],stddev=0.01)),
    'wd1': tf.Variable(tf.random_normal([2*2*384, 4096],stddev=0.01)),
    'wd2': tf.Variable(tf.random_normal([4096,4096 ],stddev=0.01)),
    'out': tf.Variable(tf.random_normal([4096,101],stddev=0.01))
}
biases = {
    'bc1': tf.Variable(tf.random_normal([96])),
    'bc2': tf.Variable(tf.random_normal([384])),
    'bc3': tf.Variable(tf.random_normal([512])),
    'bc4': tf.Variable(tf.random_normal([512])),
    'bc5': tf.Variable(tf.random_normal([384])),
    'bd1': tf.Variable(tf.random_normal([4096])),
    'bd2': tf.Variable(tf.random_normal([4096])),
    'out': tf.Variable(tf.random_normal([101]))
}



def conv2d(name, l_input, w, b,stride=1):
    return tf.nn.relu(tf.nn.bias_add(tf.nn.conv2d(l_input, w, strides=[1, stride, stride, 1], padding='VALID'),b), name=name)

def max_pool(name, l_input, k,stride=1):
    return tf.nn.max_pool(l_input, ksize=[1, k, k, 1], strides=[1, stride, stride, 1], padding='VALID', name=name)

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

def cnn_model(_X, _weights, _biases, _dropout):

    # 1st Convolution + Relu Layer
    conv1 = conv2d('conv1', _X, _weights['wc1'], _biases['bc1'],stride=2)#stride 2
    # Max Pooling (down-sampling)
    pool1 = max_pool('pool1', conv1, k=3,stride=2)
    # Apply Normalization
    norm1 = norm('norm1', pool1, lsize=5)

    # 2nd Convolution + Relu Layer
    conv2 = conv2d('conv2', norm1, _weights['wc2'], _biases['bc2'],stride=2)
    # Max Pooling (down-sampling)
    pool2 = max_pool('pool2', conv2, k=3,stride=2)
    # Apply Normalization
    norm2 = norm('norm2', pool2, lsize=5)

    # 3rd Convolution Layer
    conv3 = conv2d('conv3', norm2, _weights['wc3'], _biases['bc3'])
    # 4th conv layer
    conv4 = conv2d('conv4', conv3, _weights['wc4'], _biases['bc4'])
    # 5th conv layer
    conv5 = conv2d('conv4', conv4, _weights['wc5'], _biases['bc5'])
    # Pooling
    pool5 = max_pool('pool5', conv5, k=3,stride=2)

    # Fully connected layer
    dense1 = tf.reshape(pool5, [-1, _weights['wd1'].get_shape().as_list()[0]]) # Reshape conv3 output to fit dense layer input
    dense1 = tf.nn.relu(tf.matmul(dense1, _weights['wd1']) + _biases['bd1'], name='fc1') # Relu activation
    # Apply Dropout
    dense1 = tf.nn.dropout(dense1, _dropout)

    dense2 = tf.nn.relu(tf.matmul(dense1, _weights['wd2']) + _biases['bd2'], name='fc2') # Relu activation
    # Apply Dropout
    dense2 = tf.nn.dropout(dense2, _dropout)

    #Output, class prediction
    out = tf.matmul(dense2, _weights['out']) + _biases['out']

    return out

g_train_files1,g_train_classes=readfiles_fromtext('/home/anilil/projects/lstm/lisa-caffe-public/examples/LRCN_activity_recognition/singleframe_flow/ucf_xml_train_list.txt')
g_test_files1,g_test_classes=readfiles_fromtext('/home/anilil/projects/lstm/lisa-caffe-public/examples/LRCN_activity_recognition/singleframe_flow/ucf_xml_test_list.txt')



x = tf.placeholder(tf.float32, [None, 227,227,3])
y = tf.placeholder(tf.float32, [None, 101])
keep_prob = tf.placeholder(tf.float32)

pred= cnn_model(x,weights,biases,keep_prob)
# Define loss and optimizer
cost = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(pred, y)) # Softmax loss
optimizer = tf.train.RMSPropOptimizer(learning_rate=learning_rate).minimize(cost) # Adam Optimizer
# Evaluate model

correct_pred = tf.equal(tf.argmax(pred,1), tf.argmax(y,1))
accuracy = tf.reduce_mean(tf.cast(correct_pred, tf.float32))
Train_loss_summ=tf.scalar_summary("Train_Loss", cost)
Train_acc_summ=tf.scalar_summary("Train_Accuracy", accuracy)
test_loss_sum = tf.scalar_summary("Test_Loss", cost)
test_accuracy_sum= tf.scalar_summary("Test_Accuracy", accuracy)
merged_summary_op = tf.merge_all_summaries()
saver = tf.train.Saver()



if not os.path.isdir(log_dir):
    os.makedirs(log_dir)
# Initializing the variables
init = tf.initialize_all_variables()
f= open(log_dir+'log.log','w')
f.write ("Learning Rate : {} \n".format(learning_rate))
f.write ("Batch Size : {} \n".format(batch_size))
f.write ("Total Number of Training itters to run : {} \n".format(training_iters))
f.write ("Data from path : {}\n".format(data_root))
f.write ("Number of classes : {}\n".format(n_classes))
####################################################################################################################################################################################

f.write ("Type of Optimizer : RMSPropOptimizer \n")
##########################################################################################

f.close()
gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=0.6)


with tf.Session(config=tf.ConfigProto(gpu_options=gpu_options)) as sess:#config=tf.ConfigProto(gpu_options=gpu_options)
    sess.run(init)
    step = 0
    epoch=0
    temp=1
    summary_writer = tf.train.SummaryWriter(log_dir, graph_def=sess.graph_def)
    # Keep training until reach max iterations
    while step * batch_size < training_iters:

        batch_xs, batch_ys = train_loadbatch_from_lists(batch_size,step * batch_size)
        sess.run(optimizer, feed_dict={x: batch_xs, y: batch_ys,
                                       keep_prob: 0.5})
        if step % display_step == 0:
            # Calculate batch accuracy and loss

            tr_acc_sum,tr_loss_sum,acc,loss = sess.run([Train_acc_summ,Train_loss_summ,accuracy,cost], feed_dict={x: batch_xs, y: batch_ys,
                                                keep_prob: 0.5})

            test_data,test_label=test_loadbatch_from_lists(batch_size,step * batch_size)

            test_acc_sum,test_los_sum,test_acc ,test_cost=sess.run([test_accuracy_sum,test_loss_sum,accuracy, cost], feed_dict={x: test_data, y: test_label,
                                                keep_prob: 1})
            summary_writer.add_summary(tr_acc_sum, step * batch_size +1)
            summary_writer.add_summary(tr_loss_sum, step * batch_size +1)
            summary_writer.add_summary(test_acc_sum, step * batch_size +1)
            summary_writer.add_summary(test_los_sum, step * batch_size +1)
            f=open(log_dir+'log.log','a')
            f.write ("Iter={:10d},Train_Loss={:.6f},Train_Accuracy={:.6f},Test_Loss={:.6f},Test_Accuracy={:.6f} \n".format(step * batch_size,loss,acc,test_cost,test_acc))
            f.close()
            print "Iter={:10d},Train_Loss={:.6f},Train_Accuracy={:.6f},Test_Loss={:.6f},Test_Accuracy={:.6f} \n".format(step * batch_size,loss,acc,test_cost,test_acc)
        step += 1
        temp =temp+  batch_size
        if (temp>15000):
            temp=1
            epoch=epoch+1
            saver.save(sess,log_dir+'model.ckpt',global_step=step * batch_size+1)

    print ("Optimization Finished!")