from pyspark.sql import SparkSession
from pyspark.ml.linalg import SparseVector, VectorUDT, Vectors
from pyspark.sql.types import *
from pyspark.rdd import RDD, PipelinedRDD
import numpy as np

#preprocessor for bytes files
class preprocessor_bytes(object):
    '''
    preprocessor for .bytes files in Microsoft Malware Classification Challenge
    generates the following features:
      - counts of hexademical bigrams (e.g., '53 8F')

    parameters:
      - most_common_bigrams: boolean
        whether or not to filter out everything but the top 1000 bigrams from each class (default True)
        
    methods:
      - hashes_and_labels(X_rdd, y_rdd)
        combines rdd of hashes and rdd of lables into a single rdd
        parameters:
          - X_rdd: pyspark rdd
            rdd with hashes as rows
          - y_rdd: pyspark rdd
            rdd with labels as rows
        outputs:
          - rdd with (hash,label) as rows
          
      - get_s3_path(X_rdd)
        convert rdd of hashes to their associated S3 paths for .asm files
        parameters:
          - X_rdd: pyspark rdd
            rdd with hashes as rows
        outputs:
          - list of strings containing S3 paths to .asm files
        
      - transform(metadata, hashes_and_labels, train)
        extract features from .bytes files
        parameters:
          - metadata: pyspark rdd
            rdd with .bytes files as rows (generated using sc.WholeTextFiles)
          - hashes_and_labels: pyspark rdd (optional)
            rdd with (hash,label) as rows
        outputs:
          - spark dataframe with hash, features, and label for each row
    
    requires following non-standard python packages
      - numpy
    '''
    def __init__(self,most_common_bigrams=True):
        self.most_common_bigrams = most_common_bigrams
        self.num_features = 65536
        if self.most_common_bigrams:
            self.bigrams = list(np.load('./data/most_common_bigrams.npy'))
            self.num_features = len(self.bigrams)
    
    def _term_frequency(self,bytes):
        '''
        get term frequency of 4-hexadecimal-character words
        '''
        
        #get hash
        hash = bytes[0].encode('utf-8').rsplit('/',1)[1].split('.bytes')[0]
        
        #tokenize
        cleaned = bytes[1].replace("\r\n"," ")
        tokens = [word for word in cleaned.split() if word!='??' and len(word)==2]

        #get hex bigrams and convert hex to int
        tf = np.zeros(65536)
        for idx in xrange(len(tokens)-1):
            word = tokens[idx]+tokens[idx+1]
            converted = int(word,16)
            tf[converted] += 1
        if self.most_common_bigrams:
            tf = tf[self.bigrams]
        return (hash,tf)
    
    def hashes_and_labels(self,X_rdd,y_rdd):
        '''
        Combine rdd of hashes and rdd of lables into a single rdd
        '''
        #Check input type
        if type(X_rdd) != RDD:
            raise TypeError("Arguments must be pySpark RDDs")
        if type(y_rdd) != RDD:
            raise TypeError("Arguments must be pySpark RDDs")

        hashes = X_rdd.map(lambda hash:hash.encode("utf-8")).zipWithIndex().map(lambda x:(x[1],x[0])) 
        labels = y_rdd.map(lambda hash:hash.encode("utf-8")).zipWithIndex().map(lambda x:(x[1],x[0]))   
       
        #Join X_rdd(hashes) and y_rdd(labels)
        hash_and_labels = hashes.join(labels).map(lambda (idx,(hash,label)):(hash,label))
        return hash_and_labels
    
    def get_s3_path(self,X_rdd):
        '''
        Convert rdd of hashes to their associated S3 paths for .bytes files
        '''
        #Check input type
        if type(X_rdd) != RDD:
            raise TypeError("Arguments must be pySpark RDDs")
            
        s3_path = X_rdd.map(lambda row: 's3n://eds-uga-csci8360/data/project2/binaries/'+row+'.bytes').collect()
        return list(s3_path)

    def transform(self,bytes,hashes_and_labels=None):
        '''
        extract features from .bytes files
        '''
        #Check input type
        if type(bytes) != RDD:
            raise TypeError("Arguments must be pySpark RDDs")
        if hashes_and_labels and type(hashes_and_labels) != RDD and type(hashes_and_labels) != PipelinedRDD:
            raise TypeError("Arguments must be pySpark RDDs")
        
        #get term frequencies
        X = bytes.map(self._term_frequency).cache()
        
        #convert to sparse
        X = X.map(lambda (hash,features): (hash, SparseVector(self.num_features,np.nonzero(features)[0],features[features>0])))

        #check if labels exist
        if hashes_and_labels:
            #combine X and y into single dataframe
            data = hashes_and_labels.join(X).map(lambda (hash,(label,features)):(hash,features,label))
            schema = StructType([StructField('hash',StringType(),True),StructField('features',VectorUDT(),True),StructField('label',StringType(),True)])
            data = data.toDF(schema)
            data = data.withColumn('label',data.label.cast(DoubleType()))
        
        else:
            #if no labels, just use X
            schema = StructType([StructField('hash',StringType(),True),StructField("features", VectorUDT(), True)])
            data = X.toDF(schema)

        return data

if __name__ == '__main__':        

    #initialize spark session
    spark = SparkSession\
            .builder\
            .appName("Test")\
            .getOrCreate()
    sc = spark.sparkContext
    
    sc._jsc.hadoopConfiguration().set("fs.s3n.awsAccessKeyId", "<KEY>")
    sc._jsc.hadoopConfiguration().set("fs.s3n.awsSecretAccessKey", "<KEY>")
    
    #paths to training data
    X_file = "./data/X_train_small.txt"
    y_file = "./data/y_train_small.txt"
    X_file = sc.textFile(X_file,20)
    y_file = sc.textFile(y_file,20)
    
    #preprocess data
    preprocessor1 = preprocessor_bytes(most_common_bigrams=True)
    hashes_and_labels = preprocessor1.hashes_and_labels(X_file,y_file)
    s3_path = preprocessor1.get_s3_path(X_file)
    bytes = sc.wholeTextFiles(','.join(s3_path),20)
    #bytes = sc.wholeTextFiles('./data/bytes/*',20)
    
    data = preprocessor1.transform(bytes,hashes_and_labels)

    print data.show()

    #save to parquet
    try:
        data.write.save("./data/train_small_bytes.parquet")
    except:
        pass
        
    #paths to test data
    X_file = "./data/X_test_small.txt"
    y_file = "./data/y_test_small.txt"
    X_file = sc.textFile(X_file,20)
    y_file = sc.textFile(y_file,20)
    
    #preprocess data
    hashes_and_labels = preprocessor1.hashes_and_labels(X_file,y_file)
    s3_path = preprocessor1.get_s3_path(X_file)
    bytes = sc.wholeTextFiles(','.join(s3_path),20)
    
    data = preprocessor1.transform(bytes,hashes_and_labels)
    
    print data.show()
    
    #save to parquet
    try:
        data.write.save("./data/test_small_bytes.parquet")
    except:
        pass