from __future__ import print_function
import urllib

from pyspark.sql import SQLContext, Row,SparkSession
from pyspark.ml.feature import IDF, Tokenizer,HashingTF
from pyspark.ml.feature import VectorAssembler
from pyspark.ml import Pipeline
from pyspark.ml.classification import RandomForestClassifier
from pyspark.ml.feature import IndexToString, StringIndexer, VectorIndexer
import sys
from pyspark.ml.evaluation import MulticlassClassificationEvaluator
from pyspark.sql.types import *
from pyspark.ml.feature import CountVectorizer
from pyspark.sql.functions import lit
from pyspark.ml.feature import NGram

import re
import os


spark = SparkSession\
        .builder\
        .appName("MalwareClassification")\
        .getOrCreate()
sc=spark.sparkContext;
sqlContext=SQLContext(sc)
def main(args):
    #args[0]: Training byte parque file
    byteParqueDF=spark.read.parquet(args[0])
    byteParqueDF.createOrReplaceTempView("byteParqueDFTable")
    # args[1]: Testing byte parque file
    byteTestParqueDF = spark.read.parquet(args[1])
    byteTestParqueDF.createOrReplaceTempView("byteTestParqueDFTable")
    # args[2]: Training opcode parque file
    opcodeParqueDF=spark.read.parquet(args[2])
    opcodeParqueDF.createOrReplaceTempView("opcodeParqueDFTable")

    # args[3]: Testing opcode parque file
    opcodeTestParqueDF = spark.read.parquet(args[3])
    opcodeTestParqueDF.createOrReplaceTempView("opcodeTestParqueDFTable")

    joinedDF=spark.sql("SELECT A1.hashcodefile,A1.label,A1.byteFeatures,A2.opcodeFeatures from byteParqueDFTable A1,opcodeParqueDFTable A2 where A1.hashcodefile=A2.hashcodefile")
    joinedTestDF=spark.sql("SELECT A1.hashcodefile,A1.byteFeatures,A2.opcodeFeatures from byteTestParqueDFTable A1,opcodeTestParqueDFTable A2 where A1.hashcodefile=A2.hashcodefile")

    ###########################################################
    # Code for Assembler, joining byte and opcode features
    ###########################################################
    trainingAssembler = VectorAssembler(
        inputCols=["byteFeatures", "opcodeFeatures"],
        outputCol="features")
    testingAssembler = VectorAssembler(
        inputCols=["byteFeatures", "opcodeFeatures"],
        outputCol="features")
    trainingData=trainingAssembler.transform(joinedDF)
    testingData=testingAssembler.transform(joinedTestDF)


    # ###########################
    # Code for Random Forest Classifier
    labelIndexer = StringIndexer(inputCol="label", outputCol="indexedLabel").fit(trainingData)

        # # Train a RandomForest model.
    rf = RandomForestClassifier(labelCol="indexedLabel", featuresCol="features", numTrees=100,maxDepth=8,impurity="gini",seed=400)

        #  Convert indexed labels back to original labels.
    labelConverter = IndexToString(inputCol="prediction", outputCol="predictedLabel",
                                    labels=labelIndexer.labels)
    # # Chain indexers and forest in a Pipeline
    pipeline = Pipeline(stages=[labelIndexer,rf,labelConverter])
    #
    # # Train model.  This also runs the indexers.
    model = pipeline.fit(trainingData)
    # # # Make predictions.
    predictions = model.transform(testingData)
    # # # Select example rows to display.
    filterPredictions=predictions.select("predictedLabel","hashcodefile")

    predictionsRDD=filterPredictions.rdd

    predictionsRDD.saveAsTextFile("output.text")
    # #Select (prediction, true label) and compute test error

    evaluator = MulticlassClassificationEvaluator(
         labelCol="indexedLabel", predictionCol="prediction", metricName="accuracy")

    accuracy = evaluator.evaluate(filterPredictions)

    print("Test Error = %g" % (1.0 - accuracy))
    #
    rfModel = model.stages[1]
    print(rfModel)



if __name__ == "__main__":
    main(sys.argv)