import re
from collections import *
import os
import pandas as pd
from sklearn import feature_extraction
import sklearn
from sklearn.feature_extraction.text import TfidfTransformer
from sklearn.feature_extraction.text import CountVectorizer
transformer = TfidfTransformer()
#构建词汇表
vectorizer = CountVectorizer(analyzer='char')
ops_corus = []

#从asm文件提取opcode
def getOpcodeSequence(filename):
    print(filename)
    opcode_seq = []
    p = re.compile(r'\s([a-fA-F0-9]{2}\s)+\s*([a-z]+)')
    with open(filename,encoding='gb18030',errors='ignore') as f:
        for line in f:
            if line.startswith(".text"):
                m = re.findall(p,line)
                if m:
                    opc = m[0][1]
                    if opc != "align":
                        opcode_seq.append(opc)
    with open("opcode.txt",'w') as f2:
        f2.writelines(opcode_seq)
        f2.write("\t")
    #print(opcode_seq)
    return opcode_seq

#n-gram n=3
def getOpcodeNgram(ops, n=3):
    opngramlist = [tuple(ops[i:i+n]) for i in range(len(ops)-n)]
    opngram = Counter(opngramlist)
    return opngram

for i in range(1,4):
    basepath = "D:\MyDownloads\\train\\subtrain\\"
    map3gram = defaultdict(Counter)
    subtrain = pd.read_csv('subtrainLabels.csv')
    count = 1
    for sid in subtrain.Id:
        print ("counting the 3-gram of the {0} file...".format(str(count)))
        count += 1
        filename = basepath + sid + ".asm"
        ops = getOpcodeSequence(filename)
        op3gram = getOpcodeNgram(ops)
        ops_corus.append(' '.join(str(op3gram).strip()))
        map3gram[sid] = op3gram

#compute TFIDF
#print(ops_corus)
    ops_count = vectorizer.fit_transform(ops_corus)
#print(ops_count)
#统计每一个词语的tfdif值
    tfidf = transformer.fit_transform(ops_count)
#获取词袋模型中所有的词语
    word = vectorizer.get_feature_names()
#将tf-idf矩阵抽取出来,元素a[i][j]表示此词在i类文本中的权重
    weight = tfidf.toarray()

    cc = Counter([])
    for d in map3gram.values():
        cc += d
    selectedfeatures = {}
    tc = 0

    for k,v in cc.most_common(50*i):
        selectedfeatures[k] = v
        print (k,v)
        tc += 1
        dataframelist = []
    j = 0
    for fid,op3gram in map3gram.items():
        standard = {}
        standard["Id"] = fid
        for feature in selectedfeatures:
            if feature in op3gram:
                standard[feature] = op3gram[feature]
            else:
                standard[feature] = 0
        dataframelist.append(standard)
        j += 1
    df = pd.DataFrame(dataframelist)
    df.to_csv("TFDIF-{0}feature.csv".format(50*i),index=False)