import os
import numpy as np
import tensorflow as tf
tf.python.control_flow_ops = tf
import sample_datasets


BATCH_SIZE = 100 
tf.app.flags.DEFINE_integer('num_epochs', 100, 'Number of training epochs.')
tf.app.flags.DEFINE_float('svmC', 1, 'The C parameter of the SVM cost function.')
tf.app.flags.DEFINE_boolean('verbose', False, 'Produce verbose output.')

FLAGS = tf.app.flags.FLAGS



def dataloading(filename, AllTorOneF, NoValids):
  
  malwareset, benignset, malwarenames, benignnames, featNos = sample_datasets.load(filename, AllTorOneF)
  #malwareset, benignset = sample_datasets.normTwoSets(malwareset, benignset)
  if NoValids ==False:
    mXYtrain, mXYtest, mYtrain, mYtest, bXYtrain, bXYtest, bYtrain, bYtest= sample_datasets.divideTwoSets(malwareset, benignset, 0.16)
  else:
    mXYtrain, mXYtest, mYtrain, mYtest, bXYtrain, bXYtest, bYtrain, bYtest= sample_datasets.divideTwoSets(malwareset, benignset, 0.2)
  #concatenate malwareSet/benignSet
  #shuffle instances in the concatenated data set
  #final training/test set seperation
  numColumn = np.shape(mXYtrain)[1]
  mbXY_train = np.concatenate((mXYtrain[:,:],bXYtrain[:,:]))
  mbXY_test = np.concatenate((mXYtest[:,:],bXYtest[:,:]))
  np.random.shuffle(mbXY_train)
  np.random.shuffle(mbXY_test)


  X_train = mbXY_train[:,0:numColumn-1]
  X_test = mbXY_test[:,0:numColumn-1]
  Y_train = mbXY_train[:,numColumn-1]
  Y_test = mbXY_test[:,numColumn-1]

  
  return X_train, X_test, Y_train, Y_test, featNos

def main(arffindex):
  
  farffs = open("./arffFileList.txt",'r')
  arfflist = farffs.readlines()
  farffs.close()

  filename = "/media/ktg/New Volume/AndroidMalPaper/Arff/"
  filename += arfflist[arffindex][:-1]
  print(filename)

  fresult = open('SVM_.txt','a')

  

  with tf.device('/gpu:0'):

    X_train, X_test, Y_train, Y_test, featNos = dataloading(filename, False, False)
    print('data loading: done')

    verbose = FLAGS.verbose
    num_epochs = FLAGS.num_epochs
    train_size, tr_num_features = X_train.shape
    svmC = FLAGS.svmC

    Y_train[Y_train==0] = -1
    Y_test[Y_test==0] = -1

    x = tf.placeholder("float", shape=[None, tr_num_features])
    y = tf.placeholder("float", shape=[None,1])

    W = tf.Variable(tf.zeros([tr_num_features,1]))
    b = tf.Variable(tf.zeros([1]))
    y_raw = tf.matmul(x,W) + b

    regularization_loss = 0.5*tf.reduce_sum(tf.square(W)) 
    hinge_loss = tf.reduce_sum(tf.maximum(tf.zeros([BATCH_SIZE,1]), 
        1 - y*y_raw));
    svm_loss = regularization_loss + svmC*hinge_loss;
    train_step = tf.train.GradientDescentOptimizer(0.01).minimize(svm_loss)

    predicted_class = tf.sign(y_raw);
    correct_prediction = tf.equal(y,predicted_class)
    accuracy = tf.reduce_mean(tf.cast(correct_prediction, "float"))
  print('training: start')
  with tf.Session() as s:
    tf.initialize_all_variables().run()
    # Iterate and train.
    for step in xrange(num_epochs * train_size // BATCH_SIZE):
      if verbose:
        print step,
            
      offset = (step * BATCH_SIZE) % train_size
      batch_data = X_train[offset:(offset + BATCH_SIZE), :]
      #print(offset)
      if offset+BATCH_SIZE > Y_train.shape[0]:
        batch_labels = np.reshape(Y_train[offset:(offset + BATCH_SIZE)],(Y_train.shape[0]-offset,1))
        BATCH_SIZE = Y_train.shape[0]-offset
        train_step.run(feed_dict={x: batch_data, y: batch_labels})
        print 'loss: ', svm_loss.eval(feed_dict={x: batch_data, y: batch_labels})

      else:
        batch_labels = np.reshape(Y_train[offset:(offset + BATCH_SIZE)],(BATCH_SIZE,1))
        train_step.run(feed_dict={x: batch_data, y: batch_labels})
        print 'loss: ', svm_loss.eval(feed_dict={x: batch_data, y: batch_labels})

    print('training: done')  
    print('testing: start')
    accresult = accuracy.eval(feed_dict={x: X_test, y: Y_test})
    print('ACC: '+str(accresult))
    fresult.write(str(filename)+' '+str(accresult)+'\n')
    fresult.close()
    print('testing: done')

#if __name__ == '__main__':
#    tf.app.run()