import tensorflow as tf
from data_reader import *
import numpy as np
from uncompress import *
import os
from ae import *

slim = tf.contrib.slim

def log(message,file_path=os.path.join('cnn_logs_ae','log.txt')):

    print message
    f1=open(file_path, 'a+')
    f1.write(message)
    f1.close()



def lrelu(alpha):
    def op(inputs):
        return tf.maximum(alpha * inputs, inputs, name='leaky_relu')
    return op



def conv_net(input):

    with slim.arg_scope([slim.conv2d, slim.fully_connected], #using scope to avoid mentioning the paramters repeatdely
                                        activation_fn=lrelu(0.005),
                                        weights_initializer=tf.truncated_normal_initializer(0.0, 0.01),
                                        weights_regularizer=slim.l2_regularizer(0.0005)):
	    # net = slim.max_pool2d(input,(1,4),(1,4), padding='VALID', scope='pool_0')
	

        net = slim.conv2d(input, 512, (3,1357), 1, padding='VALID', scope='cnn_conv_1')
                        
        net = slim.max_pool2d(net, (4,1),4, padding='VALID', scope='cnn_pool_2')
                        
        net = slim.conv2d(net, 512, (5,1), 1, scope='cnn_conv_3')

        net = slim.max_pool2d(net, (4,1),4, padding='VALID', scope='cnn_pool_4')

        net = slim.flatten(net, scope='cnn_flatten_5')

        # net = slim.fully_connected(net, 1024, scope='fc_6',activation_fn=tf.nn.softmax)
	
        # net = slim.fully_connected(net, 256, scope='fc_7',activation_fn=tf.nn.softmax)

        net = slim.fully_connected(net, 2, scope='cnn_fc_8',activation_fn=tf.nn.softmax)
	
	# net = slim.fully_connected(net, 4096, scope='fc5')
	# net = slim.dropout(net, 0.5, scope='dropout6')
	# net = slim.fully_connected(net, 4096, scope='fc7')
	# net = slim.dropout(net, 0.5, scope='dropout8')
	# net = slim.fully_connected(net,2, activation_fn=None, scope='fc9')

    
    return net

def one_hot(batch_size,Y):

    B = np.zeros((batch_size,2))

    B[np.arange(batch_size),Y] = 1

    return B

if __name__=='__main__':

    os.environ['CUDA_VISIBLE_DEVICES'] = '1'
    # print one_hot(3,np.array((1,0,1)))
    # exit(0)

    x = compress(None,None,True)
    #print x
    # Training Parameters
    learning_rate = 0.00001
    # learning_rate = tf.train.exponential_decay(
    # intial_learning_rate, global_step, decay_steps,
    # decay_rate, True, name='learning_rate')


    

    ################################################################
    num_epoch = 1
    batch_size = 1
    display_step = 1
    input_size = 50
    num_classes = 2 

    X = tf.placeholder(tf.float32, [None, input_size,86796,1])

    reconstruction, compressed,_,_,_,_ = autoencoder(X)

  

    cnn_X = tf.placeholder(tf.float32, [None, input_size,1357,8])
    
    cnn_Y = tf.placeholder(tf.float32, [None, num_classes])
    #logits = conv_net(X)
    #prediction = tf.nn.softmax(logits)
    cnn_prediction = conv_net(cnn_X)
    # Define loss and optimizer
    '''loss_op = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(
	logits=logits, labels=Y))'''
    cnn_loss_op = slim.losses.softmax_cross_entropy(cnn_prediction,cnn_Y)

    tf.summary.scalar('loss',cnn_loss_op)
    cnn_optimizer = tf.train.AdagradOptimizer(learning_rate=learning_rate)
    cnn_train_op = cnn_optimizer.minimize(cnn_loss_op)


    cnn_correct_pred = tf.equal(tf.argmax(cnn_prediction, 1), tf.argmax(cnn_Y, 1))
    cnn_accuracy = tf.reduce_mean(tf.cast(cnn_correct_pred, tf.float32))
    tf.summary.scalar('accuracy',cnn_accuracy)


    cnn_init = tf.global_variables_initializer()

    data_X,data_Y = load_data()
    

    indices = np.random.permutation(np.arange(data_X.shape[0]))

    data_X = data_X[indices,:,:]

    data_Y = data_Y[indices]

    cnn_merged = tf.summary.merge_all()
    cnn_saver = tf.train.Saver()

    # print cnn_saver._var_list
    # exit(0)

    saver = tf.train.Saver(var_list=x)

    with tf.Session() as sess:

        cnn_train_writer = tf.summary.FileWriter("cnn_logs_ae/",
                                    sess.graph)

        # Run the initializer
        sess.run(cnn_init)
        
        print 'restoring ae session'
        saver.restore(sess, "ae_logs_1/save.ckpt")
        print 'done loading'   

        '''

        print 'restoring session'
        saver.restore(sess, "logs3/epoch0i180.ckpt")
        print 'done loading'
        # exit(0) '''

        i = 0
        train_acc = 0.0
        print 'started training'
        for epoch in range(num_epoch):
            for step in range(data_X.shape[0]/batch_size):
                batch_x, batch_y = data_X[step*batch_size:(step+1)*batch_size],\
                data_Y[step*batch_size:(step+1)*batch_size]

                i+=1

                batch_x = uncompress(batch_x,86796)
                batch_x = sess.run(compressed, feed_dict={X: batch_x})
                # print 'batch_x ready'
                # exit(0)
                batch_y = one_hot(batch_size,batch_y)

                batch_y = np.repeat(batch_y,50,axis=0)

                # print batch_y

                assert(batch_x.shape[0]==batch_y.shape[0])

                _,summary = sess.run([cnn_train_op,cnn_merged], feed_dict={cnn_X: batch_x, cnn_Y: batch_y})
                cnn_train_writer.add_summary(summary, i)
                if step % display_step == 0:
                    # Calculate batch loss and accuracy
                    loss, acc,summary = sess.run([cnn_loss_op, cnn_accuracy,cnn_merged], feed_dict={cnn_X: batch_x,
                                                                        cnn_Y: batch_y})
                    log("LR : "+str(learning_rate)+" Epoch : " + str(epoch) + " Step " + str(step) + ", Minibatch Loss= " + \
                        "{:.4f}".format(loss) + ", Training Accuracy= " + \
                        "{:.3f}".format(acc))
                    
                    train_acc += acc
                    log("Mean Train Accuracy:" + "{:.3f}".format(train_acc/i))

                    
                    
                    # train_writer.add_summary(summary, step)
                
                if i%20 == 0:

                    print 'saving checkpoint'
                    save_path = cnn_saver.save(sess, os.path.join('cnn_logs_ae','save.ckpt'))
                    print("Model saved in path: %s" % save_path)
                
                

        print 'done!'
        
