#imports
import re
import os
import numpy as np
from collections import Counter

#disassembled file to parse (benign, malware or unknown)
filenames = ['lab10.txt']

#TODO: This section can be made into faster, more efficient code

#convert files of words into arrays
math = open('instruction_listings/math.txt','r')
matharray = list(math.read().split())
math.close()

logic = open('instruction_listings/logic.txt','r')
logicarray = list(logic.read().split())
logic.close()

priv = open ('instruction_listings/priv.txt','r')
privarray = list(priv.read().split())
priv.close()

stack = open('instruction_listings/stack.txt','r')
stackarray = list(stack.read().split())
stack.close()

nop = open('instruction_listings/nop.txt','r')
noparray = list(nop.read().split())
nop.close()

other = open('instruction_listings/other.txt','r')
otherarray = list(other.read().split())
other.close()

memory = open('instruction_listings/memory.txt','r')
memoryarray = list(memory.read().split())
memory.close()

branch = open('instruction_listings/branch.txt','r')
brancharray = list(branch.read().split())
branch.close()

#convert arrays into dictionaries
math_dict = dict()
math_dict.update(dict.fromkeys(matharray,0))

logic_dict = dict()
logic_dict.update(dict.fromkeys(logicarray,1))

priv_dict = dict()
priv_dict.update(dict.fromkeys(privarray, 2))

stack_dict = dict()
stack_dict.update(dict.fromkeys(stackarray, 3))

nop_dict = dict()
nop_dict.update(dict.fromkeys(noparray,4))

other_dict = dict()
other_dict.update(dict.fromkeys(otherarray,5))

memory_dict = dict()
memory_dict.update(dict.fromkeys(memoryarray,6))

branch_dict = dict()
branch_dict.update(dict.fromkeys(brancharray,7))

#create big dictionary of all keys for easy processing
word_set = dict(math_dict.items() + logic_dict.items() + priv_dict.items() + stack_dict.items() + nop_dict.items() + other_dict.items() + memory_dict.items() + branch_dict.items())

#empty transition matrix
p = np.zeros((8,8), dtype = np.float)

#In the first loop of the file, lists all the states of the file
wordfound = list()

#individual loop file by file
#loop for second file?
#for line in f1: #loop can be adjusted for more files
#	for word in word_set:
#		for m in re.finditer(word,line):
#			wordfound.append(word)

#loop to go through above list of filenames instead of individual as above
for item in filenames:
    file = open(item,'r')
    for line in file:
	    for word in word_set:
		    for m in re.finditer(word,line):
			    wordfound.append(word)


#print all states
print ""
print "The Current States in the file are: "
print ""
print wordfound
print '-----------------------------------------------------------------'
print ""


statelist = map(lambda word: word_set[word], wordfound)
transition_counts = Counter(zip(statelist, statelist[1:]))

for (x,y), c in transition_counts.iteritems():
    print "(%s,%s) = %d" % (x,y,c)
    p[x,y] = c

print ""
print "THE FINAL TRANSITION MATRIX:"
print ""
print p
print ""
print ""
print "THE FINAL PROBABILITY MATRIX"
pr = p/p.sum(axis=1)[:,None]
print pr
print ""
np.save("UnknownProbabilityMatrix",pr)

