import pandas
import datetime as dt

import cv2
from bigdl.nn.layer import *
from bigdl.nn.criterion import *
from bigdl.optim.optimizer import *
from bigdl.util.common import *
from bigdl.dataset.transformer import *
from bigdl.dataset import mnist
from matplotlib.pyplot import imshow
import matplotlib.pyplot as plt
from pyspark import SparkContext
from bigdl.util import common
from bigdl.util.common import JTensor
from bigdl.dataset import mnist
sc=SparkContext.getOrCreate(conf=create_spark_conf().setMaster("local[4]").set("spark.driver.memory","14g"))
sc.setLogLevel("ERROR")
init_engine()

#constants
batch_size=256
learning_rate=0.1
noofepochs=200




import os
from PIL import Image
import numpy as np
def read_data_sets(main_path,dataset):
    #pre-proccess images 
    filelist=os.listdir(main_path+dataset)
   
    imagelist_pre = np.array([np.array(cv2.imread(main_path+dataset+fname)) for fname in filelist])
    imagelist = imagelist_pre[..., None]
    #pre-proccess Labels to match the images
    label_file = open("32by32/Train_label.txt", "r")
    labels = label_file.read().split("\r")
    correctlabels=labels
   
    for i in range(len(filelist)):
        filelist[i]=filelist[i].replace("X_train","")
	filelist[i]=filelist[i].replace("X_test","")
        filelist[i]=int(filelist[i].replace(".png",""))
        correctlabels[i]=int(labels[filelist[i]])
    
    return (imagelist,correctlabels)





def get_malware(sc, malware_path):
    (train_images,train_labels) = read_data_sets(malware_path, "train/")
    (test_images,test_labels) = read_data_sets(malware_path, "test/")
    training_mean = np.mean(train_images)
    training_std = np.std(train_images)
   
    rdd_train_images = sc.parallelize(train_images)
    rdd_train_labels = sc.parallelize(train_labels)
    rdd_test_images = sc.parallelize(test_images)
    rdd_test_labels = sc.parallelize(test_labels)

    rdd_train_sample = rdd_train_images.zip(rdd_train_labels).map(lambda (features, label):
                                        common.Sample.from_ndarray(
                                        (features - training_mean) / training_std,
                                        label))

    rdd_test_sample = rdd_test_images.map(lambda (features):
                                         common.Sample(features))

    (rdd_train,rdd_valid) = rdd_train_sample.randomSplit(weights=[0.7, 0.3], seed=2)
  

    return (rdd_train,rdd_valid, rdd_test_sample)



mnist_path = "32by32/"
(train_data, val_data, test_data) = get_malware(sc, mnist_path)


def build_model(class_num):
            model = Sequential()
            model.add(Reshape([3, 32, 32]))
            model.add(SpatialConvolution(3, 64, 3, 3).set_name('conv1'))
            model.add(ReLU())
            model.add(SpatialMaxPooling(2, 2, 2, 2).set_name('pool1'))
            model.add(ReLU())
            model.add(SpatialConvolution(64, 128, 3, 3).set_name('conv2'))
            model.add(SpatialMaxPooling(2, 2, 2, 2).set_name('pool2'))
            model.add(ReLU())
            model.add(SpatialConvolution(128, 256, 2, 2).set_name('conv3'))
            model.add(SpatialMaxPooling(2, 2, 2, 2).set_name('pool3'))
            model.add(ReLU())
            model.add(SpatialConvolution(256, 512, 2, 2).set_name('conv4'))
            model.add(Reshape([512]))
            model.add(Linear(512, 100).set_name('fc1'))
            model.add(ReLU6())
            model.add(Linear(100, class_num).set_name('output'))
            model.add(Sigmoid())
            return model




dsp_black_model=build_model(9)
optimizer = Optimizer(
    model=dsp_black_model,
    training_rdd=train_data,
    criterion=CrossEntropyCriterion(),
    optim_method=Adam(learningrate=learning_rate, learningrate_decay=0.0002),
    end_trigger=MaxEpoch(noofepochs),
    batch_size=batch_size)


optimizer.set_validation(
    batch_size=256,
    val_rdd=val_data,
    trigger=EveryEpoch(),
    val_method=[Top1Accuracy()]
)




app_name='dspblack-'+dt.datetime.now().strftime("%Y%m%d-%H%M%S")
train_summary = TrainSummary(log_dir='/tmp/bigdl_summaries',
                                     app_name=app_name)
train_summary.set_summary_trigger("Parameters", SeveralIteration(50))
val_summary = ValidationSummary(log_dir='/tmp/bigdl_summaries',
                                        app_name=app_name)
optimizer.set_train_summary(train_summary)
optimizer.set_val_summary(val_summary)
print "saving logs to ",app_name




print "Training.....!"
trained_model =  optimizer.optimize()
print "End of Traininng... :) "




print "Predictions.."
predictions = trained_model.predict_class(test_data)

evaluate_result = trained_model.evaluate(val_data, 64, [Top1Accuracy()])
print evaluate_result[0]

predictions.saveAsTextFile("Result")
