from tensorflow.contrib.learn.python.learn import metric_spec
from tensorflow.contrib.learn.python.learn.estimators import estimator
from tensorflow.contrib.learn.python.learn.estimators.svm import SVM 
from tensorflow.python.platform import app
import sample_datasets
import tensorflow as tf
import numpy as np

FLAGS = None
BATCH_SIZE = 100 


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)

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

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

    ks = list()
    for i in range(X_train.shape[1]):
      ks.append(str(i))
    cname = "feats"
    sparceCol = tf.contrib.layers.sparse_column_with_keys(column_name=cname, keys=ks)
    est = SVM(example_id_column=str(arffindex), feature_columns=[sparceCol])
    est.fit(x=X_train, y=np.reshape(Y_train,(Y_train.shape[0],1)), batch_size = BATCH_SIZE)

    metric_name = 'accuracy'
    metric = {metric_name: metric_spec.MetricSpec(eval_metrics.get_metric(metric_name), prediction_key=eval_metrics.get_prediction_key(metric_name))}
    results = est.evaluate(x=X_test, y=np.reshape(Y_test,(Y_test.shape[0],1)), batch_size=FLAGS.batch_size, metrics=metric)

    for key in sorted(results):
      print('%s: %s' % (key, results[key]))
    exit()


    